80 lines
2.3 KiB
Python
80 lines
2.3 KiB
Python
"""Simple CLI for the RAG agent.
|
||
|
||
Commands:
|
||
/add <directory> – Load all .txt/.md files from the directory into the local KB.
|
||
/search <question> – Ask the agent a question.
|
||
/quit – Exit the program.
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
from langchain_ollama import ChatOllama
|
||
from dotenv import load_dotenv
|
||
|
||
from vectorstore import create_vectorstore, load_documents
|
||
from agent import create_agent, should_use_web
|
||
|
||
# Load environment variables (TAVILY_API_KEY)
|
||
load_dotenv()
|
||
|
||
# Create or load vector store
|
||
VECTORSTORE_DIR = "./chroma_db"
|
||
vectorstore = create_vectorstore(persist_directory=VECTORSTORE_DIR)
|
||
|
||
# Create agent
|
||
agent = create_agent(vectorstore)
|
||
|
||
# Helper to print usage
|
||
USAGE = (
|
||
"Commands:\n"
|
||
" /add <directory> – Load documents into the local knowledge base.\n"
|
||
" /search <question> – Ask the agent a question.\n"
|
||
" /quit – Exit the program.\n"
|
||
)
|
||
|
||
print("RAG Agent CLI. Type /help for commands.")
|
||
|
||
while True:
|
||
try:
|
||
line = input("> ").strip()
|
||
except (EOFError, KeyboardInterrupt):
|
||
print("\nExiting.")
|
||
break
|
||
if not line:
|
||
continue
|
||
if line.lower() == "/help":
|
||
print(USAGE)
|
||
continue
|
||
if line.lower() == "/quit":
|
||
print("Bye!")
|
||
break
|
||
if line.lower().startswith("/add "):
|
||
dir_path = line[5:].strip()
|
||
if not dir_path:
|
||
print("Please provide a directory path.")
|
||
continue
|
||
if not Path(dir_path).exists():
|
||
print(f"Directory {dir_path} does not exist.")
|
||
continue
|
||
load_documents(dir_path, vectorstore)
|
||
print("Documents loaded.")
|
||
continue
|
||
if line.lower().startswith("/search "):
|
||
query = line[8:].strip()
|
||
if not query:
|
||
print("Please provide a question.")
|
||
continue
|
||
# Decide tool
|
||
tool_name = "web_search" if should_use_web(query) else "search_local_kb"
|
||
# Invoke agent
|
||
try:
|
||
result = agent.invoke({"input": query, "tool_choice": tool_name})
|
||
answer = result.get("output", "")
|
||
print("Answer:\n", answer)
|
||
except Exception as e:
|
||
print(f"Error: {e}")
|
||
continue
|
||
print("Unknown command. Type /help for usage.")
|
||
"" |