Update agent.py

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
+Provides a function ``create_agent`` that returns an ``AgentExecutor`` capable of
-very simple heuristic: if the query contains words like ``news``, ``latest``
+choosing between the local KB search and the Tavily web search.
-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.
 """
 
-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
+# Import the tools defined in tools.py
-from vectorstore import create_vectorstore
+from tools import search_local_kb, web_search
-
-# ---------------------------------------------------------------------------
-# 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
 
 # ---------------------------------------------------------------------------
 # Agent creation
 # ---------------------------------------------------------------------------
 
-def create_agent() -> AgentExecutor:
+def create_agent(vectorstore_instance) -> AgentExecutor:
-    """Instantiate the agent with the two tools.
+    """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)
+    # System prompt guiding the agent
-    return AgentExecutor(agent=agent, tools=tools, verbose=True)
+    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
+# End of module
-# ---------------------------------------------------------------------------
+# ---------------------------------------------------------------------------
-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