Build ChromaDB + Tavily RAG agent with Ollama embeddings, local/web tools, create_agent routing, and CLI ingest flow.: update main.py

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()