63 lines
2.0 KiB
Python
63 lines
2.0 KiB
Python
"""
|
||
CLI entry point for the RAG agent.
|
||
|
||
The script loads / creates the vector store, populates it from the ``documents``
|
||
folder and starts an interactive chat loop.
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
from dotenv import load_dotenv
|
||
|
||
# Load environment variables (TAVILY_API_KEY, etc.)
|
||
load_dotenv()
|
||
|
||
# Import our modules
|
||
from vectorstore import create_vectorstore, load_documents
|
||
from agent import create_agent
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helper: populate vector store
|
||
# ---------------------------------------------------------------------------
|
||
|
||
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
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main chat loop
|
||
# ---------------------------------------------------------------------------
|
||
|
||
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():
|
||
continue
|
||
# Run the agent and capture the output
|
||
result = agent.run(user_input)
|
||
print("\nAnswer:\n", result)
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# End of script
|
||
# --------------------------------------------------------------------------- |