Агент с RAG‑памятью: init_loader.py

solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/init_loader.py

@@ -4,117 +4,106 @@
 """
 init_loader.py
 
-Скрипт для загрузки файлов из указанной директории в локальное векторное хранилище Qdrant.
+Скрипт для загрузки файлов из указанной директории в локальное Qdrant‑хранилище.
 Использует:
-    - LangChain OllamaEmbeddings (model "nomic-embed-text")
-    - RecursiveCharacterTextSplitter для разбиения на чанки
-    - langchain_qdrant.QdrantVectorStore для сохранения эмбеддингов
+    * langchain.document_loaders.FileSystemLoader – чтение всех файлов
+    * RecursiveCharacterTextSplitter – разбиение на чанки
+    * OllamaEmbeddings (nomic-embed-text) – генерация эмбеддингов
+    * QdrantVectorStore – сохранение векторных представлений
 
-Параметры запуска:
-    python init_loader.py <директория_с_файлами> [--collection <имя_коллекции>]
-
-Пример:
-    python init_loader.py ./docs --collection knowledge_base
+Параметры:
+    --data-dir   Путь к директории с исходными документами
+    --collection Название коллекции в Qdrant (по умолчанию: rag_collection)
 """
 
 import argparse
 import os
 from pathlib import Path
-from typing import Iterable
+from typing import List, Dict
 
 # LangChain imports
-from langchain.embeddings.ollama import OllamaEmbeddings
+from langchain.document_loaders import DirectoryLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings.ollama import OllamaEmbeddings
 from langchain.vectorstores.qdrant import QdrantVectorStore
-from langchain.document_loaders import TextLoader, UnstructuredFileLoader
+from langchain.schema import Document
 
-# --------------------------------------------------------------------------- #
-# Конфигурация по умолчанию
-DEFAULT_COLLECTION_NAME = "rag_collection"
-QDRANT_HOST = "localhost"
-QDRANT_PORT = 6333
-EMBEDDING_MODEL = "nomic-embed-text"
 
-# --------------------------------------------------------------------------- #
-def _get_text_loaders(directory: Path) -> Iterable[TextLoader]:
+def load_documents(data_dir: Path) -> List[Document]:
     """
-    Возвращает генератор загрузчиков для всех файлов в директории.
-    Поддерживаются текстовые файлы и любые форматы, которые умеет UnstructuredFileLoader.
+    Загружает все файлы из указанной директории в список объектов Document.
+    Поддерживаются форматы *.txt, *.md, *.pdf (если установлен pdfminer).
     """
-    for file_path in directory.rglob("*"):
-        if file_path.is_file():
-            # Для простоты используем TextLoader для .txt,
-            # а остальные обрабатываем через UnstructuredFileLoader
-            if file_path.suffix.lower() == ".txt":
-                yield TextLoader(str(file_path), encoding="utf-8")
-            else:
-                yield UnstructuredFileLoader(str(file_path))
-
-# --------------------------------------------------------------------------- #
-def load_directory_to_qdrant(
-    directory: Path,
-    collection_name: str = DEFAULT_COLLECTION_NAME,
-) -> None:
-    """
-    Загружает все документы из указанной директории в Qdrant.
-
-    Args:
-        directory (Path): Путь к каталогу с документами.
-        collection_name (str): Имя коллекции в Qdrant.
-    """
-    if not directory.is_dir():
-        raise ValueError(f"Путь {directory} не является директорией")
-
-    # 1. Инициализируем эмбеддер
-    embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
-
-    # 2. Создаём текстовый splitter
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=1000,
-        chunk_overlap=200,
+    loader = DirectoryLoader(
+        str(data_dir),
+        glob="**/*",
+        suffixes=[".txt", ".md", ".pdf"],
+        show_progress=True,
     )
+    return loader.load()
 
-    # 3. Инициализируем QdrantVectorStore (создаст коллекцию при необходимости)
-    vector_store = QdrantVectorStore.from_existing_collection(
-        embedding=embeddings,
+
+def chunk_documents(docs: List[Document], chunk_size: int = 1000, overlap: int = 200) -> List[Document]:
+    """
+    Разбивает каждый документ на чанки фиксированного размера.
+    """
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=overlap,
+        separators=["\n\n", "\n", " ", ""],
+    )
+    return splitter.split_documents(docs)
+
+
+def embed_and_store(chunks: List[Document], collection_name: str) -> None:
+    """
+    Генерирует эмбеддинги для чанков и сохраняет их в Qdrant.
+    """
+    # Инициализируем Ollama embeddings
+    embedding = OllamaEmbeddings(model="nomic-embed-text")
+
+    # Создаём (или подключаемся к) коллекцию Qdrant
+    vector_store = QdrantVectorStore(
         collection_name=collection_name,
-        url=f"http://{QDRANT_HOST}:{QDRANT_PORT}",
+        embedding=embedding,
+        url="http://localhost:6333",
     )
 
-    # 4. Загружаем и разбиваем документы, затем сохраняем
-    for loader in _get_text_loaders(directory):
-        documents = loader.load()
-        if not documents:
-            continue
+    # Добавляем документы в хранилище
+    vector_store.add_documents(chunks)
 
-        # Разбиваем каждый документ на чанки
-        chunks = text_splitter.split_documents(documents)
 
-        # Добавляем векторные представления в Qdrant
-        vector_store.add_documents(chunks)
-
-    print(f"Загрузка завершена. Коллекция '{collection_name}' содержит "
-          f"{vector_store._client.count(collection_name=collection_name).count} точек.")
-
-# --------------------------------------------------------------------------- #
-def _parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description="Загрузить документы из директории в Qdrant."
-    )
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Загрузка документов в Qdrant.")
     parser.add_argument(
-        "directory",
+        "--data-dir",
         type=str,
-        help="Путь к каталогу с документами.",
+        required=True,
+        help="Путь к директории с документами.",
     )
     parser.add_argument(
         "--collection",
         type=str,
-        default=DEFAULT_COLLECTION_NAME,
-        help=f"Имя коллекции (по умолчанию: {DEFAULT_COLLECTION_NAME}).",
+        default="rag_collection",
+        help="Имя коллекции в Qdrant.",
     )
-    return parser.parse_args()
+    args = parser.parse_args()
+
+    data_dir = Path(args.data_dir).expanduser().resolve()
+    if not data_dir.is_dir():
+        raise FileNotFoundError(f"Каталог {data_dir} не найден.")
+
+    print(f"[INFO] Загружаем документы из: {data_dir}")
+    docs = load_documents(data_dir)
+    print(f"[INFO] Найдено {len(docs)} документов. Разбиваем на чанки...")
+
+    chunks = chunk_documents(docs)
+    print(f"[INFO] Получено {len(chunks)} чанков. Генерируем эмбеддинги и сохраняем в Qdrant...")
+
+    embed_and_store(chunks, args.collection)
+
+    print("[SUCCESS] Загрузка завершена.")
+
 
-# --------------------------------------------------------------------------- #
 if __name__ == "__main__":
-    args = _parse_args()
-    load_directory_to_qdrant(Path(args.directory), collection_name=args.collection)
+    main()