From 8b154f4630148290a937048ea99f9134da5e8db8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9C=D0=B0=D1=80=D0=B8=D1=8F=20=D0=91=D0=B5=D1=80=D0=B4?= =?UTF-8?q?=D0=BD=D0=B8=D0=BA=D0=BE=D0=B2=D0=B0?= Date: Wed, 27 May 2026 06:52:08 +0000 Subject: [PATCH] =?UTF-8?q?=D0=90=D0=B3=D0=B5=D0=BD=D1=82=20=D1=81=20RAG?= =?UTF-8?q?=E2=80=91=D0=BF=D0=B0=D0=BC=D1=8F=D1=82=D1=8C=D1=8E:=20init=5Fl?= =?UTF-8?q?oader.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../init_loader.py | 155 ++++++++---------- 1 file changed, 72 insertions(+), 83 deletions(-) diff --git a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/init_loader.py b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/init_loader.py index 55e8652..388ff25 100644 --- a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/init_loader.py +++ b/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) \ No newline at end of file + main() \ No newline at end of file