Update src/cli.py

This commit is contained in:
2026-06-05 10:25:18 +00:00
parent df25f76aa1
commit 7ff9b2fbb9
+56 -48
View File
@@ -1,67 +1,75 @@
"""Interactive commandline client for the RAG agent.
"""Interactive command line client for the RAG agent.
Commands:
/add add a new document to the knowledge base.
/search perform a semantic search.
/quit exit the program.
/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.
Any other input is treated as a user message and is processed by the agent.
The client uses the global agent defined in ``src.agent`` and the knowledge
base instance from ``src.tools``.
"""
from typing import List
from __future__ import annotations
from .agent import run_agent
from .tools import search_knowledge_base, add_to_knowledge_base
import sys
from pathlib import Path
from .agent import run_query
from .tools import kb
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 main() -> None:
print("Welcome to the RAG agent CLI. Type /help for commands.")
history: List[dict] = []
print("RAG Agent CLI. Type /help for commands.")
while True:
try:
user_input = input(">>> ")
except EOFError:
user_input = input("> ")
except (EOFError, KeyboardInterrupt):
print("\nExiting.")
break
if not user_input:
if not user_input.strip():
continue
if user_input.startswith("/"):
cmd, *args = user_input.split(maxsplit=1)
if cmd == "/quit":
print("Goodbye!")
break
elif cmd == "/help":
print("Commands: /add, /search, /quit")
if user_input.startswith("/add"):
parts = user_input.split(maxsplit=2)
if len(parts) != 3:
print("Usage: /add <title> <file_path>")
continue
elif cmd == "/add":
title = input("Title: ")
print("Enter content (end with a single line containing only 'END'):\n")
lines = []
while True:
line = input()
if line.strip() == "END":
break
lines.append(line)
content = "\n".join(lines)
response = add_to_knowledge_base(content, title)
print(response)
title, file_path = parts[1], parts[2]
path = Path(file_path)
if not path.is_file():
print(f"File not found: {file_path}")
continue
elif cmd == "/search":
query = input("Query: ")
results = search_knowledge_base(query, max_results=5)
if not results:
print("No results found.")
else:
for i, res in enumerate(results, 1):
print(f"{i}. Title: {res['title']}, Score: {res['score']:.4f}")
print(f" {res['content'][:200]}...\n")
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
else:
print("Unknown command. Type /help for list of commands.")
results = kb.search(query, limit=5)
if not results:
print("No results found.")
continue
# Normal user message
history.append({"role": "user", "content": user_input})
reply = run_agent(history)
print(f"Assistant: {reply}")
history.append({"role": "assistant", "content": reply})
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()