forked from zovos/vk_hackathon
113 lines
3 KiB
Python
113 lines
3 KiB
Python
from typing import Any
|
|
|
|
from qdrant_client import AsyncQdrantClient, models
|
|
|
|
from config import (
|
|
DENSE_PREFETCH_K,
|
|
QDRANT_COLLECTION_NAME,
|
|
QDRANT_DENSE_VECTOR_NAME,
|
|
QDRANT_SPARSE_VECTOR_NAME,
|
|
RETRIEVE_K,
|
|
SPARSE_PREFETCH_K,
|
|
logger,
|
|
)
|
|
from schemas import Question, SparseVector
|
|
|
|
|
|
def _build_filter(question: Question) -> models.Filter | None:
|
|
must_conditions: list[models.Condition] = []
|
|
|
|
if question.date_range:
|
|
must_conditions.append(
|
|
models.FieldCondition(
|
|
key="metadata.start",
|
|
datetime_range=models.DatetimeRange(
|
|
gte=question.date_range.from_,
|
|
lte=question.date_range.to,
|
|
),
|
|
)
|
|
)
|
|
|
|
if question.asker:
|
|
must_conditions.append(
|
|
models.FieldCondition(
|
|
key="metadata.participants",
|
|
match=models.MatchValue(value=question.asker),
|
|
)
|
|
)
|
|
|
|
return models.Filter(must=must_conditions) if must_conditions else None
|
|
|
|
|
|
async def qdrant_search(
|
|
client: AsyncQdrantClient,
|
|
primary_dense: list[float],
|
|
extra_dense: list[list[float]],
|
|
sparse_vector: SparseVector,
|
|
question: Question,
|
|
) -> list[Any]:
|
|
search_filter = _build_filter(question)
|
|
|
|
prefetch: list[models.Prefetch] = []
|
|
|
|
# Primary dense
|
|
prefetch.append(
|
|
models.Prefetch(
|
|
query=primary_dense,
|
|
using=QDRANT_DENSE_VECTOR_NAME,
|
|
limit=DENSE_PREFETCH_K,
|
|
filter=search_filter,
|
|
)
|
|
)
|
|
|
|
# Extra dense (variants / hyde) - smaller budget per query
|
|
extra_k = max(10, DENSE_PREFETCH_K // max(1, len(extra_dense)))
|
|
for vec in extra_dense:
|
|
prefetch.append(
|
|
models.Prefetch(
|
|
query=vec,
|
|
using=QDRANT_DENSE_VECTOR_NAME,
|
|
limit=extra_k,
|
|
filter=search_filter,
|
|
)
|
|
)
|
|
|
|
# Sparse
|
|
prefetch.append(
|
|
models.Prefetch(
|
|
query=models.SparseVector(
|
|
indices=sparse_vector.indices,
|
|
values=sparse_vector.values,
|
|
),
|
|
using=QDRANT_SPARSE_VECTOR_NAME,
|
|
limit=SPARSE_PREFETCH_K,
|
|
filter=search_filter,
|
|
)
|
|
)
|
|
|
|
response = await client.query_points(
|
|
collection_name=QDRANT_COLLECTION_NAME,
|
|
prefetch=prefetch,
|
|
query=models.FusionQuery(fusion=models.Fusion.RRF),
|
|
limit=RETRIEVE_K,
|
|
with_payload=True,
|
|
)
|
|
|
|
if not response.points:
|
|
logger.debug("Qdrant returned 0 points")
|
|
return []
|
|
|
|
logger.debug("Qdrant returned %d points", len(response.points))
|
|
return list(response.points)
|
|
|
|
|
|
def extract_message_ids(point: Any) -> list[str]:
|
|
payload = point.payload or {}
|
|
metadata = payload.get("metadata") or {}
|
|
message_ids = metadata.get("message_ids") or []
|
|
return [str(mid) for mid in message_ids]
|
|
|
|
|
|
def extract_page_content(point: Any) -> str:
|
|
payload = point.payload or {}
|
|
return payload.get("page_content") or ""
|