Update src/cli.py

src/cli.py

@@ -1,75 +1,62 @@
-"""Interactive command line client for the RAG agent.
+"""CLI entry point for the RAG agent.
 
-Commands:
+This module provides a :func:`run_cli` function that loads documents from a
-  /add <title> <path>   – Load a document from a file and add it to the knowledge base.
+directory into the knowledge base and then starts an interactive chat loop.
-  /search <query>       – Search the knowledge base and display results.
-  /quit                 – Exit the program.
-  /help                 – Show this help message.
-
-The client uses the global agent defined in ``src.agent`` and the knowledge
-base instance from ``src.tools``.
 """
 
 from __future__ import annotations
 
-import sys
 from pathlib import Path
+from typing import Iterable
 
-from .agent import run_query
+from src.vector_store import kb
-from .tools import kb
+from src.agent import run_query
 
-HELP_TEXT = """
+def load_documents_from_dir(directory: str | Path) -> None:
-Available commands:
+    """Load all ``.txt`` files from *directory* into the knowledge base.
-  /add <title> <file_path>   Add a document to the knowledge base.
-  /search <query>             Search the knowledge base.
-  /quit                       Exit the program.
-  /help                       Show this help message.
-"""
 
-def main() -> None:
+    Parameters
-    print("RAG Agent CLI. Type /help for commands.")
+    ----------
+    directory:
+        Path to the folder containing documents.
+    """
+    directory = Path(directory)
+    for file_path in directory.rglob("*.txt"):
+        title = file_path.stem
+        content = file_path.read_text(encoding="utf-8")
+        kb.add_document(title=title, content=content)
+        print(f"Loaded {title}")
+
+def run_cli(docs_dir: str | Path) -> None:
+    """Run an interactive CLI.
+
+    Parameters
+    ----------
+    docs_dir:
+        Directory with documents to load into the knowledge base.
+    """
+    load_documents_from_dir(docs_dir)
+
+    print("\n--- RAG Agent ready. Type your question (or /quit to exit). ---\n")
     while True:
-        try:
+        user_input = input("You: ")
-            user_input = input("> ")
+        if user_input.strip().lower() == "/quit":
-        except (EOFError, KeyboardInterrupt):
+            print("Bye!")
-            print("\nExiting.")
             break
-        if not user_input.strip():
+        if user_input.strip() == "":
             continue
-        if user_input.startswith("/add"):
+        response = run_query(user_input)
-            parts = user_input.split(maxsplit=2)
+        print(f"Agent: {response}\n")
-            if len(parts) != 3:
-                print("Usage: /add <title> <file_path>")
-                continue
-            title, file_path = parts[1], parts[2]
-            path = Path(file_path)
-            if not path.is_file():
-                print(f"File not found: {file_path}")
-                continue
-            content = path.read_text(encoding="utf-8")
-            kb.add_document(title=title, content=content)
-            print(f"Document '{title}' added.")
-        elif user_input.startswith("/search"):
-            query = user_input[len("/search"):].strip()
-            if not query:
-                print("Please provide a search query.")
-                continue
-            results = kb.search(query, limit=5)
-            if not results:
-                print("No results found.")
-                continue
-            print("Results:")
-            for i, r in enumerate(results, 1):
-                print(f"{i}. [{r['title']} - chunk {r['chunk_index']}] {r['content'][:200]}...")
-        elif user_input.startswith("/quit"):
-            print("Goodbye.")
-            break
-        elif user_input.startswith("/help"):
-            print(HELP_TEXT)
-        else:
-            # Treat as a normal user query to the agent.
-            response = run_query(user_input)
-            print(response)
 
 if __name__ == "__main__":
-    main()
+    import argparse
+
+    parser = argparse.ArgumentParser(description="RAG agent CLI")
+    parser.add_argument(
+        "--docs",
+        type=str,
+        default="docs",
+        help="Directory with documents to load into the knowledge base.",
+    )
+    args = parser.parse_args()
+    run_cli(args.docs)