Update agent.py

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"
+MODEL_NAME = "llama3"
+CHROMA_DIR = "./chroma_db"
+DOCS_DIR = "documents"
 
-vectorstore = create_vectorstore(VECTORSTORE_DIR)
+# ---------------------------------------------------------------------------
+# Load or create vector store
+# ---------------------------------------------------------------------------
+vectorstore = create_vectorstore(CHROMA_DIR)
+# Load documents only if the store is empty
+if not vectorstore._collection.count():
     load_documents(DOCS_DIR, vectorstore)
 
 # ---------------------------------------------------------------------------
-# 2. Define local search tool (needs the vectorstore)
+# Define tools
 # ---------------------------------------------------------------------------
-@tool
-def search_local_kb(query: str, top_k: int = 3) -> str:
-    """Semantic search over the local ChromaDB collection.
+class LocalKBTool(BaseTool):
+    name = "search_local_kb"
+    description = "Perform a semantic search in the local knowledge base."
 
-    Returns a formatted string containing the top_k snippets and a source tag.
+    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: <chromadb|tavily>
 """
-    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)."
-        ),
-    ]
+PROMPT = PromptTemplate(
+    input_variables=["input", "chat_history"],
+    template="""
+{chat_history}
+User: {input}
+Assistant: """
 )
 
-# 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)
+# ---------------------------------------------------------------------------
+# Agent chain
+# ---------------------------------------------------------------------------
+llm = ChatOllama(model=MODEL_NAME, temperature=0.2)
+chain = LLMChain(llm=llm, prompt=PROMPT)
 
 # ---------------------------------------------------------------------------
-# 4. Chat loop
+# 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
+    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}
+
+# ---------------------------------------------------------------------------
+# 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']}")