diff --git a/main.py b/main.py index dde65b4..bb8c6e6 100644 --- a/main.py +++ b/main.py @@ -1,20 +1,51 @@ -""" -Simple chat loop for the RAG agent. -""" +import argparse + +from dotenv import load_dotenv -import os -from langchain.agents import AgentExecutor from agent import agent +from vectorstore import create_vectorstore, load_documents -# Create executor with memory -executor = AgentExecutor(agent=agent, verbose=True) -print("RAG‑Agent ready. Type 'exit' to quit.") -while True: - user_input = input("You: ") - if user_input.lower() in {"exit", "quit"}: - break - result = executor.invoke({"input": user_input}) - # The agent returns a dict with keys 'output' and possibly tool calls. - print(f"Assistant: {result.get('output', '')}") -print("Goodbye!") +def ingest_documents(directory: str) -> None: + vectorstore = create_vectorstore() + chunk_count = load_documents(directory, vectorstore) + print(f"Loaded {chunk_count} chunks from {directory} into ChromaDB.") + + +def run_chat() -> None: + print("RAG agent ready. Type 'exit' to quit.") + while True: + user_input = input("You: ").strip() + if user_input.lower() == "exit": + print("Goodbye!") + return + if not user_input: + continue + + result = agent.invoke( + {"messages": [{"role": "user", "content": user_input}]} + ) + final_message = result["messages"][-1] + print(f"Assistant: {final_message.content}") + + +def main() -> None: + load_dotenv() + + parser = argparse.ArgumentParser() + parser.add_argument( + "--ingest", + metavar="DIRECTORY", + help="Load .txt and .md files from DIRECTORY into ChromaDB before the demo.", + ) + args = parser.parse_args() + + if args.ingest: + ingest_documents(args.ingest) + return + + run_chat() + + +if __name__ == "__main__": + main()