diff --git a/main.py b/main.py index ab7388b..289e4ed 100644 --- a/main.py +++ b/main.py @@ -1,63 +1,80 @@ -""" -CLI entry point for the RAG agent. +"""Simple CLI for the RAG agent. -The script loads / creates the vector store, populates it from the ``documents`` -folder and starts an interactive chat loop. +Commands: + /add – Load all .txt/.md files from the directory into the local KB. + /search – 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 -# Load environment variables (TAVILY_API_KEY, etc.) +from vectorstore import create_vectorstore, load_documents +from agent import create_agent, should_use_web + +# Load environment variables (TAVILY_API_KEY) load_dotenv() -# Import our modules -from vectorstore import create_vectorstore, load_documents -from agent import create_agent +# Create or load vector store +VECTORSTORE_DIR = "./chroma_db" +vectorstore = create_vectorstore(persist_directory=VECTORSTORE_DIR) -# --------------------------------------------------------------------------- -# Helper: populate vector store -# --------------------------------------------------------------------------- +# Create agent +agent = create_agent(vectorstore) -def init_vectorstore(persist_dir: str = "./chroma_db", docs_dir: str = "./documents"): - """Create or load the vector store and load documents if needed.""" - vectorstore = create_vectorstore(persist_directory=persist_dir) - # Always load documents – Chroma will deduplicate if already present. - print("Loading documents into ChromaDB…") - load_documents(docs_dir, vectorstore) - return vectorstore +# Helper to print usage +USAGE = ( + "Commands:\n" + " /add – Load documents into the local knowledge base.\n" + " /search – Ask the agent a question.\n" + " /quit – Exit the program.\n" +) -# --------------------------------------------------------------------------- -# Main chat loop -# --------------------------------------------------------------------------- +print("RAG Agent CLI. Type /help for commands.") -def main(): - print("Initializing RAG agent…") - vectorstore = init_vectorstore() - agent = create_agent(vectorstore) - - print("RAG agent ready. Type your question (or 'exit' to quit).") - while True: - try: - user_input = input("\n> ") - except (EOFError, KeyboardInterrupt): - print("\nGoodbye!") - break - if user_input.strip().lower() in {"exit", "quit", "q"}: - print("Goodbye!") - break - if not user_input.strip(): +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 - # Run the agent and capture the output - result = agent.run(user_input) - print("\nAnswer:\n", result) - -if __name__ == "__main__": - main() - -# --------------------------------------------------------------------------- -# End of script -# --------------------------------------------------------------------------- \ No newline at end of file + 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.") +"" \ No newline at end of file