add main.py

main.py

@@ -1,29 +1,82 @@
 """
-Entry point demonstrating the RAG agent.
+CLI for interacting with the RAG agent.
 
-Runs three example scenarios:
+Commands:
-1. Add a sample text directly via tool call.
+  /add <file_path>   – add document to knowledge base
-2. Search for a keyword.
+  /search <query>    – search knowledge base
-3. Load documents from a directory using init_documents.
+  /quit              – exit
+
+Optional argument: --load-dir DIR – load all .txt files from directory at start.
 """
 
-import asyncio
+import argparse
+import sys
 from pathlib import Path
+
+from rag_agent import agent, add_to_knowledge_base, search_knowledge_base
 from langchain_core.messages import HumanMessage
-from agent import agent
 
-async def demo():
-    # 1. Add sample text
-    await agent.ainvoke(
-        {"messages": [HumanMessage(content="Add document: Sample")]},
-        {"configurable": {"thread_id": "demo-1"}},
-    )
-    # 2. Search
-    result = await agent.ainvoke(
-        {"messages": [HumanMessage(content="Search for Python")]},
-        {"configurable": {"thread_id": "demo-1"}},
-    )
-    print("Search result:\n", result["messages"][-1].content)
 
-if __name__ == "__main__":  # pragma: no cover
+def load_directory(dir_path: str):
-    asyncio.run(demo())
+    """Load all .txt files from a directory into the knowledge base."""
+    p = Path(dir_path)
+    if not p.is_dir():
+        print(f"{dir_path} is not a directory")
+        return
+    for txt_file in p.rglob("*.txt"):
+        try:
+            content = txt_file.read_text(encoding="utf-8")
+            title = txt_file.stem
+            add_to_knowledge_base(content=content, title=title)
+            print(f"Loaded {txt_file}")
+        except Exception as e:
+            print(f"Failed to load {txt_file}: {e}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="RAG agent CLI")
+    parser.add_argument("--load-dir", type=str, help="Directory with .txt files to preload")
+    args = parser.parse_args()
+
+    if args.load_dir:
+        load_directory(args.load_dir)
+
+    print("Enter commands (/add <file>, /search <query>, /quit). Press Ctrl+C to exit.")
+    while True:
+        try:
+            line = input(">>> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nExiting.")
+            break
+        if not line:
+            continue
+        if line.lower() == "/quit":
+            print("Bye!")
+            break
+        if line.startswith("/add "):
+            path = line[5:].strip()
+            try:
+                content = Path(path).read_text(encoding="utf-8")
+                title = Path(path).stem
+                res = add_to_knowledge_base(content=content, title=title)
+                print(res)
+            except Exception as e:
+                print(f"Error adding file: {e}")
+        elif line.startswith("/search "):
+            query = line[8:].strip()
+            try:
+                res = search_knowledge_base(query=query, max_results=5)
+                print(res)
+            except Exception as e:
+                print(f"Error searching: {e}")
+        else:
+            # Treat any other input as a message to the agent
+            try:
+                result = agent.ainvoke({"messages": [HumanMessage(content=line)]})
+                print(result["messages"][-1].content)
+            except Exception as e:
+                print(f"Agent error: {e}")
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(main())