Implement Hybrid Agent architecture: Gemini 2.5 Lite (Router) + Gemini 3 Flash (Responder)

agentic_rag_v2_graph.py

@@ -14,7 +14,8 @@ from llm_utils import generate_with_retry
 from sql_db import query_database
 
 # Config
-MODEL_NAME = "gemini-2.5-flash"
+MODEL_FAST = "gemini-2.5-flash-lite"
+MODEL_SMART = "gemini-3-flash-preview"
 MAX_RETRIES = 2
 
 # ===============================
@@ -79,7 +80,7 @@ def text_to_sql_tool(query: str):
     - Output ONLY the SQL query. No markdown.
     - Do NOT use Markdown formatting.
     """
-    model = genai.GenerativeModel(MODEL_NAME)
+    model = genai.GenerativeModel(MODEL_SMART)
     resp = generate_with_retry(model, prompt)
     sql_query = resp.text.strip().replace("```sql", "").replace("```", "").strip() if resp else ""
     
@@ -121,7 +122,7 @@ def supervisor_node(state: AgentState):
     # Heuristic: If we already searched SQL and got results, maybe go to responder or PDF
     # But for now, let LLM decide based on history.
     
-    model = genai.GenerativeModel(MODEL_NAME)
+    model = genai.GenerativeModel(MODEL_FAST)
     resp = generate_with_retry(model, prompt)
     decision = resp.text.strip().lower() if resp else "responder"
     
@@ -153,7 +154,38 @@ def researcher_sql_node(state: AgentState):
     current_outputs = state.get("tool_outputs", []) + results
     return {**state, "tool_outputs": current_outputs}
 
-# ... (Verifier is unchanged) ...
+# 5. VERIFIER
+def verifier_node(state: AgentState):
+    """Verifies the quality of gathered information."""
+    query = state["query"]
+    tools_out = state.get("tool_outputs", [])
+    
+    # Simple verification logic
+    context = ""
+    for t in tools_out:
+        context += f"\n[{t['source'].upper()}]: {t['content']}..."
+
+    prompt = f"""
+    You are a Verifier Agent.
+    User Query: "{query}"
+    
+    Gathered Info:
+    {context}
+    
+    Task:
+    Analyze the gathered information. 
+    - Is it relevant to the query?
+    - Are there conflicts?
+    - What key details are present?
+    
+    Provide concise verification notes for the Final Responder.
+    """
+    
+    model = genai.GenerativeModel(MODEL_SMART)
+    resp = generate_with_retry(model, prompt)
+    notes = resp.text if resp else "Verification completed."
+    
+    return {**state, "verification_notes": notes}
 
 # 6. RESPONDER
 def responder_node(state: AgentState):
@@ -183,7 +215,7 @@ def responder_node(state: AgentState):
     Answer the user query. If you used SQL, summarize the data insights.
     """
     
-    model = genai.GenerativeModel(MODEL_NAME)
+    model = genai.GenerativeModel(MODEL_SMART)
     resp = generate_with_retry(model, prompt)
     answer = resp.text if resp else "I could not generate an answer."
     

llm_utils.py

@@ -11,7 +11,7 @@ class DummyResponse:
     def text(self):
         return self._text
 
-def generate_with_retry(model, prompt, retries=3, base_delay=2):
+def generate_with_retry(model, prompt, retries=5, base_delay=4):
     """
     Generates content using the Gemini model with exponential backoff for rate limits.
     Returns a dummy response if all retries fail, preventing app crashes.

main.py

@@ -21,7 +21,7 @@ load_dotenv()
 genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 
 MOCK_MODE = False  # Refactor complete - enabling real agent
-MODEL_NAME = "gemini-2.5-flash"
+MODEL_NAME = "gemini-3-flash-preview"
 MAX_FILE_SIZE = 50 * 1024 * 1024
 CACHE_TTL = 300
 
@@ -79,8 +79,12 @@ def analytics():
 @app.post("/upload")
 async def upload(files: list[UploadFile] = File(...)):
     for file in files:
-        ext = file.filename.split(".")[-1].lower()
+        filename = file.filename
+        ext = filename.split(".")[-1].lower() if "." in filename else ""
+        print(f"🔍 DEBUG: Uploading '{filename}' (Ext: {ext})")
+
         if ext not in ["pdf", "txt"]:
+            print(f"❌ REJECTED: Invalid extension '{ext}'")
             return JSONResponse(
                 status_code=400,
                 content={"error": "Only PDF and TXT files allowed"}

rag_store.py

@@ -99,11 +99,17 @@ def ingest_documents(files):
             try:
                 # Use pymupdf4llm to extract markdown with tables
                 import pymupdf4llm
-                md_text = pymupdf4llm.to_markdown(tmp_path)
+                # Get list of dicts: [{'text': '...', 'metadata': {'page': 1, ...}}]
+                pages_data = pymupdf4llm.to_markdown(tmp_path, page_chunks=True)
                 
-                for chunk in chunk_text(md_text):
-                    texts.append(chunk)
-                    meta.append({"source": file.filename, "page": "N/A"}) # pymupdf4llm merges pages by default
+                for page_obj in pages_data:
+                    p_text = page_obj["text"]
+                    p_num = page_obj["metadata"].get("page", "N/A")
+                    
+                    # Chunk within the page to preserve page context
+                    for chunk in chunk_text(p_text):
+                        texts.append(chunk)
+                        meta.append({"source": file.filename, "page": p_num})
             finally:
                 os.remove(tmp_path)
 

verify_graph_flow.py

@@ -0,0 +1,91 @@
+import unittest
+from unittest.mock import MagicMock, patch
+import os
+
+# Set dummy key if not present to avoid init errors
+if "GEMINI_API_KEY" not in os.environ:
+    os.environ["GEMINI_API_KEY"] = "dummy_key"
+if "TAVILY_API_KEY" not in os.environ:
+    os.environ["TAVILY_API_KEY"] = "dummy_key"
+
+from agentic_rag_v2_graph import build_agentic_rag_v2_graph
+
+class TestRagGraph(unittest.TestCase):
+    @patch('agentic_rag_v2_graph.genai.GenerativeModel')
+    @patch('agentic_rag_v2_graph.TavilyClient')
+    def test_web_search_flow(self, mock_tavily, mock_genai):
+        print("\n\n=== 🧪 STARTING DRY RUN GRAPH TEST ===")
+        print("Goal: Verify 'research_web' -> 'verifier' -> 'responder' flow without API calls.\n")
+
+        # === 🔧 DEMO CONFIGURATION (EDIT HERE) ===
+        # Change these values to simulate different questions!
+        DEMO_QUERY = "Who is the father of the computer?"
+        EXPECTED_WEB_CONTENT = "Charles Babbage is considered by many as the father of the computer."
+        VERIFIER_NOTE = "✅ VERIFIED: Search results confirm Charles Babbage invented the Analytical Engine."
+        FINAL_ANSWER = "Charles Babbage is the father of the computer."
+        # =========================================
+
+        # --- Setup Mocks ---
+        mock_model = MagicMock()
+        mock_genai.return_value = mock_model
+        
+        # Helper to create dummy response object
+        def create_response(text):
+            r = MagicMock()
+            r.text = text
+            return r
+
+        # Sequence of LLM outputs (Order matters!):
+        # 1. Supervisor: "research_web"
+        # 2. Verifier: "The info looks consistent."
+        # 3. Supervisor: "responder"
+        # 4. Responder: "The Answer."
+        
+        mock_model.generate_content.side_effect = [
+            create_response("research_web"),
+            create_response(VERIFIER_NOTE),
+            create_response("responder"),
+            create_response(FINAL_ANSWER)
+        ]
+        
+        # Mock Web Search Tool
+        mock_tavily_instance = MagicMock()
+        mock_tavily.return_value = mock_tavily_instance
+        mock_tavily_instance.get_search_context.return_value = EXPECTED_WEB_CONTENT
+
+        # --- Build Graph ---
+        print("🛠️  Building Graph...")
+        try:
+            graph = build_agentic_rag_v2_graph()
+            print("✅ Graph built successfully.")
+        except Exception as e:
+            self.fail(f"❌ Graph build failed: {e}")
+        
+        # --- Run Graph ---
+        initial_state = {
+            "messages": [],
+            "query": DEMO_QUERY,
+            "final_answer": "",
+            "next_node": "",
+            "current_tool": "",
+            "tool_outputs": [],
+            "verification_notes": "",
+            "retries": 0
+        }
+        
+        print("\n🏃 Invoking Graph (Mocked LLM)...")
+        result = graph.invoke(initial_state, config={"configurable": {"thread_id": "test_dry_run"}})
+        
+        # --- Assertions ---
+        print("\n\n=== 📊 TEST RESULT ANALYSIS ===")
+        print(f"Final Answer: {result['final_answer']}")
+        print(f"Verification Notes: {result['verification_notes']}")
+        
+        self.assertIn("VERIFIED", result['verification_notes'], "❌ verifier_node did not populate verification_notes!")
+        self.assertIn(FINAL_ANSWER, result['final_answer'], "❌ Responder did not fail gracefully.")
+        
+        print("\n✅ SUCCESS: The Graph followed the correct path: Supervisor -> Web -> Verifier -> Supervisor -> Responder")
+        print("✅ SUCCESS: 'verifier_node' executed and produced notes.")
+
+if __name__ == "__main__":
+    unittest.main()

verify_rag.py

@@ -1,4 +1,7 @@
 import asyncio
+import os
+from dotenv import load_dotenv
+load_dotenv()
 from agentic_rag_v2_graph import build_agentic_rag_v2_graph
 
 async def main():
@@ -21,7 +24,7 @@ async def main():
     }
     
     result = await graph.ainvoke(inputs, config=config)
-    print(f"Answer 1: {result['answer']}")
+    print(f"Answer 1: {result['final_answer']}")
     
     print("\n--- Turn 2 ---")
     inputs["query"] = "What is my name?"
@@ -33,7 +36,7 @@ async def main():
     inputs["messages"] = []
     
     result = await graph.ainvoke(inputs, config=config)
-    print(f"Answer 2: {result['answer']}")
+    print(f"Answer 2: {result['final_answer']}")
 
 if __name__ == "__main__":
     try: