From fd22bd87199247ca7d6dcac6ebf8940e5fa82e6c 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?= <kirilinka320@gmail.com>
Date: Tue, 2 Jun 2026 07:18:49 +0000
Subject: [PATCH] Update agent.py

---
 agent.py | 107 ++++++++++++++++++++++++++++++++++---------------------
 1 file changed, 67 insertions(+), 40 deletions(-)

diff --git a/agent.py b/agent.py
index 0857f18..2599ed7 100644
--- a/agent.py
+++ b/agent.py
@@ -1,57 +1,84 @@
-"""
-Main agent logic: decides whether to use local KB or web search.
+"""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.
 """
 
 import os
-from typing import Dict, Any
+from typing import List
 
 from langchain_ollama import ChatOllama
-from langchain.agents import initialize_agent, AgentType, Tool, AgentExecutor
-from langchain_core.messages import HumanMessage
+from langchain.agents import tool, AgentExecutor, ZeroShotAgent
+from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate
 
-from rag_tools import search_local_kb, web_search
 from vectorstore import create_vectorstore, load_documents
+from rag_tools import web_search
 
-# Load or create vector store
-vectorstore = create_vectorstore()
-# Load documents from the documents folder if not already loaded
-if not vectorstore._collection.count():  # type: ignore[attr-defined]
-    load_documents("./documents", vectorstore)
+# ---------------------------------------------------------------------------
+# 1. Setup vector store
+# ---------------------------------------------------------------------------
+VECTORSTORE_DIR = "./chroma_db"
+DOCS_DIR = "./documents"
 
-# Define tools
-tools = [
-    Tool(name="search_local_kb", func=search_local_kb, description="Search the local knowledge base."),
-    Tool(name="web_search", func=web_search, description="Search the web using Tavily."),
-]
+vectorstore = create_vectorstore(VECTORSTORE_DIR)
+load_documents(DOCS_DIR, vectorstore)
 
-# System prompt to guide the agent
-system_prompt = (
-    "You are an AI assistant. For questions about local documents use the 'search_local_kb' tool. "
-    "For recent news or facts not in the local docs, use 'web_search'. "
-    "Always indicate the source of your answer (chromadb or tavily)."
+# ---------------------------------------------------------------------------
+# 2. Define local search tool (needs the vectorstore)
+# ---------------------------------------------------------------------------
+@tool
+def search_local_kb(query: str, top_k: int = 3) -> str:
+    """Semantic search over the local ChromaDB collection.
+
+    Returns a formatted string containing the top_k snippets and a source tag.
+    """
+    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"
+
+# ---------------------------------------------------------------------------
+# 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 executor
-llm = ChatOllama(model="llama3")
-agent_executor = initialize_agent(
+# 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,
-    llm=llm,
-    agent=AgentType.OPENAI_FUNCTIONS,
-    verbose=True,
-    system_message=system_prompt,
+    prompt=agent_prompt,
 )
 
-def main():
-    print("Welcome to the RAG agent. Type 'exit' to quit.")
-    while True:
-        user_input = input("\nUser: ")
-        if user_input.lower() in {"exit", "quit"}:
-            print("Goodbye!")
-            break
-        # Run the agent
-        result = agent_executor.invoke({"input": user_input})
-        # The result may contain tool calls and final answer
-        print("\nAssistant:", result.get("output", ""))
+agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
+# ---------------------------------------------------------------------------
+# 4. Chat loop
+# ---------------------------------------------------------------------------
 if __name__ == "__main__":
-    main()
+    print("RAG Agent ready. Type 'exit' to quit.")
+    while True:
+        user_input = input("Запрос: ")
+        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}")