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(