Update src/cli.py

This commit is contained in:
2026-06-05 11:29:15 +00:00
parent e9293ec211
commit 416ea1a005
+48 -61
View File
@@ -1,75 +1,62 @@
"""Interactive command line client for the RAG agent.
"""CLI entry point for the RAG agent.
Commands:
/add <title> <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()
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)