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

2026-05-27 06:34:20 +00:00
parent 670ea83751
commit 3d3cb640a9
1 changed files with 120 additions and 0 deletions
@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+init_loader.py
+
+Скрипт для загрузки файлов из указанной директории в локальное векторное хранилище Qdrant.
+Использует:
+    - LangChain OllamaEmbeddings (model "nomic-embed-text")
+    - RecursiveCharacterTextSplitter для разбиения на чанки
+    - langchain_qdrant.QdrantVectorStore для сохранения эмбеддингов
+
+Параметры запуска:
+    python init_loader.py <директория_с_файлами> [--collection <имя_коллекции>]
+
+Пример:
+    python init_loader.py ./docs --collection knowledge_base
+"""
+
+import argparse
+import os
+from pathlib import Path
+from typing import Iterable
+
+# LangChain imports
+from langchain.embeddings.ollama import OllamaEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores.qdrant import QdrantVectorStore
+from langchain.document_loaders import TextLoader, UnstructuredFileLoader
+
+# --------------------------------------------------------------------------- #
+# Конфигурация по умолчанию
+DEFAULT_COLLECTION_NAME = "rag_collection"
+QDRANT_HOST = "localhost"
+QDRANT_PORT = 6333
+EMBEDDING_MODEL = "nomic-embed-text"
+
+# --------------------------------------------------------------------------- #
+def _get_text_loaders(directory: Path) -> Iterable[TextLoader]:
+    """
+    Возвращает генератор загрузчиков для всех файлов в директории.
+    Поддерживаются текстовые файлы и любые форматы, которые умеет UnstructuredFileLoader.
+    """
+    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,
+    )
+
+    # 3. Инициализируем QdrantVectorStore (создаст коллекцию при необходимости)
+    vector_store = QdrantVectorStore.from_existing_collection(
+        embedding=embeddings,
+        collection_name=collection_name,
+        url=f"http://{QDRANT_HOST}:{QDRANT_PORT}",
+    )
+
+    # 4. Загружаем и разбиваем документы, затем сохраняем
+    for loader in _get_text_loaders(directory):
+        documents = loader.load()
+        if not documents:
+            continue
+
+        # Разбиваем каждый документ на чанки
+        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."
+    )
+    parser.add_argument(
+        "directory",
+        type=str,
+        help="Путь к каталогу с документами.",
+    )
+    parser.add_argument(
+        "--collection",
+        type=str,
+        default=DEFAULT_COLLECTION_NAME,
+        help=f"Имя коллекции (по умолчанию: {DEFAULT_COLLECTION_NAME}).",
+    )
+    return parser.parse_args()
+
+# --------------------------------------------------------------------------- #
+if __name__ == "__main__":
+    args = _parse_args()
+    load_directory_to_qdrant(Path(args.directory), collection_name=args.collection)