add main.py

2026-05-28 16:43:45 +00:00
parent 7d8a7769f5
commit 33c119cd3a
1 changed files with 77 additions and 0 deletions
@@ -0,0 +1,77 @@
+"""CLI for the RAG agent with ChromaDB and Tavily web search.
+
+Commands:
+  /add <text>   - add text to local knowledge base
+  /load <file>  - load a .txt file into the knowledge base
+  /help         - show available commands
+  /quit         - exit
+  <question>    - ask the agent (auto-selects local KB or web)
+"""
+import sys
+from dotenv import load_dotenv
+
+load_dotenv()
+
+from vectorstore import add_documents  # noqa: E402
+from agent import run_agent            # noqa: E402
+
+HELP_TEXT = """
+Commands:
+  /add <text>   - add text to local ChromaDB knowledge base
+  /load <file>  - load a .txt file into the knowledge base
+  /help         - show this help
+  /quit         - exit
+  <question>    - ask the agent (uses Local KB or Web as appropriate)
+"""
+
+def main() -> None:
+    print("RAG Agent: ChromaDB + Tavily Web Search")
+    print("Type /help for commands, /quit to exit.\n")
+
+    while True:
+        try:
+            line = input("You: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nExiting.")
+            sys.exit(0)
+
+        if not line:
+            continue
+
+        if line == "/quit":
+            print("Goodbye!")
+            break
+
+        elif line == "/help":
+            print(HELP_TEXT)
+
+        elif line.startswith("/add "):
+            text = line[5:].strip()
+            if not text:
+                print("Usage: /add <text>")
+                continue
+            n = add_documents([text])
+            print(f"Added {n} chunk(s) to knowledge base.")
+
+        elif line.startswith("/load "):
+            path = line[6:].strip()
+            try:
+                with open(path, encoding="utf-8") as fh:
+                    content = fh.read()
+                n = add_documents([content])
+                print(f"Loaded '{path}': {n} chunk(s) added.")
+            except FileNotFoundError:
+                print(f"File not found: {path}")
+            except Exception as exc:
+                print(f"Error: {exc}")
+
+        else:
+            try:
+                response = run_agent(line)
+                print(f"Agent: {response}\n")
+            except Exception as exc:
+                print(f"Agent error: {exc}")
+
+
+if __name__ == "__main__":
+    main()