Update agent.py

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()