Update agent.py

agent.py

@@ -1,72 +1,128 @@
-"""
+"""Agent construction for RAG with local KB and web search.
-Agent creation for the RAG system.
 
-Provides a function ``create_agent`` that returns an ``AgentExecutor`` capable of
+The agent uses LangChain's `create_openai_functions_agent` style with a custom
-choosing between the local KB search and the Tavily web search.
+`AgentExecutor` that routes queries to either the local knowledge base or the
+web search based on a simple heuristic: if the query contains words that
+suggest recent events ("news", "today", "now", "latest"), we use the web
+search; otherwise we use the local KB.
+
+The agent returns the answer along with the source used.
 """
 
-from typing import List
+from typing import Dict, Any
 
+from langchain_core.prompts import ChatPromptTemplate
 from langchain_ollama import ChatOllama
-from langchain.agents import AgentExecutor, create_openai_functions_agent
+from langchain.agents import create_openai_functions_agent, AgentExecutor
-from langchain.tools import Tool
 
-# Import the tools defined in tools.py
+from rag_tools import search_local_kb, web_search
-from tools import search_local_kb, web_search
+from vectorstore import create_vectorstore
 
 # ---------------------------------------------------------------------------
-# Agent creation
+# Helper: decide whether to use web search
 # ---------------------------------------------------------------------------
 
-def create_agent(vectorstore_instance) -> AgentExecutor:
+WEB_KEYWORDS = {"news", "today", "now", "latest", "recent", "current"}
-    """Create an agent that can decide between local KB and web search.
+
+
+def should_use_web(query: str) -> bool:
+    words = set(query.lower().split())
+    return bool(words & WEB_KEYWORDS)
+
+# ---------------------------------------------------------------------------
+# System prompt
+# ---------------------------------------------------------------------------
+
+SYSTEM_PROMPT = "You are an assistant that answers questions. Use the local knowledge base when the question is about stored documents; otherwise use web search. Respond with the answer and indicate the source (chromadb or tavily)."
+
+# ---------------------------------------------------------------------------
+# Agent construction
+# ---------------------------------------------------------------------------
+
+def create_agent(vectorstore) -> AgentExecutor:
+    """Create an AgentExecutor that routes queries to the appropriate tool.
 
     Parameters
     ----------
-    vectorstore_instance
+    vectorstore: Chroma
-        Instance of the Chroma vector store to be used by the local search tool.
+        The vector store used by the local KB search tool.
 
     Returns
     -------
     AgentExecutor
-        Configured agent ready for use.
+        Configured agent.
     """
-    # Make the vectorstore available to the tool via the module global
+    # Attach the vectorstore to the local search tool via closure
-    import tools
+    def search_local_kb_wrapper(query: str, top_k: int = 3) -> str:
-    tools.vectorstore = vectorstore_instance
+        return search_local_kb(query, top_k=top_k, vectorstore=vectorstore)
 
-    # Define the tools
+    # Wrap the web search tool (no extra params needed)
-    tools_list: List[Tool] = [
+    def web_search_wrapper(query: str, top_k: int = 3) -> str:
-        Tool(
+        return web_search(query, top_k=top_k)
-            name="search_local_kb",
-            func=search_local_kb,
-            description="Search the local knowledge base (ChromaDB). Use when the answer is likely contained in the local documents.",
-        ),
-        Tool(
-            name="web_search",
-            func=web_search,
-            description="Search the web via Tavily. Use when the answer requires up‑to‑date information.",
-        ),
-    ]
 
-    # LLM for the agent
+    # Define tools list with updated references
-    llm = ChatOllama(model="llama3", temperature=0)
+    tools = [search_local_kb_wrapper, web_search_wrapper]
 
-    # System prompt guiding the agent
+    # LLM
-    system_prompt = (
+    llm = ChatOllama(model="llama3")
-        "You are an assistant that answers user questions. "
-        "If the answer can be found in the local knowledge base, use the tool "
-        "`search_local_kb`. If the question asks for recent or current information, "
-        "use the tool `web_search`. After obtaining the information, provide a "
-        "concise answer and state the source (`chromadb` or `tavily`)."
-    )
 
-    # Create the agent using the function calling approach
+    # Prompt template
-    agent = create_openai_functions_agent(llm=llm, tools=tools_list, system_message=system_prompt)
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", SYSTEM_PROMPT),
+        ("human", "{input}"),
+    ])
 
-    # Wrap in an executor for easy use
+    # Create agent
-    return AgentExecutor(agent=agent, tools=tools_list, verbose=True)
+    agent = create_openai_functions_agent(llm=llm, tools=tools, prompt=prompt)
+
+    # Wrap with AgentExecutor
+    return AgentExecutor(agent=agent, tools=tools, verbose=True)
 
 # ---------------------------------------------------------------------------
-# End of module
+# Simple executor for tests
-# ---------------------------------------------------------------------------
+# ---------------------------------------------------------------------------
+
+def create_agent_executor(vectorstore) -> Any:
+    """Return a simple callable that mimics the agent for testing.
+
+    The returned callable takes a dictionary with keys ``input`` and optionally
+    ``tool_choice``. It selects the appropriate tool based on the query and
+    returns a dictionary with an ``output`` key containing the result.
+    """
+    agent = create_agent(vectorstore)
+
+    def executor(inputs: Dict[str, Any]) -> Dict[str, Any]:
+        query = inputs.get("input", "")
+        tool_choice = inputs.get("tool_choice")
+        if tool_choice == "web_search":
+            result = web_search(query)
+        elif tool_choice == "search_local_kb":
+            result = search_local_kb(query, vectorstore=vectorstore)
+        else:
+            # Default routing
+            if should_use_web(query):
+                result = web_search(query)
+            else:
+                result = search_local_kb(query, vectorstore=vectorstore)
+        return {"output": result}
+
+    return executor
+
+# ---------------------------------------------------------------------------
+# Example usage
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    store = create_vectorstore()
+    agent = create_agent(store)
+    while True:
+        q = input("Query> ")
+        if q.lower() in {"exit", "quit", "q"}:
+            break
+        # Simple routing: if query contains web keywords, use web tool
+        if should_use_web(q):
+            result = agent.invoke({"input": q, "tool_choice": "web_search"})
+        else:
+            result = agent.invoke({"input": q, "tool_choice": "search_local_kb"})
+        print("Answer:", result.get("output", ""))
+""