From 6acc2077e582cd5c0c6b97565240deb2fae8f8ea 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:21:51 +0000 Subject: [PATCH] Update agent.py --- agent.py | 153 ++++++++++++++++++++++++++----------------------------- 1 file changed, 72 insertions(+), 81 deletions(-) diff --git a/agent.py b/agent.py index 12d8826..0e08605 100644 --- a/agent.py +++ b/agent.py @@ -1,15 +1,20 @@ -""" -Main agent implementation using LangChain. +"""Core logic for the RAG agent. + +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.runnables import RunnableConfig -from langchain_core.tools import BaseTool +from langchain_core.prompts import ChatPromptTemplate 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) -# --------------------------------------------------------------------------- -# Define tools -# --------------------------------------------------------------------------- -class LocalKBTool(BaseTool): - name = "search_local_kb" - description = "Perform a semantic search in the local knowledge base." +def should_use_web(query: str) -> bool: + """Return True if the query looks like it needs up‑to‑date information. - def _run(self, query: str, top_k: int = 3) -> str: # pragma: no cover - return search_local_kb(query, top_k, vectorstore) - -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: -""" - -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. + The heuristic looks for a few Russian keywords that usually indicate a + request for recent news. """ - # Simple heuristic: if the query contains words like "news", "latest", "today" use web - web_keywords = {"news", "latest", "today", "current", "recent", "update"} - if any(word in query.lower() for word in web_keywords): - result = web_search(query) - source = "tavily" - else: - result = search_local_kb(query, vectorstore=vectorstore) - source = "chromadb" - return {"answer": result, "source": source} + keywords = ["новости", "актуальные", "сегодня", "сейчас", "текущие", "текущий"] + lowered = query.lower() + return any(k in lowered for k in keywords) # --------------------------------------------------------------------------- -# 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Запрос: ") - if user_input.lower() in {"exit", "quit", "q"}: + user_query = input("\nЗапрос: ") + if user_query.lower() in {"exit", "quit", "q"}: break - output = decide_and_run(user_input) - print(f"\nОтвет:\n{output['answer']}") - print(f"Источник: {output['source']}") + # Decide which tool to use + if should_use_web(user_query): + # 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