diff --git a/cli.py b/cli.py new file mode 100644 index 0000000..d6caee1 --- /dev/null +++ b/cli.py @@ -0,0 +1,90 @@ +"""Interactive CLI for the RAG agent. + +The CLI supports the following commands: + +* ``/add `` – Load all .txt/.md files from *dir* into Qdrant. +* ``/search `` – 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() \ No newline at end of file