From 5d50a219bf023f9b52a23886f2a422a71aad5495 Mon Sep 17 00:00:00 2001
From: q <q@think>
Date: Sat, 18 Apr 2026 16:50:34 +0300
Subject: [PATCH] Clean up search: proper retry loop, batch embedding, remove
 duplicate wrapper

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 search/main.py          |  6 +++---
 search/query_builder.py |  7 +------
 search/rerank.py        | 42 ++++++++++++++++++++---------------------
 3 files changed, 25 insertions(+), 30 deletions(-)

diff --git a/search/main.py b/search/main.py
index f13cc5e..ac3acc5 100644
--- a/search/main.py
+++ b/search/main.py
@@ -25,7 +25,7 @@ from query_builder import (
     build_primary_query,
     build_sparse_query,
     embed_dense,
-    embed_dense_multi,
+    embed_dense_batch,
     embed_sparse,
 )
 from rerank import rerank_points
@@ -88,9 +88,9 @@ async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
     extra_dense_vecs: list[list[float]] = []
     if extra_texts:
         try:
-            extra_dense_vecs = await embed_dense_multi(client, extra_texts[:3])
+            extra_dense_vecs = await embed_dense_batch(client, extra_texts[:3])
         except Exception as exc:
-            logger.warning("Extra dense embedding failed, continuing without it: %s", exc)
+            logger.warning("Extra dense embedding failed, skipping: %s", exc)
 
     points = await qdrant_search(
         qdrant,
diff --git a/search/query_builder.py b/search/query_builder.py
index d7e083a..ee7338e 100644
--- a/search/query_builder.py
+++ b/search/query_builder.py
@@ -1,4 +1,3 @@
-import asyncio
 import os
 import re
 from functools import lru_cache
@@ -39,7 +38,7 @@ async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]:
 
 
 async def embed_dense_batch(client: httpx.AsyncClient, texts: list[str]) -> list[list[float]]:
-    """Single batch request for multiple texts (avoids N parallel calls)."""
+    """Single request for multiple texts — avoids N parallel calls and rate limiting."""
     response = await client.post(
         str(EMBEDDINGS_DENSE_URL),
         **get_upstream_kwargs(),
@@ -54,10 +53,6 @@ async def embed_dense_batch(client: httpx.AsyncClient, texts: list[str]) -> list
     return [item.embedding for item in payload.data]
 
 
-async def embed_dense_multi(client: httpx.AsyncClient, texts: list[str]) -> list[list[float]]:
-    return await embed_dense_batch(client, texts)
-
-
 def embed_sparse(text: str) -> SparseVector:
     vectors = list(get_sparse_model().embed([text]))
     if not vectors:
diff --git a/search/rerank.py b/search/rerank.py
index 276b414..2ad425a 100644
--- a/search/rerank.py
+++ b/search/rerank.py
@@ -1,3 +1,4 @@
+import asyncio
 from typing import Any
 
 import httpx
@@ -6,7 +7,7 @@ from config import RERANK_LIMIT, RERANKER_MODEL, RERANKER_URL, get_upstream_kwar
 from retrieval import extract_page_content
 
 
-async def get_rerank_scores(
+async def _get_rerank_scores(
     client: httpx.AsyncClient,
     query: str,
     targets: list[str],
@@ -14,8 +15,6 @@ async def get_rerank_scores(
     if not targets:
         return []
 
-    import asyncio as _asyncio
-
     for attempt in range(5):
         try:
             response = await client.post(
@@ -28,20 +27,23 @@ async def get_rerank_scores(
                     "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)
+                await asyncio.sleep(2 ** attempt)
                 continue
             raise exc
 
+        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]
+
+    logger.error("Rerank 429 after all retries, falling back")
     return []
 
 
@@ -50,24 +52,22 @@ async def rerank_points(
     query: str,
     points: list[Any],
 ) -> tuple[list[Any], list[Any]]:
-    """Return (reranked_head, retrieval_tail) so we don't lose candidates."""
     if not points:
         return [], []
 
-    rerank_candidates = points[:RERANK_LIMIT]
+    head = points[:RERANK_LIMIT]
     tail = points[RERANK_LIMIT:]
+    targets = [extract_page_content(p) for p in head]
 
-    targets = [extract_page_content(p) for p in rerank_candidates]
     try:
-        scores = await get_rerank_scores(client, query, targets)
+        scores = await _get_rerank_scores(client, query, targets)
     except Exception as exc:
         logger.warning("Rerank failed, using retrieval order: %s", exc)
-        return rerank_candidates, tail
+        return head, tail
 
-    if len(scores) != len(rerank_candidates):
+    if len(scores) != len(head):
         logger.warning("Rerank score count mismatch, using retrieval order")
-        return rerank_candidates, tail
+        return head, tail
 
-    paired = sorted(zip(scores, rerank_candidates), key=lambda x: x[0], reverse=True)
-    reranked = [p for _, p in paired]
+    reranked = [p for _, p in sorted(zip(scores, head), key=lambda x: x[0], reverse=True)]
     return reranked, tail