Добавить init_loader.py

init_loader.py

@@ -0,0 +1,109 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+init_loader.py
+
+Скрипт для загрузки файлов из указанной директории в локальное Qdrant‑хранилище.
+Использует:
+    * langchain.document_loaders.FileSystemLoader – чтение всех файлов
+    * RecursiveCharacterTextSplitter – разбиение на чанки
+    * OllamaEmbeddings (nomic-embed-text) – генерация эмбеддингов
+    * QdrantVectorStore – сохранение векторных представлений
+
+Параметры:
+    --data-dir   Путь к директории с исходными документами
+    --collection Название коллекции в Qdrant (по умолчанию: rag_collection)
+"""
+
+import argparse
+import os
+from pathlib import Path
+from typing import List, Dict
+
+# LangChain imports
+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.schema import Document
+
+
+def load_documents(data_dir: Path) -> List[Document]:
+    """
+    Загружает все файлы из указанной директории в список объектов Document.
+    Поддерживаются форматы *.txt, *.md, *.pdf (если установлен pdfminer).
+    """
+    loader = DirectoryLoader(
+        str(data_dir),
+        glob="**/*",
+        suffixes=[".txt", ".md", ".pdf"],
+        show_progress=True,
+    )
+    return loader.load()
+
+
+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,
+        embedding=embedding,
+        url="http://localhost:6333",
+    )
+
+    # Добавляем документы в хранилище
+    vector_store.add_documents(chunks)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Загрузка документов в Qdrant.")
+    parser.add_argument(
+        "--data-dir",
+        type=str,
+        required=True,
+        help="Путь к директории с документами.",
+    )
+    parser.add_argument(
+        "--collection",
+        type=str,
+        default="rag_collection",
+        help="Имя коллекции в Qdrant.",
+    )
+    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__":
+    main()