From b9d63e9251e0744bd03fa0ac5c41c0c89f5bf1d2 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:20:21 +0000 Subject: [PATCH] Update agent.py --- agent.py | 144 ++++++++++++++++++++++++++++++++----------------------- 1 file changed, 83 insertions(+), 61 deletions(-) diff --git a/agent.py b/agent.py index 2599ed7..12d8826 100644 --- a/agent.py +++ b/agent.py @@ -1,84 +1,106 @@ -"""Main agent logic. - -Creates a Chroma vector store, loads documents from the ``documents`` directory, -and runs a simple chat loop. The agent decides whether to use the local KB -or perform a web search based on the presence of the word "news" or -"latest" in the query. +""" +Main agent implementation using LangChain. """ import os -from typing import List +from typing import Dict +from langchain import LLMChain, PromptTemplate from langchain_ollama import ChatOllama -from langchain.agents import tool, AgentExecutor, ZeroShotAgent -from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate +from langchain_core.messages import HumanMessage, SystemMessage +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 -from rag_tools import web_search # --------------------------------------------------------------------------- -# 1. Setup vector store +# Configuration # --------------------------------------------------------------------------- -VECTORSTORE_DIR = "./chroma_db" -DOCS_DIR = "./documents" - -vectorstore = create_vectorstore(VECTORSTORE_DIR) -load_documents(DOCS_DIR, vectorstore) +MODEL_NAME = "llama3" +CHROMA_DIR = "./chroma_db" +DOCS_DIR = "documents" # --------------------------------------------------------------------------- -# 2. Define local search tool (needs the vectorstore) +# Load or create vector store # --------------------------------------------------------------------------- -@tool -def search_local_kb(query: str, top_k: int = 3) -> str: - """Semantic search over the local ChromaDB collection. +vectorstore = create_vectorstore(CHROMA_DIR) +# Load documents only if the store is empty +if not vectorstore._collection.count(): + load_documents(DOCS_DIR, vectorstore) - Returns a formatted string containing the top_k snippets and a source tag. +# --------------------------------------------------------------------------- +# Define tools +# --------------------------------------------------------------------------- +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 + 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. """ - retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) - docs = retriever.invoke(query) - snippets = "\n".join([f"{idx+1}. {doc.page_content[:200]}" for idx, doc in enumerate(docs)]) - return f"[Local KB]\n{snippets}\nSource: chromadb" + # 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} # --------------------------------------------------------------------------- -# 3. Agent prompt and execution -# --------------------------------------------------------------------------- -# The agent will be given two tools: search_local_kb and web_search. -# We provide a simple instruction to choose the appropriate tool. - -agent_prompt = ChatPromptTemplate.from_messages( - [ - HumanMessagePromptTemplate.from_template( - "You are an assistant that can search a local knowledge base or the web. " - "If the question is about recent events or news, use web_search. " - "Otherwise, use search_local_kb. " - "Respond with the answer and the source (chromadb or tavily)." - ), - ] -) - -# Create the agent with the two tools -tools = [search_local_kb, web_search] - -agent = ZeroShotAgent.from_llm_and_tools( - llm=ChatOllama(model="llama3", temperature=0), - tools=tools, - prompt=agent_prompt, -) - -agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True) - -# --------------------------------------------------------------------------- -# 4. Chat loop +# CLI loop # --------------------------------------------------------------------------- if __name__ == "__main__": print("RAG Agent ready. Type 'exit' to quit.") while True: - user_input = input("Запрос: ") + user_input = input("\nЗапрос: ") if user_input.lower() in {"exit", "quit", "q"}: - print("Bye!") break - try: - result = agent_executor.invoke({"input": user_input}) - print(result["output"]) - except Exception as e: - print(f"Error: {e}") + output = decide_and_run(user_input) + print(f"\nОтвет:\n{output['answer']}") + print(f"Источник: {output['source']}")