From e244b8bf307a62b6d230f604794343dc7e061a29 Mon Sep 17 00:00:00 2001
From: q <q@think>
Date: Sat, 18 Apr 2026 16:47:52 +0300
Subject: [PATCH] Fix search reliability: batch dense embedding, graceful
 extra-query fallback, rerank 429 retry
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- embed_dense_multi now sends one batch request (N texts → 1 API call) instead of N parallel
  requests, avoiding rate-limit errors when question has variants/hyde
- Extra dense embeddings (variants/hyde) wrapped in try/except so primary query always succeeds
- Reranker now retries up to 5 times with exponential backoff on 429, matching Lotus reference

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 search/main.py          | 14 ++++++++------
 search/query_builder.py | 19 +++++++++++++++++--
 search/rerank.py        | 41 ++++++++++++++++++++++++++++-------------
 3 files changed, 53 insertions(+), 21 deletions(-)

diff --git a/search/main.py b/search/main.py
index 27fbd92..f13cc5e 100644
--- a/search/main.py
+++ b/search/main.py
@@ -80,16 +80,18 @@ async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
     extra_texts = build_extra_dense_queries(question)
     sparse_text = build_sparse_query(question)
 
-    async def _no_extra() -> list:
-        return []
-
-    extra_task = embed_dense_multi(client, extra_texts) if extra_texts else _no_extra()
-    primary_dense, extra_dense_vecs, sparse_vec = await asyncio.gather(
+    primary_dense, sparse_vec = await asyncio.gather(
         _embed_dense_with_retry(client, primary_query),
-        extra_task,
         asyncio.to_thread(embed_sparse, sparse_text),
     )
 
+    extra_dense_vecs: list[list[float]] = []
+    if extra_texts:
+        try:
+            extra_dense_vecs = await embed_dense_multi(client, extra_texts[:3])
+        except Exception as exc:
+            logger.warning("Extra dense embedding failed, continuing without it: %s", exc)
+
     points = await qdrant_search(
         qdrant,
         primary_dense,
diff --git a/search/query_builder.py b/search/query_builder.py
index 1a9fa61..d7e083a 100644
--- a/search/query_builder.py
+++ b/search/query_builder.py
@@ -38,9 +38,24 @@ async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]:
     return payload.data[0].embedding
 
 
+async def embed_dense_batch(client: httpx.AsyncClient, texts: list[str]) -> list[list[float]]:
+    """Single batch request for multiple texts (avoids N parallel calls)."""
+    response = await client.post(
+        str(EMBEDDINGS_DENSE_URL),
+        **get_upstream_kwargs(),
+        json={
+            "model": os.getenv("EMBEDDINGS_DENSE_MODEL", EMBEDDINGS_DENSE_MODEL),
+            "input": texts,
+        },
+    )
+    response.raise_for_status()
+    payload = DenseEmbeddingResponse.model_validate(response.json())
+    payload.data.sort(key=lambda x: x.index)
+    return [item.embedding for item in payload.data]
+
+
 async def embed_dense_multi(client: httpx.AsyncClient, texts: list[str]) -> list[list[float]]:
-    tasks = [embed_dense(client, t) for t in texts]
-    return list(await asyncio.gather(*tasks))
+    return await embed_dense_batch(client, texts)
 
 
 def embed_sparse(text: str) -> SparseVector:
diff --git a/search/rerank.py b/search/rerank.py
index a7c38b8..276b414 100644
--- a/search/rerank.py
+++ b/search/rerank.py
@@ -14,20 +14,35 @@ async def get_rerank_scores(
     if not targets:
         return []
 
-    response = await client.post(
-        str(RERANKER_URL),
-        **get_upstream_kwargs(),
-        json={
-            "model": RERANKER_MODEL,
-            "encoding_format": "float",
-            "text_1": query,
-            "text_2": targets,
-        },
-    )
-    response.raise_for_status()
+    import asyncio as _asyncio
 
-    data = response.json().get("data") or []
-    return [float(sample["score"]) for sample in data]
+    for attempt in range(5):
+        try:
+            response = await client.post(
+                str(RERANKER_URL),
+                **get_upstream_kwargs(),
+                json={
+                    "model": RERANKER_MODEL,
+                    "encoding_format": "float",
+                    "text_1": query,
+                    "text_2": targets,
+                },
+            )
+            if response.status_code == 429:
+                wait = 2 ** attempt
+                logger.warning("Rerank 429, retry %d/5 in %ds", attempt + 1, wait)
+                await _asyncio.sleep(wait)
+                continue
+            response.raise_for_status()
+            data = response.json().get("data") or []
+            return [float(sample["score"]) for sample in data]
+        except Exception as exc:
+            if attempt < 4:
+                await _asyncio.sleep(2 ** attempt)
+                continue
+            raise exc
+
+    return []
 
 
 async def rerank_points(