46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
import os
|
|
from pathlib import Path
|
|
from agent import create_rag_agent
|
|
from utils import load_documents, init_qdrant_client
|
|
|
|
# Load initial documents from docs/ directory
|
|
DOCS_DIR = Path("docs")
|
|
if DOCS_DIR.exists():
|
|
init_qdrant_client()
|
|
load_documents(DOCS_DIR)
|
|
|
|
agent = create_rag_agent()
|
|
|
|
print("RAG Agent ready. Use /add <file>, /search <query>, or /quit to exit.")
|
|
while True:
|
|
user_input = input("> ").strip()
|
|
if not user_input:
|
|
continue
|
|
if user_input.lower() == "/quit":
|
|
print("Goodbye!")
|
|
break
|
|
if user_input.startswith("/add"):
|
|
parts = user_input.split(maxsplit=1)
|
|
if len(parts) != 2:
|
|
print("Usage: /add <file_path>")
|
|
continue
|
|
file_path = Path(parts[1])
|
|
if not file_path.exists():
|
|
print(f"File {file_path} does not exist.")
|
|
continue
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
content = f.read()
|
|
result = agent.invoke({"messages": [{"role": "user", "content": f"Add document: {parts[1]}\n{content}"}]})
|
|
print(result.messages[-1].content)
|
|
elif user_input.startswith("/search"):
|
|
query = user_input[len("/search"):].strip()
|
|
if not query:
|
|
print("Usage: /search <query>")
|
|
continue
|
|
result = agent.invoke({"messages": [{"role": "user", "content": f"Search for: {query}"}]})
|
|
print(result.messages[-1].content)
|
|
else:
|
|
# Regular chat with agent
|
|
result = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
|
|
print(result.messages[-1].content)
|