Update main.py

main.py

@@ -1,80 +1,82 @@
-"""Simple CLI for the RAG agent.
+"""RAG agent CLI.
 
-Commands:
-  /add <directory>   – Load all .txt/.md files from the directory into the local KB.
-  /search <question> – Ask the agent a question.
-  /quit              – Exit the program.
+This script demonstrates a simple chat loop with a LangChain agent that
+searches either a local Chroma vector store or the web via Tavily.  The
+agent automatically decides which tool to use based on the user query.
+
+Prerequisites:
+  * Ollama must be running locally with the ``llama3`` model and the
+    ``nomic-embed-text`` embedding model.
+  * A valid Tavily API key must be set in the environment variable
+    ``TAVILY_API_KEY``.
+  * The ``documents`` directory should contain the source text files.
 """
 
 import os
-import sys
 from pathlib import Path
 
 from langchain_ollama import ChatOllama
-from dotenv import load_dotenv
+from langchain.agents import create_agent
 
 from vectorstore import create_vectorstore, load_documents
-from agent import create_agent, should_use_web
+from tools import search_local_kb, web_search
 
-# Load environment variables (TAVILY_API_KEY)
-load_dotenv()
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+CHROMA_DIR = "./chroma_db"
+DOCS_DIR = "./documents"
 
-# Create or load vector store
-VECTORSTORE_DIR = "./chroma_db"
-vectorstore = create_vectorstore(persist_directory=VECTORSTORE_DIR)
+# ---------------------------------------------------------------------------
+# Initialise vector store and load documents
+# ---------------------------------------------------------------------------
+print("Initializing Chroma vector store…")
+vectorstore = create_vectorstore(persist_directory=CHROMA_DIR)
+print("Loading documents…")
+load_documents(DOCS_DIR, vectorstore)
 
-# Create agent
-agent = create_agent(vectorstore)
-
-# Helper to print usage
-USAGE = (
-    "Commands:\n"
-    "  /add <directory>   – Load documents into the local knowledge base.\n"
-    "  /search <question> – Ask the agent a question.\n"
-    "  /quit              – Exit the program.\n"
+# ---------------------------------------------------------------------------
+# Agent setup
+# ---------------------------------------------------------------------------
+# System prompt that tells the model how to choose a tool.
+SYSTEM_PROMPT = (
+    "You are an AI assistant that can answer questions using two tools. "
+    "If the answer requires up‑to‑date information, use the web_search tool. "
+    "Otherwise, use the search_local_kb tool. "
+    "When you call a tool, the tool will return the answer. "
+    "Respond with the final answer and include the source tag (chromadb or tavily)."
 )
 
-print("RAG Agent CLI. Type /help for commands.")
+llm = ChatOllama(model="llama3", temperature=0)
 
+# Tools list
+TOOLS = [search_local_kb, web_search]
+
+agent = create_agent(
+    model=llm,
+    tools=TOOLS,
+    system_prompt=SYSTEM_PROMPT,
+)
+
+# ---------------------------------------------------------------------------
+# Chat loop
+# ---------------------------------------------------------------------------
+print("\n--- RAG Agent CLI ---")
+print("Type 'exit' or 'quit' to end.")
 while True:
+    user_input = input("\nUser: ")
+    if user_input.lower() in {"exit", "quit", "q"}:
+        print("Goodbye!")
+        break
+    # Invoke the agent
     try:
-        line = input("> ").strip()
-    except (EOFError, KeyboardInterrupt):
-        print("\nExiting.")
-        break
-    if not line:
-        continue
-    if line.lower() == "/help":
-        print(USAGE)
-        continue
-    if line.lower() == "/quit":
-        print("Bye!")
-        break
-    if line.lower().startswith("/add "):
-        dir_path = line[5:].strip()
-        if not dir_path:
-            print("Please provide a directory path.")
-            continue
-        if not Path(dir_path).exists():
-            print(f"Directory {dir_path} does not exist.")
-            continue
-        load_documents(dir_path, vectorstore)
-        print("Documents loaded.")
-        continue
-    if line.lower().startswith("/search "):
-        query = line[8:].strip()
-        if not query:
-            print("Please provide a question.")
-            continue
-        # Decide tool
-        tool_name = "web_search" if should_use_web(query) else "search_local_kb"
-        # Invoke agent
-        try:
-            result = agent.invoke({"input": query, "tool_choice": tool_name})
-            answer = result.get("output", "")
-            print("Answer:\n", answer)
-        except Exception as e:
-            print(f"Error: {e}")
-        continue
-    print("Unknown command. Type /help for usage.")
-""
+        response = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
+        # The response is a dict with a "messages" key
+        assistant_msg = next(
+            m for m in response["messages"] if m["role"] == "assistant"
+        )
+        print("\nAssistant:", assistant_msg["content"].strip())
+    except Exception as e:
+        print("Error:", e)
+
+"""End of main.py"""