Upload 15 files

rag_pipeline.py

@@ -454,11 +454,13 @@ class RAGPipeline:
             ),
         ])
 
-        self.compose_prompt = ChatPromptTemplate.from_messages([
+        # SHORT compose prompt (for WhatsApp)
+        self.compose_prompt_short = ChatPromptTemplate.from_messages([
             (
                 "system",
                 ANTI_HALLUCINATION_SYSTEM + "\n\n" +
                 "You are Káàntà AI, a senior Nigerian tax consultant. Build expert answers ONLY from provided facts.\n\n"
+                "RESPONSE STYLE: BRIEF - Answer in 3-10 concise sentences for WhatsApp. Lead with the key answer immediately. Focus on the most critical information only.\n\n"
                 "WRITING STYLE: Lead with specific numbers and percentages. Remove [F1] fact IDs from final output.\n\n"
                 "PROHIBITED CONTENT:\n"
                 "- DO NOT add generic compliance warnings like 'consult a tax professional' or 'comply with regulations'\n"
@@ -471,7 +473,7 @@ class RAGPipeline:
                 "{answer_schema}\n"
                 "Rules:\n"
                 "- Every detail must reference at least one fact_id.\n"
-                "- \"explainer\" items expand on the mechanics; \"key_points\" focus on actions/implications and must NOT repeat explainer text.\n"
+                "- \"explainer\" items should be 1-9 sentences max; \"key_points\" 1 sentence each - BRIEF for WhatsApp\n"
                 "- Set ask_for_income=true only when personalized calculations are impossible without it.\n"
                 "- Keep wording concise and practical."
             ),
@@ -483,6 +485,48 @@ class RAGPipeline:
             ),
         ])
 
+        # LONG compose prompt (for PDF reports)
+        self.compose_prompt_long = ChatPromptTemplate.from_messages([
+            (
+                "system",
+                ANTI_HALLUCINATION_SYSTEM + "\n\n" +
+                "You are Káàntà AI, a senior Nigerian tax consultant. Build expert answers ONLY from provided facts.\n\n"
+                "RESPONSE STYLE: COMPREHENSIVE REPORT for PDF - Provide detailed explanation with:\n"
+                "- Thorough concept explanation with background context\n"
+                "- Multiple real-world examples with step-by-step calculations\n"
+                "- Tables comparing different scenarios (e.g., income brackets, tax rates)\n"
+                "- Numerical breakdowns showing how amounts are derived\n"
+                "- Specific references to Nigerian tax laws (e.g., 'Per Finance Act 2023, Section X...')\n"
+                "- Practical implications and edge cases\n"
+                "Format professionally for a PDF report with clear sections.\n\n"
+                "WRITING STYLE: Lead with specific numbers and percentages. Remove [F1] fact IDs from final output.\n\n"
+                "PROHIBITED CONTENT:\n"
+                "- DO NOT add generic compliance warnings like 'consult a tax professional' or 'comply with regulations'\n"
+                "- DO NOT add administrative penalty warnings unless they are specifically mentioned in the facts\n"
+                "- Focus on answering the user's question with facts, not generic advice\n\n"
+                "Workflow:\n"
+                "1. Inside <analysis></analysis>, plan the unique insights you will cover. List fact IDs you will use. "
+                "If a fact would appear twice, mark it as DUPLICATE and drop the repeat.\n"
+                "2. Inside <final></final>, output JSON that matches this schema:\n"
+                "{answer_schema}\n"
+                "Rules:\n"
+                "- Every detail must reference at least one fact_id.\n"
+                "- \"explainer\" items should provide comprehensive explanations with examples and calculations\n"
+                "- \"key_points\" should cover actions, implications, edge cases, and scenarios - detailed for PDF\n"
+                "- Set ask_for_income=true only when personalized calculations are impossible without it.\n"
+                "- Provide thorough, professional report-quality content."
+            ),
+            (
+                "human",
+                "Question:\n{question}\n\n"
+                "Verified facts (JSON):\n{facts_json}\n"
+                "Follow the required tag structure."
+            ),
+        ])
+
+        # Keep backward compatibility - default to short
+        self.compose_prompt = self.compose_prompt_short
+
         self.chain = self._build_chain()
         print("RAG pipeline ready")
 
@@ -780,10 +824,16 @@ class RAGPipeline:
             )
         return cleaned
 
-    def _compose_from_facts(self, question: str, facts: List[Dict[str, Any]]) -> Optional[str]:
+    def _compose_from_facts(self, question: str, facts: List[Dict[str, Any]], response_type: str = 'short') -> Optional[str]:
         if not facts:
             return None
 
+        # Select appropriate compose prompt based on response_type
+        if response_type.lower() == 'long':
+            compose_prompt = self.compose_prompt_long
+        else:
+            compose_prompt = self.compose_prompt_short
+
         facts_json = json.dumps({"facts": facts}, ensure_ascii=False)
         payload = {
             "question": question,
@@ -791,7 +841,7 @@ class RAGPipeline:
             "answer_schema": ANSWER_SCHEMA_TEXT,
         }
 
-        raw = (self.compose_prompt | self.llm | StrOutputParser()).invoke(payload)
+        raw = (compose_prompt | self.llm | StrOutputParser()).invoke(payload)
         _, final_json = self._extract_analysis_and_final(raw)
         structured = self._safe_json_parse(final_json)
         if not structured:
@@ -860,7 +910,7 @@ class RAGPipeline:
 
         return final_output
 
-    def _fact_guided_answer(self, question: str) -> str:
+    def _fact_guided_answer(self, question: str, response_type: str = 'short') -> str:
         docs = self._retrieve(question)
         snippets = self._prepare_context_snippets(docs)
         if not snippets:
@@ -868,7 +918,7 @@ class RAGPipeline:
 
         try:
             facts = self._harvest_facts(question, snippets)
-            response = self._compose_from_facts(question, facts)
+            response = self._compose_from_facts(question, facts, response_type=response_type)
             if response:
                 return response
         except Exception as exc:
@@ -1113,11 +1163,22 @@ class RAGPipeline:
         return "qa"
 
     # Stub for a future extractor chain - currently route extractor requests to QA chain with strict rules
-    def _extract_structured(self, question: str) -> str:
-        return self._fact_guided_answer(question)
+    def _extract_structured(self, question: str, response_type: str = 'short') -> str:
+        return self._fact_guided_answer(question, response_type=response_type)
+
+    def query(self, question: str, verbose: bool = False, response_type: str = 'short') -> str:
+        """
+        Route and answer the question with persona-aware responses.
+
+        Args:
+            question: User's tax question
+            verbose: If True, print debug information
+            response_type: 'short' for WhatsApp messages (3-4 sentences),
+                          'long' for PDF reports (comprehensive with examples)
 
-    def query(self, question: str, verbose: bool = False) -> str:
-        """Route and answer the question with persona-aware responses."""
+        Returns:
+            Formatted answer based on response_type
+        """
         # First, check if question is tax-related
         if not self._is_tax_related_question(question):
             return (
@@ -1179,9 +1240,9 @@ class RAGPipeline:
         if task == "summarize":
             return self._summarize_chapter(question)
         elif task == "extract":
-            return self._extract_structured(question)
+            return self._extract_structured(question, response_type=response_type)
         else:
-            return self._fact_guided_answer(question)
+            return self._fact_guided_answer(question, response_type=response_type)
 
 
 def main():

test_response_types.py

@@ -0,0 +1,140 @@
+#!/usr/bin/env python3
+"""
+Test script to demonstrate SHORT vs LONG response types.
+Shows the difference between WhatsApp (brief) and PDF (comprehensive) outputs.
+"""
+
+import sys
+from pathlib import Path
+from rag_pipeline import RAGPipeline, DocumentStore
+
+
+def test_response_types():
+    """Test both SHORT and LONG response types with the same question."""
+
+    print("=" * 80)
+    print("RESPONSE TYPE COMPARISON TEST")
+    print("=" * 80)
+
+    # Initialize RAG pipeline
+    print("\nInitializing RAG pipeline...")
+    vector_store_path = Path("vector_store")
+    doc_store = DocumentStore(
+        persist_dir=vector_store_path,
+        embedding_model="BAAI/bge-large-en-v1.5"
+    )
+
+    src = Path("data")
+    pdfs = doc_store.discover_pdfs(src)
+    doc_store.build_vector_store(pdfs, force_rebuild=False)
+
+    rag = RAGPipeline(
+        doc_store=doc_store,
+        model="llama-3.3-70b-versatile",
+        temperature=0.1,
+        top_k=15,
+    )
+
+    print("✓ RAG pipeline initialized\n")
+
+    # Test question
+    question = "What are the personal income tax rates in Nigeria?"
+
+    print("=" * 80)
+    print("TEST QUESTION:")
+    print(question)
+    print("=" * 80)
+
+    # Test SHORT response (WhatsApp)
+    print("\n" + "=" * 80)
+    print("SHORT RESPONSE (for WhatsApp)")
+    print("=" * 80)
+    print("\nExpected: 3-4 concise sentences, immediate answer, key facts only\n")
+
+    try:
+        short_answer = rag.query(question, verbose=False, response_type='short')
+        print(short_answer)
+
+        # Quality checks for SHORT
+        print("\n" + "-" * 80)
+        print("SHORT RESPONSE QUALITY CHECKS:")
+        word_count = len(short_answer.split())
+        sentence_count = short_answer.count('.') + short_answer.count('?') + short_answer.count('!')
+        has_numbers = any(char.isdigit() for char in short_answer)
+
+        print(f"  Word count: {word_count} (target: 50-150 words for brief)")
+        print(f"  Sentence count: ~{sentence_count}")
+        print(f"  Contains numbers: {has_numbers}")
+
+        if word_count <= 200:
+            print("  ✓ PASS: Response is concise")
+        else:
+            print("  ⚠️ WARNING: Response may be too long for WhatsApp")
+
+    except Exception as e:
+        print(f"❌ ERROR: {e}")
+        import traceback
+        traceback.print_exc()
+
+    # Test LONG response (PDF)
+    print("\n\n" + "=" * 80)
+    print("LONG RESPONSE (for PDF Report)")
+    print("=" * 80)
+    print("\nExpected: Comprehensive with examples, calculations, tables, law references\n")
+
+    try:
+        long_answer = rag.query(question, verbose=False, response_type='long')
+        print(long_answer)
+
+        # Quality checks for LONG
+        print("\n" + "-" * 80)
+        print("LONG RESPONSE QUALITY CHECKS:")
+        word_count = len(long_answer.split())
+        has_examples = 'example' in long_answer.lower() or 'instance' in long_answer.lower()
+        has_calculations = '×' in long_answer or 'calculate' in long_answer.lower()
+        has_law_refs = 'section' in long_answer.lower() or 'act' in long_answer.lower()
+        has_numbers = any(char.isdigit() for char in long_answer)
+
+        print(f"  Word count: {word_count} (target: 300+ words for comprehensive)")
+        print(f"  Contains examples: {has_examples}")
+        print(f"  Contains calculations: {has_calculations}")
+        print(f"  Contains law references: {has_law_refs}")
+        print(f"  Contains numbers: {has_numbers}")
+
+        if word_count >= 300:
+            print("  ✓ PASS: Response is comprehensive")
+        else:
+            print("  ⚠️ WARNING: Response may be too brief for PDF report")
+
+        if has_examples and has_numbers:
+            print("  ✓ PASS: Response includes examples and numbers")
+        else:
+            print("  ⚠️ WARNING: Response may lack examples or numbers")
+
+    except Exception as e:
+        print(f"❌ ERROR: {e}")
+        import traceback
+        traceback.print_exc()
+
+    # Summary
+    print("\n\n" + "=" * 80)
+    print("COMPARISON SUMMARY")
+    print("=" * 80)
+    print("\nKEY DIFFERENCES:")
+    print("  SHORT (WhatsApp):")
+    print("    - 3-4 sentences")
+    print("    - Immediate answer")
+    print("    - Key facts only")
+    print("    - No examples or detailed calculations")
+    print("")
+    print("  LONG (PDF Report):")
+    print("    - Multiple paragraphs")
+    print("    - Detailed explanations")
+    print("    - Examples with step-by-step calculations")
+    print("    - Law references and edge cases")
+    print("    - Professional report format")
+    print("=" * 80)
+
+
+if __name__ == "__main__":
+    test_response_types()