Add cli.py
This commit is contained in:
@@ -0,0 +1,90 @@
|
|||||||
|
"""Interactive CLI for the RAG agent.
|
||||||
|
|
||||||
|
The CLI supports the following commands:
|
||||||
|
|
||||||
|
* ``/add <dir>`` – Load all .txt/.md files from *dir* into Qdrant.
|
||||||
|
* ``/search <query>`` – Run the agent on *query* and print the answer.
|
||||||
|
* ``/quit`` – Exit the program.
|
||||||
|
|
||||||
|
The vector store is created on first use and persisted in
|
||||||
|
``./qdrant_db``.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# Load environment variables (e.g., TAVILY_API_KEY)
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
load_dotenv()
|
||||||
|
|
||||||
|
# Local modules
|
||||||
|
from .vector_store import create_vectorstore, load_documents
|
||||||
|
from .agent import create_agent
|
||||||
|
from . import tools
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Global state
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
VECTORSTORE = None
|
||||||
|
AGENT = None
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Helper functions
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
def init_vectorstore() -> None:
|
||||||
|
global VECTORSTORE
|
||||||
|
if VECTORSTORE is None:
|
||||||
|
VECTORSTORE = create_vectorstore()
|
||||||
|
# Expose the store to the tools module so that the decorated
|
||||||
|
# functions can access it via the global name ``vectorstore``.
|
||||||
|
tools.vectorstore = VECTORSTORE
|
||||||
|
|
||||||
|
def init_agent() -> None:
|
||||||
|
global AGENT
|
||||||
|
if AGENT is None:
|
||||||
|
AGENT = create_agent()
|
||||||
|
|
||||||
|
def add_documents(directory: str) -> None:
|
||||||
|
init_vectorstore()
|
||||||
|
load_documents(directory, VECTORSTORE)
|
||||||
|
|
||||||
|
def run_query(query: str) -> None:
|
||||||
|
init_agent()
|
||||||
|
# The agent expects a dictionary with the key ``input``.
|
||||||
|
result = AGENT.invoke({"input": query})
|
||||||
|
print("\nAnswer:\n", result)
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# CLI loop
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
print("RAG Agent CLI. Type '/quit' to exit.")
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
line = input("> ")
|
||||||
|
except (EOFError, KeyboardInterrupt):
|
||||||
|
print("\nBye!")
|
||||||
|
break
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
if line.strip() == "/quit":
|
||||||
|
print("Bye!")
|
||||||
|
break
|
||||||
|
if line.startswith("/add "):
|
||||||
|
_, dir_path = line.split(" ", 1)
|
||||||
|
dir_path = dir_path.strip()
|
||||||
|
if not os.path.isdir(dir_path):
|
||||||
|
print(f"{dir_path} is not a directory.")
|
||||||
|
continue
|
||||||
|
add_documents(dir_path)
|
||||||
|
print("Documents added.")
|
||||||
|
elif line.startswith("/search "):
|
||||||
|
_, query = line.split(" ", 1)
|
||||||
|
run_query(query.strip())
|
||||||
|
else:
|
||||||
|
print("Unknown command. Use /add, /search, or /quit.")
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Reference in New Issue
Block a user