# -*- coding: utf-8 -*- """ Сервис для взаимодействия с векторной базой данных Qdrant. Обеспечивает создание гибридной коллекции (dense + sparse), добавление точек, гибридный поиск с RRF, полнотекстовый поиск и удаление по doc_id. ВНИМАНИЕ: Этот сервис НЕ использует генерацию эмбеддингов на лету через models.Document. Вместо этого он получает готовые dense-векторы из EmbeddingService (GigaChat) и sparse-векторы из FastEmbed (Qdrant/bm25). ВАЖНО: В версиях qdrant-client 1.10+ метод search() заменён на query_points(). В этом файле используется правильный API для вашей версии. """ import logging from typing import List, Dict, Any, Optional from uuid import uuid4 from qdrant_client import QdrantClient, models logger = logging.getLogger(__name__) class QdrantService: def __init__( self, host: str, port: int, grpc_port: int, collection_name: str, vector_size: int, distance: str, prefer_grpc: bool = False ): """ Инициализация сервиса Qdrant. Аргументы: host: хост Qdrant (обычно localhost) port: HTTP порт Qdrant (обычно 6333) grpc_port: gRPC порт Qdrant (обычно 6334) collection_name: имя коллекции в Qdrant vector_size: размерность плотных векторов (1024 для GigaChat Embeddings) distance: метрика расстояния (Cosine или Euclid) prefer_grpc: использовать ли gRPC вместо HTTP (по умолчанию False) """ self.host = host self.port = port self.collection_name = collection_name self.vector_size = vector_size self.distance = models.Distance.COSINE if distance.lower() == "cosine" else models.Distance.EUCLID # Имена векторных полей в коллекции self.dense_vector_name = "dense" self.sparse_vector_name = "sparse" # Создаём клиент Qdrant self.client = QdrantClient(host=host, port=port, grpc_port=grpc_port, prefer_grpc=prefer_grpc) self._ensure_collection() logger.info(f"QdrantService инициализирован: коллекция {collection_name}, размер dense {vector_size}") def _ensure_collection(self) -> None: """ Проверяет существование коллекции. Если нет — создаёт её с поддержкой dense и sparse векторов. """ if not self.client.collection_exists(self.collection_name): logger.info(f"Коллекция {self.collection_name} не найдена. Создаю новую (dense + sparse)...") try: # Конфигурация dense-векторов (семантический поиск) vectors_config = { self.dense_vector_name: models.VectorParams( size=self.vector_size, distance=self.distance, ) } # Конфигурация sparse-векторов (ключевой поиск BM25 с IDF) sparse_vectors_config = { self.sparse_vector_name: models.SparseVectorParams( modifier=models.Modifier.IDF ) } self.client.create_collection( collection_name=self.collection_name, vectors_config=vectors_config, sparse_vectors_config=sparse_vectors_config ) logger.info(f"Коллекция {self.collection_name} создана (dense size={self.vector_size}, sparse с IDF)") except Exception as e: # Если библиотека не поддерживает sparse (старая версия), создаём только dense logger.warning(f"Не удалось создать sparse-векторы: {e}. Создаю коллекцию только с dense.") self.client.create_collection( collection_name=self.collection_name, vectors_config={ self.dense_vector_name: models.VectorParams( size=self.vector_size, distance=self.distance, ) } ) logger.info(f"Коллекция {self.collection_name} создана (только dense size={self.vector_size})") # После создания коллекции добавляем payload-индексы self._create_payload_indexes() else: # Коллекция уже существует – проверяем, есть ли sparse-поле (логируем предупреждение, если нет) try: collection_info = self.client.get_collection(self.collection_name) if not hasattr(collection_info, 'config') or not hasattr(collection_info.config, 'params'): logger.debug("Не удалось проверить наличие sparse-поля – пропускаем") else: sparse_exists = (hasattr(collection_info.config.params, 'sparse_vectors') and self.sparse_vector_name in collection_info.config.params.sparse_vectors) if not sparse_exists: logger.warning( f"Коллекция {self.collection_name} не содержит sparse-векторного поля. " "Гибридный поиск будет работать только с dense. " "Рекомендуется пересоздать коллекцию с помощью --force или вручную." ) except Exception as e: logger.debug(f"Не удалось проверить конфигурацию коллекции: {e}") logger.debug(f"Коллекция {self.collection_name} уже существует") def _create_payload_indexes(self) -> None: """ Создаёт индексы для полей payload для ускорения фильтрации и поиска. Индексы нужны для полей: - user_jid: для доступа к личным документам (keyword) - room_jid: для доступа к документам комнаты (keyword) - is_global: для доступа к глобальным документам (keyword) - doc_id: для быстрого удаления всех чанков документа (integer) - text: для полнотекстового поиска (text с настройками tokenizer=whitespace, lowercase, phrase_matching) """ logger.info("Создаю payload-индексы для коллекции...") # Индексы для фильтрации (keyword) for field in ["user_jid", "room_jid", "is_global"]: try: self.client.create_payload_index( collection_name=self.collection_name, field_name=field, field_schema="keyword" ) logger.debug(f"Индекс keyword для поля '{field}' создан") except Exception as e: logger.warning(f"Не удалось создать индекс для '{field}': {e}") # Индекс для doc_id (integer) try: self.client.create_payload_index( collection_name=self.collection_name, field_name="doc_id", field_schema="integer" ) logger.debug("Индекс integer для поля 'doc_id' создан") except Exception as e: logger.warning(f"Не удалось создать индекс для 'doc_id': {e}") # Полнотекстовый индекс для текста с детальными параметрами try: # Используем расширенную конфигурацию (tokenizer, lowercase, phrase_matching) self.client.create_payload_index( collection_name=self.collection_name, field_name="text", field_schema=models.TextIndexParams( tokenizer="whitespace", lowercase=True, min_token_len=1, max_token_len=0, # 0 = без ограничения phrase_matching=True ) ) logger.info("Полнотекстовый индекс для поля 'text' создан (tokenizer=whitespace, lowercase, phrase_matching)") except Exception as e: # Fallback для старых версий qdrant-client (без TextIndexParams) logger.warning(f"Не удалось создать расширенный индекс text: {e}. Пробую упрощённый.") try: self.client.create_payload_index( collection_name=self.collection_name, field_name="text", field_schema="text" ) logger.info("Полнотекстовый индекс для поля 'text' создан (упрощённый)") except Exception as e2: logger.error(f"Не удалось создать индекс для 'text': {e2}") logger.info(f"Payload-индексы для коллекции {self.collection_name} созданы") def add_chunk(self, dense: List[float], sparse: Any, payload: Dict[str, Any]) -> str: """ Добавляет точку (чанк) в коллекцию с двумя типами векторов: dense и sparse. Аргументы: dense: плотный вектор (1024) от GigaChat sparse: разрежённый вектор от FastEmbed (SparseEmbedding) или None payload: метаданные чанка Возвращает: str: уникальный идентификатор точки """ point_id = str(uuid4()) # Преобразуем sparse-вектор из формата FastEmbed в формат Qdrant sparse_vector = None if sparse is not None: if hasattr(sparse, 'indices') and hasattr(sparse, 'values'): # FastEmbed SparseEmbedding -> SparseVector sparse_vector = models.SparseVector( indices=sparse.indices.tolist(), values=sparse.values.tolist() ) elif isinstance(sparse, dict) and 'indices' in sparse: # Уже словарь sparse_vector = models.SparseVector(**sparse) else: logger.warning(f"Неизвестный формат sparse: {type(sparse)}") # Если sparse-вектор не задан, передаём пустой (чтобы не нарушить схему коллекции) if sparse_vector is None: sparse_vector = models.SparseVector(indices=[], values=[]) point = models.PointStruct( id=point_id, vector={ self.dense_vector_name: dense, self.sparse_vector_name: sparse_vector }, payload=payload ) self.client.upsert(collection_name=self.collection_name, points=[point]) logger.debug(f"Чанк добавлен: id={point_id}, doc_id={payload.get('doc_id')}") return point_id def search_hybrid( self, dense_vector: List[float], sparse_vector: Any, user_jid: str, room_jid: str = None, top_k: int = 15, rrf_k: int = 60 ) -> List[Dict]: """ Гибридный поиск: объединяет результаты dense и sparse поиска через Reciprocal Rank Fusion (RRF). Аргументы: dense_vector: плотный вектор запроса (1024) от GigaChat sparse_vector: разрежённый вектор запроса от FastEmbed (SparseEmbedding) или None user_jid: JID пользователя для фильтрации доступа room_jid: JID комнаты (None для личного чата) top_k: количество возвращаемых результатов rrf_k: константа сглаживания в RRF (обычно 60) Возвращает: List[Dict]: список результатов с ключами score, payload, id """ # --- 1. Фильтр доступа --- if room_jid: # В комнате: документы комнаты ИЛИ глобальные filter_cond = models.Filter( should=[ models.FieldCondition(key="room_jid", match=models.MatchValue(value=room_jid)), models.FieldCondition(key="is_global", match=models.MatchValue(value=True)) ] ) else: # В личном чате: личные документы пользователя ИЛИ глобальные filter_cond = models.Filter( should=[ models.FieldCondition(key="user_jid", match=models.MatchValue(value=user_jid)), models.FieldCondition(key="is_global", match=models.MatchValue(value=True)) ] ) # --- 2. Dense-поиск (семантический) --- dense_resp = self.client.query_points( collection_name=self.collection_name, query=dense_vector, using=self.dense_vector_name, limit=top_k, query_filter=filter_cond, with_payload=True ) # --- 3. Sparse-поиск (ключевой) с преобразованием формата --- sparse_resp = None if sparse_vector is not None: # Преобразуем sparse-вектор аналогично методу add_chunk sparse_query = None if hasattr(sparse_vector, 'indices') and hasattr(sparse_vector, 'values'): # FastEmbed SparseEmbedding -> SparseVector sparse_query = models.SparseVector( indices=sparse_vector.indices.tolist(), values=sparse_vector.values.tolist() ) elif isinstance(sparse_vector, dict) and 'indices' in sparse_vector: sparse_query = models.SparseVector(**sparse_vector) elif isinstance(sparse_vector, models.SparseVector): sparse_query = sparse_vector else: logger.warning(f"Неизвестный формат sparse_vector: {type(sparse_vector)}") if sparse_query is not None: sparse_resp = self.client.query_points( collection_name=self.collection_name, query=sparse_query, using=self.sparse_vector_name, limit=top_k, query_filter=filter_cond, with_payload=True ) # --- 4. Reciprocal Rank Fusion (RRF) --- scores = {} # point_id -> {score, payload, id} for rank, hit in enumerate(dense_resp.points): rrf_score = 1.0 / (rrf_k + rank + 1) scores[hit.id] = { "score": rrf_score, "payload": hit.payload, "id": hit.id } if sparse_resp: for rank, hit in enumerate(sparse_resp.points): rrf_score = 1.0 / (rrf_k + rank + 1) if hit.id in scores: scores[hit.id]["score"] += rrf_score else: scores[hit.id] = { "score": rrf_score, "payload": hit.payload, "id": hit.id } # --- 5. Сортировка и возврат top_k --- sorted_items = sorted(scores.values(), key=lambda x: x["score"], reverse=True) return sorted_items[:top_k] def search_keywords(self, query_text: str, user_jid: str, room_jid: str = None, top_k: int = 15) -> List[Dict]: """ Выполняет полнотекстовый поиск по полю text (для точных запросов). Аргументы: query_text: текст запроса (ищется как подстрока в text) user_jid: JID пользователя для фильтрации room_jid: JID комнаты (None для личного чата) top_k: количество возвращаемых результатов Возвращает: List[Dict]: список результатов с ключами score, payload, id Используется для точных запросов типа "выведи пункт 1.1", где семантический поиск может не справиться. """ # Формируем фильтр доступа if room_jid: filter_cond = models.Filter( must=[ models.FieldCondition(key="room_jid", match=models.MatchValue(value=room_jid)) ], should=[ models.FieldCondition(key="is_global", match=models.MatchValue(value=True)) ] ) else: filter_cond = models.Filter( must=[ models.FieldCondition(key="user_jid", match=models.MatchValue(value=user_jid)) ], should=[ models.FieldCondition(key="is_global", match=models.MatchValue(value=True)) ] ) # Добавляем условие полнотекстового поиска (MatchText) filter_cond.must.append( models.FieldCondition( key="text", match=models.MatchText(text=query_text) ) ) # Получаем точки через scroll points = [] offset = None while len(points) < top_k: scroll_result = self.client.scroll( collection_name=self.collection_name, scroll_filter=filter_cond, limit=top_k, offset=offset, with_payload=True ) points.extend(scroll_result[0]) offset = scroll_result[1] if offset is None: break # Простая сортировка по количеству совпавших ключевых слов (score не от Qdrant, а наша эвристика) query_words = set(query_text.lower().split()) for point in points: text = point.payload.get('text', '').lower() matches = sum(1 for word in query_words if word in text) point._match_score = matches points.sort(key=lambda p: getattr(p, '_match_score', 0), reverse=True) return [ { "score": getattr(p, '_match_score', 0), "payload": p.payload, "id": p.id } for p in points[:top_k] ] def delete_by_doc_id(self, doc_id: int) -> None: """ Удаляет все точки (чанки), принадлежащие указанному документу. Аргументы: doc_id: идентификатор документа в PostgreSQL """ filter_cond = models.Filter( must=[ models.FieldCondition(key="doc_id", match=models.MatchValue(value=doc_id)) ] ) points = [] offset = None while True: scroll_result = self.client.scroll( collection_name=self.collection_name, scroll_filter=filter_cond, limit=100, offset=offset ) points.extend(scroll_result[0]) offset = scroll_result[1] if offset is None: break ids = [p.id for p in points] if ids: self.client.delete(collection_name=self.collection_name, points_selector=ids) logger.info(f"Удалено {len(ids)} чанков для doc_id={doc_id}") def delete_collection(self) -> None: """Удаляет всю коллекцию (опасная операция, используется редко).""" if self.client.collection_exists(self.collection_name): self.client.delete_collection(self.collection_name) logger.warning(f"Коллекция {self.collection_name} удалена")