From 883afee0c3ff8b443a43dc72456aa824a5c16599 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:15:38 +0000 Subject: [PATCH] Update agent.py --- agent.py | 141 ++++++++++++++++++++++++++++--------------------------- 1 file changed, 73 insertions(+), 68 deletions(-) diff --git a/agent.py b/agent.py index a7997dd..3c4bed8 100644 --- a/agent.py +++ b/agent.py @@ -1,18 +1,16 @@ -"""Main script for the RAG agent with ChromaDB and Tavily. +"""Main RAG agent implementation. -The script: -1. Loads or creates the Chroma vector store. -2. Loads documents from the `documents/` folder. -3. Sets up the LangChain agent with two tools: `search_local_kb` and `web_search`. -4. Runs a simple CLI loop. +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. """ import os from pathlib import Path from langchain_ollama import ChatOllama -from langchain.agents import initialize_agent, AgentType -from langchain.tools import Tool +from langchain.agents import AgentExecutor, create_openai_tools_agent +from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate from vectorstore import create_vectorstore, load_documents from rag_tools import search_local_kb, web_search @@ -20,80 +18,87 @@ from rag_tools import search_local_kb, web_search # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- -CHROMA_DIR = "./chroma_db" -DOCS_DIR = "./documents" -MODEL = "llama3" +VECTORSTORE_DIR = Path("./chroma_db") +DOCUMENTS_DIR = Path("./documents") # --------------------------------------------------------------------------- -# Helper: load or create vector store +# Initialise vector store and retriever # --------------------------------------------------------------------------- -vectorstore = create_vectorstore(persist_directory=CHROMA_DIR) +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.") -# Load documents – we always load; Chroma will deduplicate by ID if same content -print("Loading documents into ChromaDB (if not already present)...") -load_documents(DOCS_DIR, vectorstore) -print("Documents loaded.") +# Global retriever for tool access +vectorstore_retriever = vectorstore.as_retriever(search_kwargs={"k": 3}) # --------------------------------------------------------------------------- -# Define tools – pass the vectorstore to the local search tool +# LLM and prompt # --------------------------------------------------------------------------- -# We wrap the tool functions to include the vectorstore argument +llm = ChatOllama(model="llama3") -def local_kb_tool(query: str, top_k: int = 3): - return search_local_kb(query=query, top_k=top_k, vectorstore=vectorstore) +system_prompt = """You are an AI assistant that can answer questions using two sources: -# Create LangChain Tool objects -local_tool = Tool( - name="search_local_kb", - func=local_kb_tool, - description="Semantic search in the local knowledge base. Use when the answer is in the local documents.", -) -web_tool = Tool( - name="web_search", - func=web_search, - description="Search the web using Tavily. Use for up‑to‑date facts or news.", -) +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 # --------------------------------------------------------------------------- -llm = ChatOllama(model=MODEL, temperature=0.0) +# Tools are automatically discovered via the @tool decorator in rag_tools.py +tools = [search_local_kb, web_search] -system_prompt = ( - "You are an assistant that answers user questions. " - "If the answer is likely to be in the local knowledge base, use the tool " - "search_local_kb. If the answer requires up‑to‑date information, use the " - "web_search tool. After retrieving information, provide the answer and " - "state the source: either 'chromadb' or 'tavily'." -) - -agent = initialize_agent( - tools=[local_tool, web_tool], - llm=llm, - agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION, - verbose=True, - prefix=system_prompt, -) +agent = create_openai_tools_agent(llm=llm, tools=tools, prompt=prompt) +agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True) # --------------------------------------------------------------------------- -# CLI loop +# Public API # --------------------------------------------------------------------------- -print("\nRAG Agent ready. Type your question (or 'exit' to quit).\n") -while True: - try: - query = input("Query: ") - except (KeyboardInterrupt, EOFError): - print("\nExiting.") - break - if query.strip().lower() in {"exit", "quit", "q"}: - print("Exiting.") - break - if not query.strip(): - continue - # Run the agent - try: - result = agent.run(query) - print(f"\nAnswer:\n{result}\n") - except Exception as e: - print(f"Error: {e}") - continue + +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.") + while True: + try: + user_input = input("\nQuery: ") + except (KeyboardInterrupt, EOFError): + print("\nExiting.") + break + if user_input.lower() in {"exit", "quit"}: + print("Goodbye!") + break + response = answer_query(user_input) + print("\nAnswer:\n", response)