Update main.py

main.py

@@ -1,63 +1,80 @@
-"""
-CLI entry point for the RAG agent.
+"""Simple CLI for the RAG agent.
 
-The script loads / creates the vector store, populates it from the ``documents``
-folder and starts an interactive chat loop.
+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.
 """
 
 import os
 import sys
 from pathlib import Path
 
+from langchain_ollama import ChatOllama
 from dotenv import load_dotenv
 
-# Load environment variables (TAVILY_API_KEY, etc.)
+from vectorstore import create_vectorstore, load_documents
+from agent import create_agent, should_use_web
+
+# Load environment variables (TAVILY_API_KEY)
 load_dotenv()
 
-# Import our modules
-from vectorstore import create_vectorstore, load_documents
-from agent import create_agent
+# Create or load vector store
+VECTORSTORE_DIR = "./chroma_db"
+vectorstore = create_vectorstore(persist_directory=VECTORSTORE_DIR)
 
-# ---------------------------------------------------------------------------
-# Helper: populate vector store
-# ---------------------------------------------------------------------------
+# Create agent
+agent = create_agent(vectorstore)
 
-def init_vectorstore(persist_dir: str = "./chroma_db", docs_dir: str = "./documents"):
-    """Create or load the vector store and load documents if needed."""
-    vectorstore = create_vectorstore(persist_directory=persist_dir)
-    # Always load documents – Chroma will deduplicate if already present.
-    print("Loading documents into ChromaDB…")
-    load_documents(docs_dir, vectorstore)
-    return 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"
+)
 
-# ---------------------------------------------------------------------------
-# Main chat loop
-# ---------------------------------------------------------------------------
+print("RAG Agent CLI. Type /help for commands.")
 
-def main():
-    print("Initializing RAG agent…")
-    vectorstore = init_vectorstore()
-    agent = create_agent(vectorstore)
-
-    print("RAG agent ready. Type your question (or 'exit' to quit).")
-    while True:
+while True:
     try:
-            user_input = input("\n> ")
+        line = input("> ").strip()
     except (EOFError, KeyboardInterrupt):
-            print("\nGoodbye!")
+        print("\nExiting.")
         break
-        if user_input.strip().lower() in {"exit", "quit", "q"}:
-            print("Goodbye!")
-            break
-        if not user_input.strip():
+    if not line:
         continue
-        # Run the agent and capture the output
-        result = agent.run(user_input)
-        print("\nAnswer:\n", result)
-
-if __name__ == "__main__":
-    main()
-
-# ---------------------------------------------------------------------------
-# End of script
-# ---------------------------------------------------------------------------
+    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.")
+""