.ai_update/changes.md

@@ -0,0 +1,27 @@
+# AI Update Log
+
+## Scope
+- File: `search/main.py`
+- Purpose: fixed and extended retrieval/rerank pipeline according to TODO items.
+
+## Done Changes
+- Switched base query selection to `question.search_text` with fallback to `question.text`.
+- Added support for `question.variants` as additional query variants in retrieval.
+- Added support for `question.hyde` as additional dense-only queries.
+- Added support for `question.keywords` as the primary source for sparse query text with fallback to current query variant.
+- Stopped losing retrieval candidates after rerank: rerank is applied to head (`RERANK_LIMIT`), tail candidates are preserved.
+- Added deduplication of retrieval points by Qdrant point id before rerank.
+- Implemented score aggregation by `message_id` (sum of chunk scores mapped to same message).
+- Limited final response to `top-50` message ids via `FINAL_TOP_K = 50`.
+- Final output message ids are now selected from aggregated scores (sorted by score desc, tie-break by message_id).
+
+## Notes
+- Earlier step introduced direct `message_ids` deduplication before response.
+- Current logic supersedes this by ranking and selecting unique `message_id` values from aggregated scores.
+
+## Verification
+- Syntax check passed after each main change: `python -m py_compile search/main.py`.
+
+## How To Use This Log
+- Treat this file as the source of truth for already completed `search/main.py` tasks.
+- On next tasks, read this file first to avoid duplicate edits.

search/main.py

@@ -175,6 +175,7 @@ DENSE_PREFETCH_K = 10
 SPRASE_PREFETCH_K = 30
 RETRIEVE_K = 20
 RERANK_LIMIT = 10
+FINAL_TOP_K = 50
 
 async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]:
     # Dense endpoint ожидает OpenAI-compatible body с input как списком строк.
@@ -359,6 +360,25 @@ def collect_hyde_queries(question: Question, base_queries: list[str]) -> list[st
     return hyde_queries
 
 
+def build_sparse_query_text(question: Question, fallback_query: str) -> str:
+    keywords: list[str] = []
+    seen: set[str] = set()
+
+    for keyword in question.keywords or []:
+        normalized = keyword.strip()
+        if not normalized:
+            continue
+        if normalized in seen:
+            continue
+        seen.add(normalized)
+        keywords.append(normalized)
+
+    if keywords:
+        return " ".join(keywords)
+
+    return fallback_query
+
+
 def deduplicate_points(points: list[Any]) -> list[Any]:
     unique_points: list[Any] = []
     seen_ids: set[str] = set()
@@ -375,6 +395,13 @@ def deduplicate_points(points: list[Any]) -> list[Any]:
     return unique_points
 
 
+def extract_point_score(point: Any) -> float:
+    score = getattr(point, "score", 0.0)
+    if score is None:
+        return 0.0
+    return float(score)
+
+
 def extract_message_ids(point: Any) -> list[str]:
     payload = point.payload or {}
     metadata = payload.get("metadata") or {}
@@ -414,21 +441,49 @@ async def rerank_points(
     client: httpx.AsyncClient,
     query: str,
     points: list[Any],
-) -> list[Any]:
-    rerank_candidates = points[:10]
+) -> list[tuple[Any, float]]:
+    rerank_candidates = points[:RERANK_LIMIT]
+    tail_candidates = points[RERANK_LIMIT:]
     rerank_targets = [point.payload.get("page_content") for point in rerank_candidates]
     scores = await get_rerank_scores(client, query, rerank_targets)
 
     reranked_candidates = [
-        point
-        for _, point in sorted(
+        (point, float(score))
+        for score, point in sorted(
             zip(scores, rerank_candidates, strict=True),
             key=lambda item: item[0],
             reverse=True,
         )
     ]
+    tail_with_scores = [
+        (point, extract_point_score(point))
+        for _, point in sorted(
+            [(extract_point_score(point), point) for point in tail_candidates],
+            key=lambda item: item[0],
+            reverse=True,
+        )
+    ]
 
-    return reranked_candidates
+    return reranked_candidates + tail_with_scores
+
+
+def aggregate_message_scores(scored_points: list[tuple[Any, float]]) -> dict[str, float]:
+    aggregated_scores: dict[str, float] = {}
+
+    for point, point_score in scored_points:
+        point_message_ids = set(extract_message_ids(point))
+        for message_id in point_message_ids:
+            aggregated_scores[message_id] = aggregated_scores.get(message_id, 0.0) + point_score
+
+    return aggregated_scores
+
+
+def select_top_message_ids(aggregated_scores: dict[str, float], limit: int) -> list[str]:
+    sorted_items = sorted(
+        aggregated_scores.items(),
+        key=lambda item: (-item[1], item[0]),
+    )
+    return [message_id for message_id, _ in sorted_items[:limit]]
 
 
 # Ваш сервис должен имплементировать оба этих метода
@@ -451,7 +506,8 @@ async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
     all_points: list[Any] = []
     for query_variant in queries:
         dense_vector = await embed_dense(client, query_variant)
-        sparse_vector = await embed_sparse(query_variant)
+        sparse_query_text = build_sparse_query_text(payload.question, query_variant)
+        sparse_vector = await embed_sparse(sparse_query_text)
         points = await qdrant_search(qdrant, dense_vector, sparse_vector, payload.question)
         if points:
             all_points.extend(list(points))
@@ -466,11 +522,9 @@ async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
     if not best_points:
         return SearchAPIResponse(results=[])
 
-    best_points = await rerank_points(client, query, list(best_points))
-
-    message_ids: list[str] = [] 
-    for point in best_points:
-        message_ids += extract_message_ids(point)
+    scored_points = await rerank_points(client, query, list(best_points))
+    aggregated_scores = aggregate_message_scores(scored_points)
+    message_ids = select_top_message_ids(aggregated_scores, FINAL_TOP_K)
 
     return SearchAPIResponse(
         results=[SearchAPIItem(message_ids=message_ids)]