fix: load_docs.py — Агент с RAG-памятью

load_docs.py

@@ -1,49 +1,19 @@
-import os
-import argparse
-from dotenv import load_dotenv
-from langchain_openai import OpenAIEmbeddings
-from langchain_qdrant import QdrantVectorStore
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+import sys
+from pathlib import Path
+
+from utils import splitter, vector_store
 from langchain_core.documents import Document
-from qdrant_client import QdrantClient
 
-# Загрузка переменных окружения
-load_dotenv()
-
-# Инициализация эмбеддингов
-embeddings = OpenAIEmbeddings(
-    model="text-embedding-3-small",
-    base_url="https://openrouter.ai/api/v1",
-    api_key=os.getenv("OPENAI_API_KEY"),
-)
-
-# Инициализация Qdrant
-client = QdrantClient(url="http://localhost:6333")
-collection_name = "knowledge_base"
-vector_store = QdrantVectorStore(
-    client=client,
-    collection_name=collection_name,
-    embeddings=embeddings,
-)
-
-# Чанкинг
-splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-
-def load_directory(dir_path: str):
-    """Загружает все .txt файлы из директории в векторную базу."""
-    for root, dirs, files in os.walk(dir_path):
-        for file in files:
-            if file.lower().endswith(".txt"):
-                path = os.path.join(root, file)
-                with open(path, "r", encoding="utf-8") as f:
-                    text = f.read()
-                chunks = splitter.split_text(text)
-                docs = [Document(page_content=chunk, metadata={"title": file}) for chunk in chunks]
-                vector_store.add_documents(docs)
-                print(f"Added {len(docs)} chunks from {file}")
+def load_documents(dir_path: str):
+    for file_path in Path(dir_path).rglob("*.txt"):
+        content = file_path.read_text(encoding="utf-8")
+        chunks = splitter.split_text(content)
+        docs = [Document(page_content=chunk, metadata={"title": file_path.name}) for chunk in chunks]
+        vector_store.add_documents(docs)
+        print(f"Added {len(docs)} chunks from {file_path}")
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Load documents into Qdrant.")
-    parser.add_argument("directory", help="Path to directory with .txt files")
-    args = parser.parse_args()
-    load_directory(args.directory)
+    if len(sys.argv) != 2:
+        print("Usage: python load_docs.py <directory>")
+        sys.exit(1)
+    load_documents(sys.argv[1])