Update src/cli.py
This commit is contained in:
+48
-61
@@ -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)
|
||||
Reference in New Issue
Block a user