fix(usage): embed query once across doc_types; address review round 1
Round-1 claude-review findings: - 🔴 Multi-doc_type search billed N embedding calls as 1. nc_semantic_search loops search() once per doc_type on one BM25HybridSearchAlgorithm instance, and each call re-embedded the query, so only the last query_token_count was recorded. Cache the dense embedding per query on the (per-request) instance so the query is embedded — and metered — exactly once regardless of how many doc_types are searched. This also removes the redundant per-type embed work and avoids billing a user N× for one logical query. - 🟡 Ollama embed() now delegates to embed_with_usage() so single and batch embeds use the same /api/embed endpoint (was the legacy /api/embeddings), keeping _detect_dimension and other embed() callers consistent. - 🟢 round() instead of truncating int() when coercing provider-reported token counts (forward-compatible if a provider ever returns a float). Tests: per-instance query-embedding cache (embedded once across 3 doc_types; re-embeds on a different query). Deferred (stated on the PR): mistral/openai single-embed dual path (changes tested error/request semantics on the cloud-critical path — separate refactor), bedrock boto3 sync-in-async (pre-existing; no new invoke_model calls per doc). Deck #67. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
64318f0b25
commit
a0bb5642cb
@@ -194,7 +194,7 @@ class BedrockProvider(Provider):
|
||||
|
||||
token_count = response_body.get("inputTextTokenCount")
|
||||
tokens = (
|
||||
int(token_count)
|
||||
round(token_count)
|
||||
if isinstance(token_count, (int, float))
|
||||
else self._estimate_tokens([text])
|
||||
)
|
||||
|
||||
@@ -209,7 +209,7 @@ class MistralProvider(Provider):
|
||||
# gives an int, but test doubles / partial responses can surface a
|
||||
# non-numeric attribute — fall back to the estimate there.
|
||||
tokens = (
|
||||
int(total_tokens)
|
||||
round(total_tokens)
|
||||
if isinstance(total_tokens, (int, float))
|
||||
else self._estimate_tokens(batch)
|
||||
)
|
||||
|
||||
@@ -82,17 +82,12 @@ class OllamaProvider(Provider):
|
||||
Raises:
|
||||
NotImplementedError: If embeddings not enabled (no embedding_model)
|
||||
"""
|
||||
if not self.supports_embeddings:
|
||||
raise NotImplementedError(
|
||||
"Embedding not supported - no embedding_model configured"
|
||||
)
|
||||
|
||||
response = await self.client.post(
|
||||
f"{self.base_url}/api/embeddings",
|
||||
json={"model": self.embedding_model, "prompt": text},
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()["embedding"]
|
||||
# Delegate to embed_with_usage so single and batch embeds use the same
|
||||
# /api/embed endpoint (the legacy /api/embeddings differs in payload and
|
||||
# omits prompt_eval_count). _detect_dimension() and other embed() callers
|
||||
# therefore stay consistent with the search/indexing path.
|
||||
embedding, _ = await self.embed_with_usage(text)
|
||||
return embedding
|
||||
|
||||
async def embed_batch(
|
||||
self, texts: list[str], batch_size: int = 32
|
||||
@@ -166,7 +161,7 @@ class OllamaProvider(Provider):
|
||||
|
||||
prompt_eval = data.get("prompt_eval_count")
|
||||
total_tokens += (
|
||||
int(prompt_eval)
|
||||
round(prompt_eval)
|
||||
if isinstance(prompt_eval, (int, float))
|
||||
else self._estimate_tokens(batch)
|
||||
)
|
||||
|
||||
@@ -233,7 +233,7 @@ class OpenAIProvider(Provider):
|
||||
# gives an int, but test doubles / partial responses can surface a
|
||||
# non-numeric attribute — fall back to the estimate there.
|
||||
tokens = (
|
||||
int(total_tokens)
|
||||
round(total_tokens)
|
||||
if isinstance(total_tokens, (int, float))
|
||||
else self._estimate_tokens(batch)
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user