Fix search reliability: batch dense embedding, graceful extra-query fallback, rerank 429 retry
- 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>
This commit is contained in:
parent
878971bb57
commit
68e2dc22c4
3 changed files with 53 additions and 21 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -14,6 +14,10 @@ async def get_rerank_scores(
|
|||
if not targets:
|
||||
return []
|
||||
|
||||
import asyncio as _asyncio
|
||||
|
||||
for attempt in range(5):
|
||||
try:
|
||||
response = await client.post(
|
||||
str(RERANKER_URL),
|
||||
**get_upstream_kwargs(),
|
||||
|
|
@ -24,10 +28,21 @@ 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)
|
||||
continue
|
||||
raise exc
|
||||
|
||||
return []
|
||||
|
||||
|
||||
async def rerank_points(
|
||||
|
|
|
|||
Loading…
Reference in a new issue