From 416ea1a005413a0e9ea1a0f2dc35e67264437864 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A0=D0=BE=D0=BC?= =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= Date: Fri, 5 Jun 2026 11:29:15 +0000 Subject: [PATCH] Update src/cli.py --- src/cli.py | 109 +++++++++++++++++++++++------------------------------ 1 file changed, 48 insertions(+), 61 deletions(-) diff --git a/src/cli.py b/src/cli.py index 6bade2b..2b66041 100644 --- a/src/cli.py +++ b/src/cli.py @@ -1,75 +1,62 @@ -"""Interactive command line client for the RAG agent. +"""CLI entry point for the RAG agent. -Commands: - /add <path> – Load a document from a file and add it to the knowledge base. - /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``. +This module provides a :func:`run_cli` function that loads documents from a +directory into the knowledge base and then starts an interactive chat loop. """ from __future__ import annotations -import sys from pathlib import Path +from typing import Iterable -from .agent import run_query -from .tools import kb +from src.vector_store import kb +from src.agent import run_query -HELP_TEXT = """ -Available commands: - /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 load_documents_from_dir(directory: str | Path) -> None: + """Load all ``.txt`` files from *directory* into the knowledge base. -def main() -> None: - print("RAG Agent CLI. Type /help for commands.") + Parameters + ---------- + 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("> ") - except (EOFError, KeyboardInterrupt): - print("\nExiting.") + user_input = input("You: ") + if user_input.strip().lower() == "/quit": + print("Bye!") break - if not user_input.strip(): + if user_input.strip() == "": continue - if user_input.startswith("/add"): - parts = user_input.split(maxsplit=2) - 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) + response = run_query(user_input) + print(f"Agent: {response}\n") if __name__ == "__main__": - main() \ No newline at end of file + 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) \ No newline at end of file