Публикация полного набора файлов: add main.py

main.py

@@ -0,0 +1,51 @@
+import argparse
+
+from dotenv import load_dotenv
+
+from agent import agent
+from vectorstore import create_vectorstore, load_documents
+
+
+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()