Update file qdrant_service.py
This commit is contained in:
Markov Andrey
@@ -0,0 +1,452 @@
|
||||
# -*- 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} удалена")
|
||||
Reference in New Issue
Block a user