From 997f2acd1a9439af344de839e9adf5ef701d1273 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A0=D0=BE=D0=BC?= =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= Date: Thu, 4 Jun 2026 20:08:13 +0000 Subject: [PATCH] Update agent.py --- agent.py | 147 +++++++++++++++++-------------------------------------- 1 file changed, 44 insertions(+), 103 deletions(-) diff --git a/agent.py b/agent.py index e08d30f..43c056c 100644 --- a/agent.py +++ b/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 -`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 uses a simple tool‑based architecture. The LLM is a local +`ChatOllama` model (llama3). Two tools are available: -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 -from vectorstore import create_vectorstore +# Import tools without relative import to allow top‑level import +import tools # --------------------------------------------------------------------------- -# Helper: decide whether to use web search +# Helper: create a tool list # --------------------------------------------------------------------------- -WEB_KEYWORDS = {"news", "today", "now", "latest", "recent", "current"} - - -def should_use_web(query: str) -> bool: - words = set(query.lower().split()) - return bool(words & WEB_KEYWORDS) +def get_tools() -> List[BaseTool]: + """Return the list of tools used by the agent.""" + return [tools.search_local_kb, tools.web_search] # --------------------------------------------------------------------------- # 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)." +SYSTEM_PROMPT = ( + "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: - """Create an AgentExecutor that routes queries to the appropriate tool. +def create_agent() -> Runnable: + """Create a runnable agent. - Parameters - ---------- - vectorstore: Chroma - The vector store used by the local KB search tool. - - Returns - ------- - AgentExecutor - Configured agent. + 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. """ - # Attach the vectorstore to the local search tool via closure - def search_local_kb_wrapper(query: str, top_k: int = 3) -> str: - return search_local_kb(query, top_k=top_k, vectorstore=vectorstore) + llm = ChatOllama(model="llama3", temperature=0.0) - # Wrap the web search tool (no extra params needed) - def web_search_wrapper(query: str, top_k: int = 3) -> str: - return web_search(query, top_k=top_k) + # AgentExecutor can accept a system message via agent_kwargs + from langchain.agents import AgentExecutor - # Define tools list with updated references - tools = [search_local_kb_wrapper, web_search_wrapper] + agent = AgentExecutor.from_llm_and_tools( + 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 - 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", "")) -"" \ No newline at end of file + return agent \ No newline at end of file