From da7e626532ff69edb9ce3f95e0648f9fee6bf1e7 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:17:15 +0000 Subject: [PATCH] Update agent.py --- agent.py | 127 ++++++++++++++++++------------------------------------- 1 file changed, 40 insertions(+), 87 deletions(-) diff --git a/agent.py b/agent.py index 3c4bed8..0857f18 100644 --- a/agent.py +++ b/agent.py @@ -1,104 +1,57 @@ -"""Main RAG agent implementation. - -The agent can answer questions using either the local ChromaDB knowledge base -or live web search via Tavily. The decision of which tool to use is made by -the LLM itself based on the prompt. +""" +Main agent logic: decides whether to use local KB or web search. """ import os -from pathlib import Path +from typing import Dict, Any from langchain_ollama import ChatOllama -from langchain.agents import AgentExecutor, create_openai_tools_agent -from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate +from langchain.agents import initialize_agent, AgentType, Tool, AgentExecutor +from langchain_core.messages import HumanMessage -from vectorstore import create_vectorstore, load_documents from rag_tools import search_local_kb, web_search +from vectorstore import create_vectorstore, load_documents -# --------------------------------------------------------------------------- -# Configuration -# --------------------------------------------------------------------------- -VECTORSTORE_DIR = Path("./chroma_db") -DOCUMENTS_DIR = Path("./documents") +# 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) -# --------------------------------------------------------------------------- -# Initialise vector store and retriever -# --------------------------------------------------------------------------- -vectorstore = create_vectorstore(str(VECTORSTORE_DIR)) -# Load documents on first run – this is idempotent -if not any(VECTORSTORE_DIR.iterdir()): - print("Loading documents into ChromaDB…") - load_documents(str(DOCUMENTS_DIR), vectorstore) - print("Documents loaded.") +# 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."), +] -# Global retriever for tool access -vectorstore_retriever = vectorstore.as_retriever(search_kwargs={"k": 3}) +# 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)." +) -# --------------------------------------------------------------------------- -# LLM and prompt -# --------------------------------------------------------------------------- +# Create the agent executor llm = ChatOllama(model="llama3") +agent_executor = initialize_agent( + tools=tools, + llm=llm, + agent=AgentType.OPENAI_FUNCTIONS, + verbose=True, + system_message=system_prompt, +) -system_prompt = """You are an AI assistant that can answer questions using two sources: - -1. A local knowledge base (ChromaDB). Use the tool ``search_local_kb`` when the - answer can be found in the documents. -2. Live web search (Tavily). Use the tool ``web_search`` when the answer requires - up‑to‑date information. - -After retrieving the information, answer the user question and explicitly -state the source you used: either ``chromadb`` or ``tavily``. - -If you are unsure, ask for clarification. Do not provide fabricated data. -""" - -prompt = ChatPromptTemplate.from_messages([ - SystemMessagePromptTemplate.from_template(system_prompt), - HumanMessagePromptTemplate.from_template("{input}") -]) - -# --------------------------------------------------------------------------- -# Agent setup -# --------------------------------------------------------------------------- -# Tools are automatically discovered via the @tool decorator in rag_tools.py -tools = [search_local_kb, web_search] - -agent = create_openai_tools_agent(llm=llm, tools=tools, prompt=prompt) -agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True) - -# --------------------------------------------------------------------------- -# Public API -# --------------------------------------------------------------------------- - -def answer_query(query: str) -> str: - """Return the agent's answer for *query*. - - Parameters - ---------- - query: str - The user's question. - - Returns - ------- - str - The agent's response. - """ - result = agent_executor.invoke({"input": query}) - return result["output"] - -# --------------------------------------------------------------------------- -# CLI entry point -# --------------------------------------------------------------------------- -if __name__ == "__main__": - print("RAG Agent ready. Type 'exit' to quit.") +def main(): + print("Welcome to the RAG agent. Type 'exit' to quit.") while True: - try: - user_input = input("\nQuery: ") - except (KeyboardInterrupt, EOFError): - print("\nExiting.") - break + user_input = input("\nUser: ") if user_input.lower() in {"exit", "quit"}: print("Goodbye!") break - response = answer_query(user_input) - print("\nAnswer:\n", response) + # Run the agent + result = agent_executor.invoke({"input": user_input}) + # The result may contain tool calls and final answer + print("\nAssistant:", result.get("output", "")) + +if __name__ == "__main__": + main()