Update cli.py

This commit is contained in:
2026-06-04 20:02:26 +00:00
parent 79dac766c7
commit 6f3f64a4b4
+38 -57
View File
@@ -1,68 +1,49 @@
"""Simple CLI for interacting with the RAG agent.
"""Simple interactive CLI for the RAG agent.
Commands:
/add <title> <file_path> - Add a document to the knowledge base.
/search <query> [max] - Search the knowledge base.
/quit - Exit the CLI.
/add title content add a document to the knowledge base.
/search query perform a semantic search.
/quit exit.
any other text will be sent to the agent for normal answering.
"""
import sys
from pathlib import Path
# Import the tools directly. They are decorated with @tool but can be called like normal functions.
from rag_tools import add_to_knowledge_base, search_knowledge_base
from .agent import create_agent_executor
agent = create_agent_executor()
def main():
print("Welcome to the RAG Agent CLI. Type /help for commands.")
while True:
try:
user_input = input(">> ")
except (EOFError, KeyboardInterrupt):
print("\nExiting.")
break
if not user_input:
print("RAG Agent CLI. Type /quit to exit.")
while True:
try:
user_input = input(">>> ")
except EOFError:
break
if not user_input:
continue
if user_input.lower() == "/quit":
print("Goodbye!")
break
if user_input.startswith("/add "):
# Expected format: /add title content
parts = user_input.split(" ", 2)
if len(parts) < 3:
print("Usage: /add title content")
continue
if user_input.startswith("/help"):
print("Commands:\n /add <title> <file_path> - Add a document to the knowledge base.\n /search <query> [max] - Search the knowledge base.\n /quit - Exit the CLI.")
continue
if user_input.startswith("/quit"):
print("Goodbye!")
break
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]
try:
content = Path(file_path).read_text(encoding="utf-8")
except Exception as e:
print(f"Error reading file: {e}")
continue
print("Adding document...", end=" ")
result = add_to_knowledge_base(content=content, title=title)
print(result)
continue
if user_input.startswith("/search"):
parts = user_input.split(maxsplit=2)
if len(parts) < 2:
print("Usage: /search <query> [max_results]")
continue
query = parts[1]
max_results = 5
if len(parts) == 3:
try:
max_results = int(parts[2])
except ValueError:
print("max_results must be an integer.")
continue
print("Searching...", end=" ")
result = search_knowledge_base(query=query, max_results=max_results)
print(result)
continue
print("Unknown command. Type /help for a list of commands.")
title, content = parts[1], parts[2]
# Directly call the tool via the agent
result = agent.run({"input": f"Add document: {title} {content}"})
print(result)
continue
if user_input.startswith("/search "):
query = user_input[8:].strip()
result = agent.run({"input": f"Search for: {query}"})
print(result)
continue
# Normal conversation
result = agent.run({"input": user_input})
print(result)
if __name__ == "__main__":
main()
# The CLI is already running in the main thread
pass