Update agent.py

agent.py

@@ -1,128 +1,69 @@
-"""Agent construction for RAG with local KB and web search.
+"""Agent implementation using LangChain 1.x.
 
-The agent uses LangChain's `create_openai_functions_agent` style with a custom
+The agent uses a simple tool‑based architecture.  The LLM is a local
-`AgentExecutor` that routes queries to either the local knowledge base or the
+`ChatOllama` model (llama3).  Two tools are available:
-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.
+* ``search_local_kb`` – semantic search in Qdrant.
+* ``web_search`` – web search via Tavily.
+
+The system prompt instructs the LLM to decide which tool to use based on the
+question.  The response always contains a marker indicating the source
+(`chromadb`/`tavily`).  The marker is added by the LLM itself.
 """
 
-from typing import Dict, Any
+from __future__ import annotations
+
+from typing import List
 
-from langchain_core.prompts import ChatPromptTemplate
 from langchain_ollama import ChatOllama
-from langchain.agents import create_openai_functions_agent, AgentExecutor
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import Runnable
+from langchain_core.tools import BaseTool
 
-from rag_tools import search_local_kb, web_search
+# Import tools without relative import to allow top‑level import
-from vectorstore import create_vectorstore
+import tools
 
 # ---------------------------------------------------------------------------
-# Helper: decide whether to use web search
+# Helper: create a tool list
 # ---------------------------------------------------------------------------
 
-WEB_KEYWORDS = {"news", "today", "now", "latest", "recent", "current"}
+def get_tools() -> List[BaseTool]:
-
+    """Return the list of tools used by the agent."""
-
+    return [tools.search_local_kb, tools.web_search]
-def should_use_web(query: str) -> bool:
-    words = set(query.lower().split())
-    return bool(words & WEB_KEYWORDS)
 
 # ---------------------------------------------------------------------------
 # System prompt
 # ---------------------------------------------------------------------------
-
+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)."
+    "You are an AI assistant that can answer questions using two sources: "
+    "1) a local knowledge base (Qdrant) and 2) the web via Tavily. "
+    "If the answer can be found in the local KB, use the `search_local_kb` tool. "
+    "If the answer requires up‑to‑date information, use the `web_search` tool. "
+    "Return the answer followed by a source marker on a new line: "
+    "`Source: chromadb` or `Source: tavily`."
+)
 
 # ---------------------------------------------------------------------------
 # Agent construction
 # ---------------------------------------------------------------------------
 
-def create_agent(vectorstore) -> AgentExecutor:
+def create_agent() -> Runnable:
-    """Create an AgentExecutor that routes queries to the appropriate tool.
+    """Create a runnable agent.
 
-    Parameters
+    The agent is a simple chain: system prompt → user message → tool calls → LLM
-    ----------
+    response.  It uses the default tool‑calling behaviour of LangChain 1.x.
-    vectorstore: Chroma
-        The vector store used by the local KB search tool.
-
-    Returns
-    -------
-    AgentExecutor
-        Configured agent.
     """
-    # Attach the vectorstore to the local search tool via closure
+    llm = ChatOllama(model="llama3", temperature=0.0)
-    def search_local_kb_wrapper(query: str, top_k: int = 3) -> str:
-        return search_local_kb(query, top_k=top_k, vectorstore=vectorstore)
 
-    # Wrap the web search tool (no extra params needed)
+    # AgentExecutor can accept a system message via agent_kwargs
-    def web_search_wrapper(query: str, top_k: int = 3) -> str:
+    from langchain.agents import AgentExecutor
-        return web_search(query, top_k=top_k)
 
-    # Define tools list with updated references
+    agent = AgentExecutor.from_llm_and_tools(
-    tools = [search_local_kb_wrapper, web_search_wrapper]
+        llm=llm,
+        tools=get_tools(),
+        verbose=True,
+        # Provide the system prompt so the LLM knows how to behave
+        agent_kwargs={"system_message": SYSTEM_PROMPT},
+    )
 
-    # LLM
+    return agent
-    llm = ChatOllama(model="llama3")
-
-    # Prompt template
-    prompt = ChatPromptTemplate.from_messages([
-        ("system", SYSTEM_PROMPT),
-        ("human", "{input}"),
-    ])
-
-    # Create agent
-    agent = create_openai_functions_agent(llm=llm, tools=tools, prompt=prompt)
-
-    # Wrap with AgentExecutor
-    return AgentExecutor(agent=agent, tools=tools, verbose=True)
-
-# ---------------------------------------------------------------------------
-# 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", ""))
-""