452 lines
22 KiB
Python
452 lines
22 KiB
Python
# -*- 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} удалена") |