feat: implement multi-agent RAG with self-correction and web search

Added Supervisor-Worker architecture, Skeptical Verifier node, and Tavily API integration for external grounding.

agentic_rag_v2_graph.py

@@ -1,296 +1,230 @@
-from typing import TypedDict, List, Optional, Annotated
+import os
+import time
+from typing import TypedDict, List, Optional, Annotated, Literal
 import google.generativeai as genai
 from langgraph.graph import StateGraph, END
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph.message import add_messages
-from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
-import time
-import random
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
+from tavily import TavilyClient
 
 from rag_store import search_knowledge
 from eval_logger import log_eval
 from llm_utils import generate_with_retry
 
+# Config
 MODEL_NAME = "gemini-2.5-flash"
 MAX_RETRIES = 2
 
-
-
-
-def format_history(messages: List[BaseMessage]) -> str:
-    history_str = ""
-    for msg in messages:
-        role = "User" if isinstance(msg, HumanMessage) else "Assistant"
-        history_str += f"{role}: {msg.content}\n"
-    return history_str
-
 # ===============================
 # STATE
 # ===============================
 class AgentState(TypedDict):
     messages: Annotated[List[BaseMessage], add_messages]
     query: str
-    refined_query: str
-    decision: str
-    retrieved_chunks: List[dict]
-    retrieval_quality: str
+    final_answer: str
+    
+    # Internal routing & scratchpad
+    next_node: str
+    current_tool: str
+    tool_outputs: List[dict]  # list of {source: 'pdf'|'web', content: ..., score: ...}
+    verification_notes: str
     retries: int
-    answer: Optional[str]
-    confidence: float
-    answer_known: bool
-
 
 # ===============================
-# LLM DECISION NODE (PLANNER)
+# TOOLS
 # ===============================
-def llm_decision_node(state: AgentState) -> AgentState:
-    history = format_history(state.get("messages", []))
-    prompt = f"""
-You are an AI agent deciding whether a question requires document retrieval.
-Answer ONLY one word:
-- use_rag
-- no_rag
-
-Conversation History:
-{history}
-
-Current Question:
-{state["query"]}
-"""
-    model = genai.GenerativeModel(MODEL_NAME)
-    resp = generate_with_retry(model, prompt)
-
-    decision = "use_rag"
-    if resp and "no_rag" in resp.text.lower():
-        decision = "no_rag"
-
-    return {**state, "decision": decision}
+def pdf_search_tool(query: str):
+    """Searches internal PDF knowledge base."""
+    results = search_knowledge(query, top_k=4)
+    # Format for consumption
+    return [
+        {
+            "source": "internal_pdf",
+            "content": r["text"],
+            "metadata": r["metadata"],
+            "score": r.get("score", 0)
+        }
+        for r in results
+    ]
 
+def web_search_tool(query: str):
+    """Searches the web using Tavily."""
+    api_key = os.getenv("TAVILY_API_KEY")
+    if not api_key:
+        return [{"source": "external_web", "content": "Error: TAVILY_API_KEY not found.", "score": 0}]
+    
+    try:
+        tavily = TavilyClient(api_key=api_key)
+        # Search context first for cleaner text
+        context = tavily.get_search_context(query=query, search_depth="advanced")
+        return [{
+            "source": "external_web", 
+            "content": context, 
+            "score": 0.8 # Arbitrary confidence for web
+        }]
+    except Exception as e:
+        return [{"source": "external_web", "content": f"Web search error: {str(e)}", "score": 0}]
 
 # ===============================
-# RETRIEVAL NODE (TOOL)
+# NODES
 # ===============================
-def retrieve_node(state: AgentState) -> AgentState:
-    q = state["refined_query"] or state["query"]
-    chunks = search_knowledge(q)
-    return {**state, "retrieved_chunks": chunks}
-
 
-# ===============================
-# GRADE DOCUMENTS NODE (GRADER)
-# ===============================
-def grade_documents_node(state: AgentState) -> AgentState:
-    """
-    Determines whether the retrieved documents are relevant to the question.
-    """
+# 1. SUPERVISOR
+def supervisor_node(state: AgentState):
+    """Decides whether to research (and which tool) or answer."""
     query = state["query"]
-    retrieved_docs = state["retrieved_chunks"]
+    history_len = len(state.get("messages", []))
+    
+    # If we already have tools output, check if we need more or are done
+    tools_out = state.get("tool_outputs", [])
     
-    filtered_docs = []
-    for doc in retrieved_docs:
-        prompt = f"""
-        You are a grader assessing relevance of a retrieved document to a user question.
-        
-        Retrieved document:
-        {doc['text']}
-        
-        User question:
-        {query}
-        
-        If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant.
-        Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.
-        
-        Answer ONLY 'yes' or 'no'.
-        """
-        model = genai.GenerativeModel(MODEL_NAME)
-        resp = generate_with_retry(model, prompt)
-        score = resp.text.strip().lower() if resp else "no"
-        
-        if "yes" in score:
-            filtered_docs.append(doc)
-            
-    return {**state, "retrieved_chunks": filtered_docs}
-
-
-# ===============================
-# RETRIEVAL EVALUATION (CRITIC)
-# ===============================
-def evaluate_retrieval_node(state: AgentState) -> AgentState:
-    if not state["retrieved_chunks"]:
-        return {**state, "retrieval_quality": "bad"}
-
-    context_sample = "\n".join(c["text"][:200] for c in state["retrieved_chunks"][:3])
-
     prompt = f"""
-Evaluate whether the following retrieved context is sufficient
-to answer the question.
-
-Answer ONLY one word:
-- good
-- bad
-
-Question:
-{state["query"]}
-
-Context:
-{context_sample}
-"""
-
-    model = genai.GenerativeModel(MODEL_NAME)
-    resp = generate_with_retry(model, prompt)
-
-    quality = "bad"
-    if resp and "good" in resp.text.lower():
-        quality = "good"
+    You are a Supervisor Agent.
+    User Query: "{query}"
+    
+    Current Gathered Info Count: {len(tools_out)}
+    
+    Decide next step:
+    1. "research_pdf": If we haven't checked internal docs yet.
+    2. "research_web": If PDF info is missing/insufficient and we haven't checked web yet.
+    3. "responder": If we have enough info OR we have tried everything.
+    
+    Return ONLY one of: research_pdf, research_web, responder
+    """
+    
+    # Simple heuristic to save calls, or use LLM? 
+    # Prompt says "Planning Node: The LLM must decide".
+    
+    # We can force PDF first to be efficient
+    has_pdf = any(t["source"] == "internal_pdf" for t in tools_out)
+    if not has_pdf:
+        return {**state, "next_node": "research_pdf"}
         
-    return {**state, "retrieval_quality": quality}
-
-
-# ===============================
-# QUERY REFINEMENT (SELF-CORRECTION)
-# ===============================
-def refine_query_node(state: AgentState) -> AgentState:
-    history = format_history(state.get("messages", []))
-    prompt = f"""
-Rewrite the following question to improve document retrieval.
-Be concise and factual.
-
-Conversation History:
-{history}
-
-Original question:
-{state["query"]}
-"""
-
     model = genai.GenerativeModel(MODEL_NAME)
     resp = generate_with_retry(model, prompt)
+    decision = resp.text.strip().lower() if resp else "responder"
     
-    refined = resp.text.strip() if resp else state["query"]
-
-    return {
-        **state,
-        "refined_query": refined,
-        "retries": state["retries"] + 1
-    }
-
+    if "pdf" in decision: return {**state, "next_node": "research_pdf"}
+    if "web" in decision: return {**state, "next_node": "research_web"}
+    
+    return {**state, "next_node": "responder"}
 
-# ===============================
-# ANSWER WITH RAG (HIGH CONF)
-# ===============================
-def answer_with_rag_node(state: AgentState) -> AgentState:
-    context = "\n\n".join(c["text"] for c in state["retrieved_chunks"])
-    history = format_history(state.get("messages", []))
+# 2. RESEARCHER (PDF)
+def researcher_pdf_node(state: AgentState):
+    query = state["query"]
+    results = pdf_search_tool(query)
+    
+    # Append to tool_outputs
+    current_outputs = state.get("tool_outputs", []) + results
+    
+    # Log
+    log_eval(query, len(results), 0.9, len(results) > 0, source_type="internal_pdf")
+    
+    return {**state, "tool_outputs": current_outputs}
 
+# 3. RESEARCHER (WEB)
+def researcher_web_node(state: AgentState):
+    query = state["query"]
+    results = web_search_tool(query)
+    
+    current_outputs = state.get("tool_outputs", []) + results
+    
+    # Log
+    log_eval(query, 1, 0.7, True, source_type="external_web")
+    
+    return {**state, "tool_outputs": current_outputs}
+
+# 4. VERIFIER
+def verifier_node(state: AgentState):
+    """Cross-references Web findings against PDF context."""
+    tool_outputs = state.get("tool_outputs", [])
+    web_content = [t for t in tool_outputs if t["source"] == "external_web"]
+    pdf_content = [t for t in tool_outputs if t["source"] == "internal_pdf"]
+    
+    if not web_content:
+        return state # Nothing to verify
+        
+    # If we skipped PDF for some reason, let's quick-check it now for verification context
+    if not pdf_content:
+        pdf_content = pdf_search_tool(state["query"])
+        
+    web_text = "\n".join([c["content"] for c in web_content])
+    pdf_text = "\n".join([c["content"] for c in pdf_content])
+    
     prompt = f"""
-Answer using ONLY the context below.
-If the answer is not present, say "I don't know".
-
-Context:
-{context}
-
-Conversation History:
-{history}
-
-Question:
-{state["query"]}
-"""
-
+    You are a Skeptical Verifier.
+    
+    Query: {state["query"]}
+    
+    INTERNAL PDF KNOWLEDGE:
+    {pdf_text[:2000]}
+    
+    EXTERNAL WEB FINDINGS:
+    {web_text[:2000]}
+    
+    Task:
+    Check if the External Web Findings contradict the Internal PDF Knowledge.
+    If Web says 'X' and PDF says 'Y', report the conflict.
+    
+    Output a brief "Verification Note". If no conflict, say "No conflict".
+    """
+    
     model = genai.GenerativeModel(MODEL_NAME)
     resp = generate_with_retry(model, prompt)
-    answer_text = resp.text if resp else "Error generating answer due to quota limits."
-
-    answer_known = "i don't know" not in answer_text.lower()
-    confidence = min(0.95, 0.6 + (0.1 * len(state["retrieved_chunks"])))
-
-    log_eval(
-        query=state["query"],
-        retrieved_count=len(state["retrieved_chunks"]),
-        confidence=confidence,
-        answer_known=answer_known
-    )
-
-    # Append interaction to memory
-    new_messages = [
-        HumanMessage(content=state["query"]),
-        AIMessage(content=answer_text)
-    ]
-
-    return {
-        **state,
-        "messages": new_messages,
-        "answer": answer_text,
-        "confidence": confidence,
-        "answer_known": answer_known
-    }
-
+    note = resp.text.strip() if resp else "Verification failed."
+    
+    current_notes = state.get("verification_notes", "")
+    new_notes = f"{current_notes}\n[Verification]: {note}"
+    
+    return {**state, "verification_notes": new_notes}
 
-# ===============================
-# ANSWER WITHOUT RAG
-# ===============================
-def answer_direct_node(state: AgentState) -> AgentState:
-    history = format_history(state.get("messages", []))
+# 5. RESPONDER
+def responder_node(state: AgentState):
+    query = state["query"]
+    tools_out = state.get("tool_outputs", [])
+    notes = state.get("verification_notes", "")
+    
+    # Check if we found nothing
+    if not tools_out and state["retries"] < 1:
+        # Self-correction: Rewrite
+        prompt = f"Rewrite this query to be more specific: {query}"
+        model = genai.GenerativeModel(MODEL_NAME)
+        resp = generate_with_retry(model, prompt)
+        new_query = resp.text.strip() if resp else query
+        return {**state, "query": new_query, "retries": state["retries"] + 1, "next_node": "supervisor"} # Loop back
+        
+    context = ""
+    for t in tools_out:
+        context += f"\n[{t['source'].upper()}]: {t['content'][:500]}..."
+        
     prompt = f"""
-Conversation History:
-{history}
-
-Answer clearly and concisely:
-{state['query']}
-"""
-
+    You are the Final Responder.
+    User Query: {query}
+    
+    Gathered Info:
+    {context}
+    
+    Verification Notes (Conflicts?):
+    {notes}
+    
+    Instructions:
+    1. Answer the user query based on gathered info.
+    2. If there are conflicts (e.g. PDF vs Web), explicitly mention them and trust PDF more but note the Web claim.
+    3. Cite sources (Internal PDF vs External Web).
+    """
+    
     model = genai.GenerativeModel(MODEL_NAME)
     resp = generate_with_retry(model, prompt)
-    answer_text = resp.text if resp else "Error generating answer due to quota limits."
-
-    log_eval(
-        query=state["query"],
-        retrieved_count=0,
-        confidence=0.4,
-        answer_known=True
-    )
-
-    # Append interaction to memory
-    new_messages = [
-        HumanMessage(content=state["query"]),
-        AIMessage(content=answer_text)
-    ]
-
-    return {
-        **state,
-        "messages": new_messages,
-        "answer": answer_text,
-        "confidence": 0.4,
-        "answer_known": True
-    }
-
-
-# ===============================
-# NO ANSWER
-# ===============================
-def no_answer_node(state: AgentState) -> AgentState:
-    log_eval(
-        query=state["query"],
-        retrieved_count=0,
-        confidence=0.0,
-        answer_known=False
-    )
-    
-    answer_text = "I don't know based on the provided documents."
+    answer = resp.text if resp else "I could not generate an answer."
     
-    # Append interaction to memory
-    new_messages = [
-        HumanMessage(content=state["query"]),
-        AIMessage(content=answer_text)
-    ]
-
     return {
-        **state,
-        "messages": new_messages,
-        "answer": answer_text,
-        "confidence": 0.0,
-        "answer_known": False
+        **state, 
+        "final_answer": answer, 
+        "messages": [AIMessage(content=answer)],
+        "next_node": "end"
     }
 
-
 # ===============================
 # GRAPH BUILDER
 # ===============================
@@ -298,57 +232,40 @@ def build_agentic_rag_v2_graph():
     graph = StateGraph(AgentState)
     memory = MemorySaver()
 
-    graph.add_node("decide", llm_decision_node)
-    graph.add_node("retrieve", retrieve_node)
-    graph.add_node("grade", grade_documents_node)
-    graph.add_node("evaluate", evaluate_retrieval_node)
-    graph.add_node("refine", refine_query_node)
-    graph.add_node("answer_rag", answer_with_rag_node)
-    graph.add_node("answer_direct", answer_direct_node)
-    graph.add_node("no_answer", no_answer_node)
+    graph.add_node("supervisor", supervisor_node)
+    graph.add_node("research_pdf", researcher_pdf_node)
+    graph.add_node("research_web", researcher_web_node)
+    graph.add_node("verifier", verifier_node)
+    graph.add_node("responder", responder_node)
 
-    graph.set_entry_point("decide")
+    graph.set_entry_point("supervisor")
 
+    # Routing
     graph.add_conditional_edges(
-        "decide",
-        lambda s: s["decision"],
+        "supervisor",
+        lambda s: s["next_node"],
         {
-            "use_rag": "retrieve",
-            "no_rag": "answer_direct"
+            "research_pdf": "research_pdf",
+            "research_web": "research_web",
+            "responder": "responder"
         }
     )
-
-    graph.add_edge("retrieve", "grade")
-
-    def check_relevance(state):
-        if not state["retrieved_chunks"]:
-            if state["retries"] >= MAX_RETRIES:
-                return "no_answer"
-            return "rewrite"
-        return "evaluate"
-
-    graph.add_conditional_edges(
-        "grade",
-        check_relevance,
-        {
-            "rewrite": "refine",
-            "evaluate": "evaluate",
-            "no_answer": "no_answer"
-        }
-    )
-
+    
+    # Research PDF -> Supervisor (to decide if Web is needed)
+    graph.add_edge("research_pdf", "supervisor")
+    
+    # Research Web -> Verifier -> Supervisor
+    graph.add_edge("research_web", "verifier")
+    graph.add_edge("verifier", "supervisor")
+    
+    # Responder -> Maybe loop back if self-correction triggered?
     graph.add_conditional_edges(
-        "evaluate",
-        lambda s: "retry" if s["retrieval_quality"] == "bad" and s["retries"] < MAX_RETRIES else "answer",
+        "responder",
+        lambda s: "supervisor" if s["next_node"] == "supervisor" else "end",
         {
-            "retry": "refine",
-            "answer": "answer_rag"
+            "supervisor": "supervisor",
+            "end": END
         }
     )
 
-    graph.add_edge("refine", "retrieve")
-    graph.add_edge("answer_rag", END)
-    graph.add_edge("answer_direct", END)
-    graph.add_edge("no_answer", END)
-
     return graph.compile(checkpointer=memory)

eval_logger.py

@@ -7,14 +7,16 @@ def log_eval(
     query: str,
     retrieved_count: int,
     confidence: float,
-    answer_known: bool
+    answer_known: bool,
+    source_type: str = "internal_pdf"  # Added source_type
 ):
     record = {
         "timestamp": time(),
         "query": query,
         "retrieved_count": retrieved_count,
         "confidence": confidence,
-        "answer_known": answer_known
+        "answer_known": answer_known,
+        "source_type": source_type
     }
 
     with open(LOG_FILE, "a", encoding="utf-8") as f:

main.py

@@ -8,10 +8,11 @@ from pydantic import BaseModel
 from dotenv import load_dotenv
 import google.generativeai as genai
 
-from rag_store import ingest_documents, get_all_chunks, clear_database
+from rag_store import ingest_documents, get_all_chunks, clear_database, search_knowledge
 from analytics import get_analytics
 from agentic_rag_v2_graph import build_agentic_rag_v2_graph
 from llm_utils import generate_with_retry
+import asyncio
 
 # =========================================================
 # ENV + MODEL
@@ -19,6 +20,7 @@ from llm_utils import generate_with_retry
 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"
 MAX_FILE_SIZE = 50 * 1024 * 1024
 CACHE_TTL = 300
@@ -115,6 +117,49 @@ async def ask(data: PromptRequest):
         if now - ts < CACHE_TTL:
             return cached
 
+    # ==========================
+    # 🟧 MOCK MODE (NO API)
+    # ==========================
+    if MOCK_MODE:
+        await asyncio.sleep(0.5)  # Simulate latency
+        
+        mock_answer = ""
+        mock_citations = []
+        
+        if "summary" in key or "summarize" in key:
+            # Local summary mock
+            chunks = get_all_chunks(limit=3)
+            mock_answer = "⚠️ **MOCK SUMMARY** ⚠️\n\n(API Quota Exhausted - Showing direct database content)\n\n"
+            for c in chunks:
+                # Show full text to avoid breaking markdown tables
+                mock_answer += f"### Chunk from {c['metadata']['source']}\n{c['text']}\n\n---\n\n"
+        else:
+            # Local retrieval mock
+            retrieved = search_knowledge(query)
+            mock_answer = "⚠️ **MOCK RESPONSE** ⚠️\n\n(API Quota Exhausted)\n\nI found the following relevant information in your documents using local search (Exact text match):\n\n"
+            
+            seen_sources = set()
+            for r in retrieved:
+                # Add citation
+                meta = r["metadata"]
+                if (meta["source"], meta["page"]) not in seen_sources:
+                     mock_citations.append(meta)
+                     seen_sources.add((meta["source"], meta["page"]))
+                
+                # Show full text of the relevant chunk
+                mock_answer += f"> **Source: {meta['source']}**\n\n{r['text']}\n\n---\n"
+            
+            if not retrieved:
+                mock_answer += "No relevant documents found in the local index."
+
+        response = {
+            "answer": mock_answer,
+            "confidence": 0.85,
+            "citations": mock_citations
+        }
+        answer_cache[key] = (now, response)
+        return response
+
     # ==========================
     # 🟦 SUMMARY (BYPASS AGENT)
     # ==========================
@@ -142,27 +187,48 @@ async def ask(data: PromptRequest):
     # ==========================
     # 🟩 AGENTIC RAG (LLM + EVALUATION)
     # ==========================
-    result = agentic_graph.invoke({
+    # ==========================
+    # 🟩 AGENTIC RAG (MULTI-TOOL SUPERVISOR)
+    # ==========================
+    # Initialize state for new graph
+    initial_state = {
         "messages": [],
         "query": query,
-        "refined_query": "",
-        "decision": "",
-        "retrieved_chunks": [],
-        "retrieval_quality": "",
-        "retries": 0,
-        "answer": None,
-        "confidence": 0.0,
-        "answer_known": False
-    }, config={"configurable": {"thread_id": data.thread_id}})
-
-    response = {
-        "answer": result["answer"],
-        "confidence": result["confidence"],
-        "citations": list({
-            (c["metadata"]["source"], c["metadata"]["page"]): c["metadata"]
-            for c in result.get("retrieved_chunks", [])
-        }.values())
+        "final_answer": "",
+        "next_node": "",
+        "current_tool": "",
+        "tool_outputs": [],
+        "verification_notes": "",
+        "retries": 0
     }
+    
+    try:
+        result = agentic_graph.invoke(initial_state, config={"configurable": {"thread_id": data.thread_id}})
+        
+        # Extract citations from tool outputs
+        citations = []
+        seen = set()
+        for t in result.get("tool_outputs", []):
+            src = t.get("source", "unknown")
+            # If it's a PDF, it has metadata
+            if src == "internal_pdf":
+                meta = t.get("metadata", {})
+                key_ = (meta.get("source"), meta.get("page"))
+                if key_ not in seen:
+                    citations.append(meta)
+                    seen.add(key_)
+            # If it's Web, just cite the source
+            elif src == "external_web":
+                citations.append({"source": "Tavily Web Search", "page": "Web"})
+
+        response = {
+            "answer": result.get("final_answer", "No answer produced."),
+            "confidence": 0.9 if result.get("tool_outputs") else 0.1,
+            "citations": citations
+        }
+        
+        answer_cache[key] = (now, response)
+        return response
 
-    answer_cache[key] = (now, response)
-    return response
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": f"Agent execution failed: {str(e)}"})