vk_hackathon/search/retrieval.py
q f9f1b332bd Fix 500 error: replace datetime_range with range in Qdrant filter
qdrant-client 1.15.1 does not support datetime_range in FieldCondition.
Use models.Range with string comparison (same as Lotus reference).
Also wrap date filter in try-except to prevent crash on bad date format.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-18 22:01:34 +03:00

119 lines
3.2 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:
try:
must_conditions.append(
models.FieldCondition(
key="metadata.end",
range=models.Range(gte=question.date_range.from_),
)
)
must_conditions.append(
models.FieldCondition(
key="metadata.start",
range=models.Range(lte=question.date_range.to),
)
)
except Exception as e:
logger.warning("Date filter failed: %s", e)
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 ""