diff --git a/qdrant_client.py b/qdrant_client.py new file mode 100644 index 0000000..6c93f17 --- /dev/null +++ b/qdrant_client.py @@ -0,0 +1,175 @@ +# qdrant_client.py +""" +Модуль инициализации клиента Qdrant и простых операций с ним. +Используется в агенте для RAG‑памяти. + +Поддержка: +- Создание/получение коллекции +- Добавление документов (с чанками) с эмбеддингами Ollama +- Семантический поиск по запросу + +Требуемые зависимости: + pip install langchain langchain-qdrant langchain-ollama qdrant-client +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Iterable, List, Tuple + +from langchain_ollama import OllamaEmbeddings +from langchain_qdrant import QdrantVectorStore +from langchain_text_splitters import RecursiveCharacterTextSplitter +from qdrant_client import QdrantClient as _QdrantClient +from qdrant_client.http.models import Distance, VectorParams + +# --------------------------------------------------------------------------- # +# Конфигурация (можно переопределить через переменные окружения) +# --------------------------------------------------------------------------- # + +QDRANT_HOST: str = os.getenv("QDRANT_HOST", "localhost") +QDRANT_PORT: int = int(os.getenv("QDRANT_PORT", "6333")) +COLLECTION_NAME: str = os.getenv("QDRANT_COLLECTION", "rag_knowledge") + +EMBEDDING_MODEL: str = os.getenv("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text") +TEXT_SPLITTER_PARAMS = { + "chunk_size": 1000, + "chunk_overlap": 200, +} + +# --------------------------------------------------------------------------- # +# Класс клиента Qdrant +# --------------------------------------------------------------------------- # + +class QdrantClient: + """ + Обёртка над qdrant_client и langchain_qdrant для упрощения работы с коллекцией. + """ + + def __init__(self, host: str = QDRANT_HOST, port: int = QDRANT_PORT, + collection_name: str = COLLECTION_NAME) -> None: + self.host = host + self.port = port + self.collection_name = collection_name + + # Инициализируем низкоуровневый клиент + self.client = _QdrantClient(host=self.host, port=self.port) + + # Создаём коллекцию при необходимости + if not self._collection_exists(): + self._create_collection() + + # Обёртка LangChain для упрощённого добавления/поиска + self.vector_store = QdrantVectorStore( + client=self.client, + collection_name=self.collection_name, + embedding=OllamaEmbeddings(model=EMBEDDING_MODEL), + ) + + # ----------------------------------------------------------------------- # + # Внутренние вспомогательные методы + # ----------------------------------------------------------------------- # + + def _collection_exists(self) -> bool: + """Проверяем наличие коллекции.""" + try: + self.client.get_collection(name=self.collection_name) + return True + except Exception: + return False + + def _create_collection(self) -> None: + """Создаём коллекцию с параметрами по умолчанию.""" + vector_params = VectorParams(size=384, distance=Distance.COSINE) + self.client.create_collection( + collection_name=self.collection_name, + vectors_config=vector_params + ) + + # ----------------------------------------------------------------------- # + # Публичные методы + # ----------------------------------------------------------------------- # + + def add_documents(self, documents: Iterable[str], titles: Iterable[str]) -> None: + """ + Добавляет документы в коллекцию после разбиения на чанки. + + :param documents: Итерируемый список текстов. + :param titles: Заголовки документов (один к одному с documents). + """ + splitter = RecursiveCharacterTextSplitter(**TEXT_SPLITTER_PARAMS) + all_chunks: List[Tuple[str, str]] = [] + + for doc_text, title in zip(documents, titles): + chunks = splitter.split_text(doc_text) + # Для каждого чанка сохраняем заголовок как метаданные + all_chunks.extend([(chunk, title) for chunk in chunks]) + + texts, metas = zip(*all_chunks) + + self.vector_store.add_texts( + texts=texts, + metadatas=[{"title": meta} for meta in metas] + ) + + def search(self, query: str, limit: int = 5) -> List[dict]: + """ + Семантический поиск по запросу. + + :param query: Текст запроса. + :param limit: Количество результатов. + :return: Список словарей с полями 'text', 'score' и метаданными. + """ + results = self.vector_store.similarity_search_with_score(query, k=limit) + return [ + { + "text": r[0].page_content, + "score": r[1], + "metadata": r[0].metadata + } + for r in results + ] + + # ----------------------------------------------------------------------- # + # Утилиты для загрузки файлов из директории + # ----------------------------------------------------------------------- # + + def load_from_directory(self, directory: str | Path) -> None: + """ + Загружает все текстовые файлы из указанной папки в коллекцию. + + :param directory: Путь к каталогу. + """ + dir_path = Path(directory) + if not dir_path.is_dir(): + raise ValueError(f"Папка {directory} не существует") + + docs = [] + titles = [] + + for file in dir_path.glob("**/*.txt"): + text = file.read_text(encoding="utf-8") + docs.append(text) + titles.append(file.stem) + + self.add_documents(docs, titles) + + +# --------------------------------------------------------------------------- # +# Пример использования (можно удалить при импорте в другие модули) +# --------------------------------------------------------------------------- # + +if __name__ == "__main__": + client = QdrantClient() + # Загрузка из папки ./data + client.load_from_directory("./data") + + # Поиск по запросу + query_text = "Какой смысл жизни?" + results = client.search(query_text, limit=3) + for i, res in enumerate(results, 1): + print(f"\nРезультат {i}:") + print(f"Score: {res['score']:.4f}") + print(f"Title: {res['metadata'].get('title')}") + print(f"Text snippet: {res['text'][:200]}...") \ No newline at end of file