From 3f4c34e17a8e2c74ec1cc92c8c46af638ea8337a 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: Tue, 2 Jun 2026 07:46:48 +0000 Subject: [PATCH] Update agent.py --- agent.py | 146 +++++++++++++++++++++---------------------------------- 1 file changed, 56 insertions(+), 90 deletions(-) diff --git a/agent.py b/agent.py index c92d20e..1352f20 100644 --- a/agent.py +++ b/agent.py @@ -1,106 +1,72 @@ -"""Core logic for the RAG agent. +""" +Agent creation for the RAG system. -The agent decides whether to use the local knowledge base or Tavily based on a -very simple heuristic: if the query contains words like ``news``, ``latest`` -or ``today`` it is routed to the web search; otherwise the local KB is used. - -The decision logic can be replaced with a more sophisticated router if -desired. +Provides a function ``create_agent`` that returns an ``AgentExecutor`` capable of +choosing between the local KB search and the Tavily web search. """ -from __future__ import annotations +from typing import List -from typing import Tuple - -from langchain.agents import AgentExecutor, create_openai_tools_agent from langchain_ollama import ChatOllama -from langchain_core.prompts import ChatPromptTemplate +from langchain.agents import AgentExecutor, create_openai_functions_agent +from langchain.tools import Tool -from rag_tools import search_local_kb, web_search -from vectorstore import create_vectorstore - -# --------------------------------------------------------------------------- -# Prompt template -# --------------------------------------------------------------------------- -SYSTEM_PROMPT = """You are an AI assistant that can answer questions using either a local knowledge base or real‑time web search. - -When answering, always include the source of the information: -- "chromadb" for local knowledge base results. -- "tavily" for web search results. - -If you are uncertain, say "I don't know" but still mention the source you used. -""" - -USER_PROMPT = """Question: {question}\n -When you respond, first state the source (chromadb or tavily) and then provide the answer. -""" - -prompt = ChatPromptTemplate.from_messages([ - ("system", SYSTEM_PROMPT), - ("user", USER_PROMPT), -]) - -# --------------------------------------------------------------------------- -# Decision logic -# --------------------------------------------------------------------------- -WEB_KEYWORDS = {"news", "latest", "today", "current", "recent"} - - -def choose_tool(question: str) -> Tuple[str, callable]: - """Return the name of the tool and the function to call. - - Parameters - ---------- - question: str - The user query. - - Returns - ------- - Tuple[str, callable] - The tool name and the corresponding function. - """ - lowered = question.lower() - if any(word in lowered for word in WEB_KEYWORDS): - return "web_search", web_search - return "search_local_kb", search_local_kb +# Import the tools defined in tools.py +from tools import search_local_kb, web_search # --------------------------------------------------------------------------- # Agent creation # --------------------------------------------------------------------------- -def create_agent() -> AgentExecutor: - """Instantiate the agent with the two tools. +def create_agent(vectorstore_instance) -> AgentExecutor: + """Create an agent that can decide between local KB and web search. - The LLM used is Ollama's ``llama3``. + Parameters + ---------- + vectorstore_instance + Instance of the Chroma vector store to be used by the local search tool. + + Returns + ------- + AgentExecutor + Configured agent ready for use. """ - tools = [search_local_kb, web_search] + # Make the vectorstore available to the tool via the module global + import tools + tools.vectorstore = vectorstore_instance + + # Define the tools + tools_list: List[Tool] = [ + Tool( + 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 llm = ChatOllama(model="llama3", temperature=0) - # Build an agent that knows about the tools and uses the custom prompt - agent = create_openai_tools_agent(llm=llm, tools=tools, prompt=prompt) - return AgentExecutor(agent=agent, tools=tools, verbose=True) + + # System prompt guiding the agent + system_prompt = ( + "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 + agent = create_openai_functions_agent(llm=llm, tools=tools_list, system_message=system_prompt) + + # Wrap in an executor for easy use + return AgentExecutor(agent=agent, tools=tools_list, verbose=True) # --------------------------------------------------------------------------- -# CLI loop -# --------------------------------------------------------------------------- -if __name__ == "__main__": - # Ensure the vector store is loaded once - store = create_vectorstore() - # Load documents if the store is empty - if not store.get_index_info(): - from vectorstore import load_documents - load_documents("documents", store) - - agent = create_agent() - print("RAG Agent ready. Type 'exit' to quit.") - while True: - try: - question = input("\nЗапрос: ") - except EOFError: - break - if question.strip().lower() in {"exit", "quit"}: - break - # The agent will automatically call the chosen tool via the prompt. - # We simply pass the question to the agent. - result = agent.invoke({"input": question}) - # The agent's output already contains the source. - print("Ответ:", result["output"]) # noqa: T201 +# End of module +# --------------------------------------------------------------------------- \ No newline at end of file