Update agent.py

agent.py

@@ -1,15 +1,20 @@
-"""
+"""Core logic for the RAG agent.
-Main agent implementation using LangChain.
+
+The agent decides whether to use the local knowledge base or Tavily based on simple heuristics:
+
+- If the query contains words that usually refer to recent events (e.g. "новости", "актуальные", "сегодня", "сейчас"), we route to Tavily.
+- Otherwise we assume the answer can be found in the local KB.
+
+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 typing import Dict
+from typing import Dict, Any
 
-from langchain import LLMChain, PromptTemplate
+from langchain.agents import AgentExecutor, create_openai_tools_agent
 from langchain_ollama import ChatOllama
-from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.runnables import RunnableConfig
-from langchain_core.tools import BaseTool
 
 from rag_tools import search_local_kb, web_search
 from vectorstore import create_vectorstore, load_documents
@@ -17,90 +22,76 @@ from vectorstore import create_vectorstore, load_documents
 # ---------------------------------------------------------------------------
 # Configuration
 # ---------------------------------------------------------------------------
-MODEL_NAME = "llama3"
+# Directory containing the documents to index
+DOCUMENTS_DIR = "./documents"
+# Persistence directory for Chroma
 CHROMA_DIR = "./chroma_db"
-DOCS_DIR = "documents"
+# The LLM used by the agent
+LLM = ChatOllama(model="llama3")
 
 # ---------------------------------------------------------------------------
-# Load or create vector store
+# Helper functions
 # ---------------------------------------------------------------------------
-vectorstore = create_vectorstore(CHROMA_DIR)
-# Load documents only if the store is empty
-if not vectorstore._collection.count():
-    load_documents(DOCS_DIR, vectorstore)
 
-# ---------------------------------------------------------------------------
+def should_use_web(query: str) -> bool:
-# Define tools
+    """Return True if the query looks like it needs up‑to‑date information.
-# ---------------------------------------------------------------------------
-class LocalKBTool(BaseTool):
-    name = "search_local_kb"
-    description = "Perform a semantic search in the local knowledge base."
 
-    def _run(self, query: str, top_k: int = 3) -> str:  # pragma: no cover
+    The heuristic looks for a few Russian keywords that usually indicate a
-        return search_local_kb(query, top_k, vectorstore)
+    request for recent news.
-
-class WebSearchTool(BaseTool):
-    name = "web_search"
-    description = "Search the web using Tavily."
-
-    def _run(self, query: str, top_k: int = 3) -> str:  # pragma: no cover
-        return web_search(query, top_k)
-
-tools = [LocalKBTool(), WebSearchTool()]
-
-# ---------------------------------------------------------------------------
-# Prompt template
-# ---------------------------------------------------------------------------
-SYSTEM_PROMPT = """
-You are an AI assistant that answers user questions.
-- If the answer can be found in the local knowledge base, use the tool `search_local_kb`.
-- If the answer requires up‑to‑date information, use the tool `web_search`.
-After providing the answer, always state the source in the format:
-
-Source: <chromadb|tavily>
-"""
-
-PROMPT = PromptTemplate(
-    input_variables=["input", "chat_history"],
-    template="""
-{chat_history}
-User: {input}
-Assistant: """
-)
-
-# ---------------------------------------------------------------------------
-# Agent chain
-# ---------------------------------------------------------------------------
-llm = ChatOllama(model=MODEL_NAME, temperature=0.2)
-chain = LLMChain(llm=llm, prompt=PROMPT)
-
-# ---------------------------------------------------------------------------
-# Helper to decide which tool to use
-# ---------------------------------------------------------------------------
-def decide_and_run(query: str) -> Dict[str, str]:
-    """Use the LLM to decide whether to use local KB or web search.
-
-    Returns a dict with keys: answer, source.
     """
-    # Simple heuristic: if the query contains words like "news", "latest", "today" use web
+    keywords = ["новости", "актуальные", "сегодня", "сейчас", "текущие", "текущий"]
-    web_keywords = {"news", "latest", "today", "current", "recent", "update"}
+    lowered = query.lower()
-    if any(word in query.lower() for word in web_keywords):
+    return any(k in lowered for k in keywords)
-        result = web_search(query)
-        source = "tavily"
-    else:
-        result = search_local_kb(query, vectorstore=vectorstore)
-        source = "chromadb"
-    return {"answer": result, "source": source}
 
 # ---------------------------------------------------------------------------
-# CLI loop
+# Agent setup
 # ---------------------------------------------------------------------------
-if __name__ == "__main__":
+# 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.") ,
+    ("user", "{input}"),
+])
+
+# Agent – we use the simple tool‑calling agent
+agent = create_openai_tools_agent(llm=LLM, tools=TOOLS, prompt=prompt)
+executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True)
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+def main() -> None:
     print("RAG Agent ready. Type 'exit' to quit.")
     while True:
-        user_input = input("\nЗапрос: ")
+        user_query = input("\nЗапрос: ")
-        if user_input.lower() in {"exit", "quit", "q"}:
+        if user_query.lower() in {"exit", "quit", "q"}:
             break
-        output = decide_and_run(user_input)
+        # Decide which tool to use
-        print(f"\nОтвет:\n{output['answer']}")
+        if should_use_web(user_query):
-        print(f"Источник: {output['source']}")
+            # Explicitly call the web search tool
+            result = web_search(user_query)
+            source = "tavily"
+        else:
+            result = search_local_kb(user_query, vectorstore=vectorstore)
+            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