feat(usage): meter embedding tokens as embeddings_queries on both paths

embeddings_queries now records the embedding request's token count (the unit
upstream providers bill on) instead of an operation count, and fires on the
indexing path too. Previously only semantic search recorded it (value=1), so a
re-indexing run produced no embeddings_queries events at all — only pages_chunks.

- Provider layer: additive embed_with_usage / embed_batch_with_usage surface the
  per-request token count (Mistral/OpenAI usage.total_tokens, Bedrock Titan
  inputTextTokenCount, Ollama prompt_eval_count); a char-based estimate is the
  fallback (Simple, and any provider/response without a token field). Gateway and
  EmbeddingService forward through. The count travels as a return value / a
  per-request SearchAlgorithm attribute — never on the singleton — so concurrent
  indexing + search can't mis-attribute bills.
- Indexing (vector/processor.py): records embeddings_queries (value=batch tokens)
  alongside the existing pages_chunks event.
- Search (server/semantic.py): value is now the query embedding's token count,
  relayed from BM25HybridSearchAlgorithm via query_token_count.

The astrolabe_embeddings_queries Stripe meter (sum aggregation) now sums tokens
with no CP/Terraform change. The meter "queries"->tokens naming/unit
clarification (homelab-terraform #254) + CP rollup/portal copy is a follow-up.

Deck #67.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Chris Coutinho
2026-06-08 00:53:58 +02:00
co-authored by Claude Opus 4.8
parent fe17994c4d
commit 64318f0b25
17 changed files with 647 additions and 35 deletions
+49 -6
View File
@@ -122,17 +122,42 @@ class MistralProvider(Provider):
async def embed_batch(self, texts: list[str]) -> list[list[float]]:
"""Generate embeddings for multiple texts, chunking by ``BATCH_SIZE``."""
embeddings, _ = await self.embed_batch_with_usage(texts)
return embeddings
async def embed_with_usage(self, text: str) -> tuple[list[float], int]:
"""Embed one text, reporting the Mistral request's token count."""
embeddings, tokens = await self.embed_batch_with_usage([text])
if not embeddings:
raise RuntimeError(
f"Mistral embeddings API returned no embedding for model "
f"{self.embedding_model}"
)
return embeddings[0], tokens
async def embed_batch_with_usage(
self, texts: list[str]
) -> tuple[list[list[float]], int]:
"""Embed multiple texts, summing the Mistral-reported token usage.
Returns ``(embeddings, total_tokens)`` where ``total_tokens`` is the
sum of ``response.usage.total_tokens`` across the ``BATCH_SIZE`` sub-
requests (the unit Mistral bills on). Used by the usage-metering hooks
to record ``embeddings_queries`` by tokens (Deck #67).
"""
if not self.supports_embeddings:
raise NotImplementedError(_NO_EMBEDDING_MODEL_MSG)
if not texts:
return []
return [], 0
all_embeddings: list[list[float]] = []
total_tokens = 0
for i in range(0, len(texts), BATCH_SIZE):
batch = texts[i : i + BATCH_SIZE]
batch_embeddings = await self._embed_batch_request(batch)
batch_embeddings, batch_tokens = await self._embed_batch_request(batch)
all_embeddings.extend(batch_embeddings)
total_tokens += batch_tokens
if self._dimension is None and batch_embeddings:
self._dimension = len(batch_embeddings[0])
@@ -142,11 +167,18 @@ class MistralProvider(Provider):
self.embedding_model,
)
return all_embeddings
return all_embeddings, total_tokens
@_retry_429
async def _embed_batch_request(self, batch: list[str]) -> list[list[float]]:
"""Single batch request with rate-limit retry."""
async def _embed_batch_request(
self, batch: list[str]
) -> tuple[list[list[float]], int]:
"""Single batch request with rate-limit retry.
Returns ``(embeddings, token_count)``; ``token_count`` comes from the
response's ``usage.total_tokens`` and falls back to a char-based
estimate if the API omits usage.
"""
assert self.embedding_model is not None
response = await self.client.embeddings.create_async(
model=self.embedding_model,
@@ -170,7 +202,18 @@ class MistralProvider(Provider):
f"Mistral embeddings API returned {len(result)} embeddings "
f"for {len(batch)} inputs"
)
return result
usage = getattr(response, "usage", None)
total_tokens = getattr(usage, "total_tokens", None) if usage else None
# Guard on numeric type (not just ``is not None``): a real response
# gives an int, but test doubles / partial responses can surface a
# non-numeric attribute — fall back to the estimate there.
tokens = (
int(total_tokens)
if isinstance(total_tokens, (int, float))
else self._estimate_tokens(batch)
)
return result, tokens
def get_dimension(self) -> int:
if not self.supports_embeddings: