Update agent.py

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:
+
+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:
-        query = input("Query: ")
+            user_input = input("\nQuery: ")
         except (KeyboardInterrupt, EOFError):
             print("\nExiting.")
             break
-    if query.strip().lower() in {"exit", "quit", "q"}:
-        print("Exiting.")
+        if user_input.lower() in {"exit", "quit"}:
+            print("Goodbye!")
             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
+        response = answer_query(user_input)
+        print("\nAnswer:\n", response)