Update agent.py

agent.py

@@ -1,97 +1,106 @@
 """Core logic for the RAG agent.
 
-The agent decides whether to use the local knowledge base or Tavily based on simple heuristics:
+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.
 
-- If the query contains words that usually refer to recent events (e.g. "новости", "актуальные", "сегодня", "сейчас"), we route to Tavily.
+The decision logic can be replaced with a more sophisticated router if
-- Otherwise we assume the answer can be found in the local KB.
+desired.
-
-The agent uses LangChain's :class:`langchain.agents.AgentExecutor` with a custom prompt that
-instructs the LLM to specify the source in the response.
 """
 
-import os
+from __future__ import annotations
-from typing import Dict, Any
+
+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 rag_tools import search_local_kb, web_search
-from vectorstore import create_vectorstore, load_documents
+from vectorstore import create_vectorstore
 
 # ---------------------------------------------------------------------------
-# Configuration
+# Prompt template
 # ---------------------------------------------------------------------------
-# Directory containing the documents to index
+SYSTEM_PROMPT = """You are an AI assistant that can answer questions using either a local knowledge base or real‑time web search.
-DOCUMENTS_DIR = "./documents"
-# Persistence directory for Chroma
-CHROMA_DIR = "./chroma_db"
-# The LLM used by the agent
-LLM = ChatOllama(model="llama3")
 
-# ---------------------------------------------------------------------------
+When answering, always include the source of the information:
-# Helper functions
+- "chromadb" for local knowledge base results.
-# ---------------------------------------------------------------------------
+- "tavily" for web search results.
 
-def should_use_web(query: str) -> bool:
+If you are uncertain, say "I don't know" but still mention the source you used.
-    """Return True if the query looks like it needs up‑to‑date information.
+"""
 
-    The heuristic looks for a few Russian keywords that usually indicate a
+USER_PROMPT = """Question: {question}\n
-    request for recent news.
+When you respond, first state the source (chromadb or tavily) and then provide the answer.
-    """
+"""
-    keywords = ["новости", "актуальные", "сегодня", "сейчас", "текущие", "текущий"]
-    lowered = query.lower()
-    return any(k in lowered for k in keywords)
 
-# ---------------------------------------------------------------------------
-# Agent setup
-# ---------------------------------------------------------------------------
-# Load or create the vector store
-vectorstore = create_vectorstore(persist_directory=CHROMA_DIR)
-# If the store is empty, load documents from the documents directory
-if not vectorstore._collection.count():  # type: ignore[attr-defined]
-    load_documents(DOCUMENTS_DIR, vectorstore)
-
-# Tools – we pass the vectorstore instance to the local search tool via a closure
-local_kb_tool = search_local_kb
-local_kb_tool.__globals__["vectorstore"] = vectorstore  # inject vectorstore
-
-TOOLS = [local_kb_tool, web_search]
-
-# Prompt template – the LLM is instructed to specify the source.
 prompt = ChatPromptTemplate.from_messages([
-    ("system", "You are an AI assistant that answers user questions.") ,
+    ("system", SYSTEM_PROMPT),
-    ("user", "{input}"),
+    ("user", USER_PROMPT),
 ])
 
-# Agent – we use the simple tool‑calling agent
+# ---------------------------------------------------------------------------
-agent = create_openai_tools_agent(llm=LLM, tools=TOOLS, prompt=prompt)
+# Decision logic
-executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True)
+# ---------------------------------------------------------------------------
+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
 
 # ---------------------------------------------------------------------------
-# CLI
+# Agent creation
 # ---------------------------------------------------------------------------
 
-def main() -> None:
+def create_agent() -> AgentExecutor:
+    """Instantiate the agent with the two tools.
+
+    The LLM used is Ollama's ``llama3``.
+    """
+    tools = [search_local_kb, web_search]
+    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)
+
+# ---------------------------------------------------------------------------
+# 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:
-        user_query = input("\nЗапрос: ")
+        try:
-        if user_query.lower() in {"exit", "quit", "q"}:
+            question = input("\nЗапрос: ")
+        except EOFError:
             break
-        # Decide which tool to use
+        if question.strip().lower() in {"exit", "quit"}:
-        if should_use_web(user_query):
+            break
-            # Explicitly call the web search tool
+        # The agent will automatically call the chosen tool via the prompt.
-            result = web_search(user_query)
+        # We simply pass the question to the agent.
-            source = "tavily"
+        result = agent.invoke({"input": question})
-        else:
+        # The agent's output already contains the source.
-            result = search_local_kb(user_query, vectorstore=vectorstore)
+        print("Ответ:", result["output"])  # noqa: T201
-            source = "chromadb"
-        # Ask the LLM to format the answer
-        formatted = LLM.invoke({"input": user_query + "\n\nAnswer: " + result})
-        print("\nОтвет:", formatted.content)
-        print("Источник:", source)
-
-if __name__ == "__main__":
-    main()
-
-# End of agent.py