Update main.py

main.py

@@ -1,25 +1,63 @@
-"""Entry point for the RAG agent.
+"""
+CLI entry point for the RAG agent.
 
-Running ``python main.py`` starts a simple CLI that accepts user queries and
+The script loads / creates the vector store, populates it from the ``documents``
-prints the agent's answer along with the source.
+folder and starts an interactive chat loop.
 """
 
-from agent import create_agent, create_vectorstore
+import os
-from vectorstore import load_documents
+import sys
+from pathlib import Path
 
-if __name__ == "__main__":
+from dotenv import load_dotenv
-    # Load or create the vector store
+
-    store = create_vectorstore()
+# Load environment variables (TAVILY_API_KEY, etc.)
-    if not store.get_index_info():
+load_dotenv()
-        load_documents("documents", store)
+
-    agent = create_agent()
+# Import our modules
-    print("RAG Agent ready. Type 'exit' to quit.")
+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:
-            query = input("\nЗапрос: ")
+            user_input = input("\n> ")
-        except EOFError:
+        except (EOFError, KeyboardInterrupt):
+            print("\nGoodbye!")
             break
-        if query.strip().lower() in {"exit", "quit"}:
+        if user_input.strip().lower() in {"exit", "quit", "q"}:
+            print("Goodbye!")
             break
-        result = agent.invoke({"input": query})
+        if not user_input.strip():
-        print("Ответ:", result["output"])  # noqa: T201
+            continue
+        # Run the agent and capture the output
+        result = agent.run(user_input)
+        print("\nAnswer:\n", result)
+
+if __name__ == "__main__":
+    main()
+
+# ---------------------------------------------------------------------------
+# End of script
+# ---------------------------------------------------------------------------