69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
"""Simple CLI for interacting with 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.
|
|
"""
|
|
|
|
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
|
|
|
|
|
|
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:
|
|
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.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|