feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'

main.py

@@ -0,0 +1,65 @@
+import os
+import sys
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+from vectorstore import create_vectorstore, load_documents, collection_exists
+from agent import create_agent
+
+def main():
+    # Load environment variables
+    load_dotenv()
+
+    # Configuration
+    qdrant_path = os.getenv("QDRANT_PATH", "./qdrant_db")
+    embedding_model = os.getenv("EMBEDDING_MODEL", "nomic-embed-text")
+    llm_model = os.getenv("LLM_MODEL", "llama3")
+    tavily_api_key = os.getenv("TAVILY_API_KEY")
+    if not tavily_api_key:
+        print("Error: TAVILY_API_KEY not set in .env")
+        sys.exit(1)
+
+    # Create vector store
+    vectorstore = create_vectorstore(
+        persist_directory=qdrant_path,
+        collection_name="documents",
+        embedding_model=embedding_model,
+    )
+
+    # Load documents if collection is empty
+    if not collection_exists(vectorstore):
+        print("Loading documents into Qdrant...")
+        docs_dir = Path("documents")
+        if not docs_dir.exists():
+            print(f"Documents directory '{docs_dir}' not found.")
+            sys.exit(1)
+        load_documents(str(docs_dir), vectorstore)
+        print("Documents loaded.")
+    else:
+        print("Qdrant collection already exists. Skipping document load.")
+
+    # Create agent
+    agent = create_agent(vectorstore, tavily_api_key, llm_model=llm_model)
+
+    print("Chat agent ready. Type 'exit' to quit.")
+    while True:
+        try:
+            user_input = input("\nYou: ")
+        except (KeyboardInterrupt, EOFError):
+            print("\nExiting.")
+            break
+
+        if user_input.strip().lower() in {"exit", "quit"}:
+            print("Goodbye!")
+            break
+
+        # Run agent
+        try:
+            response = agent.run(user_input)
+            print(f"\nAssistant: {response}")
+        except Exception as e:
+            print(f"Error: {e}")
+
+if __name__ == "__main__":
+    main()