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:
co-authored by
Claude Opus 4.8
parent
fe17994c4d
commit
64318f0b25
@@ -224,3 +224,15 @@ class GatewayProvider(OpenAIProvider):
|
||||
async def embed_batch(self, texts: list[str]) -> list[list[float]]:
|
||||
await self._ensure_bearer()
|
||||
return await super().embed_batch(texts)
|
||||
|
||||
async def embed_batch_with_usage(
|
||||
self, texts: list[str]
|
||||
) -> tuple[list[list[float]], int]:
|
||||
# Only the batch usage-variant is overridden: OpenAIProvider's
|
||||
# embed_with_usage() routes through embed_batch_with_usage(), so a
|
||||
# single embed_with_usage() call already lands here and refreshes the
|
||||
# bearer exactly once (overriding both would double-ensure). This
|
||||
# differs from embed()/embed_batch() above, where the single embed() is
|
||||
# self-contained and therefore needs its own override.
|
||||
await self._ensure_bearer()
|
||||
return await super().embed_batch_with_usage(texts)
|
||||
|
||||
@@ -52,6 +52,22 @@ class EmbeddingService:
|
||||
"""
|
||||
return await self.provider.embed_batch(texts)
|
||||
|
||||
async def embed_with_usage(self, text: str) -> tuple[list[float], int]:
|
||||
"""Embed one text and report the request's token count.
|
||||
|
||||
Returns ``(embedding, token_count)`` for usage metering (Deck #67).
|
||||
"""
|
||||
return await self.provider.embed_with_usage(text)
|
||||
|
||||
async def embed_batch_with_usage(
|
||||
self, texts: list[str]
|
||||
) -> tuple[list[list[float]], int]:
|
||||
"""Embed multiple texts and report the total token count.
|
||||
|
||||
Returns ``(embeddings, token_count)`` for usage metering (Deck #67).
|
||||
"""
|
||||
return await self.provider.embed_batch_with_usage(texts)
|
||||
|
||||
def get_dimension(self) -> int:
|
||||
"""
|
||||
Get embedding dimension.
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Unified provider interface for embeddings and text generation."""
|
||||
|
||||
import math
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
@@ -55,6 +56,40 @@ class Provider(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _estimate_tokens(texts: list[str]) -> int:
|
||||
"""Best-effort token estimate when a provider returns no usage data.
|
||||
|
||||
Uses a coarse ~4-chars-per-token heuristic so the billable token
|
||||
value stays non-zero and monotone with input size for local/dev
|
||||
providers (Simple, Ollama without ``prompt_eval_count``). Real
|
||||
providers override ``*_with_usage`` to report exact counts.
|
||||
"""
|
||||
return math.ceil(sum(len(t) for t in texts) / 4)
|
||||
|
||||
async def embed_with_usage(self, text: str) -> tuple[list[float], int]:
|
||||
"""Embed one text and report the request's token count.
|
||||
|
||||
Returns ``(embedding, token_count)``. The default delegates to
|
||||
:meth:`embed` and estimates the tokens; providers that surface real
|
||||
usage from their embedding response override this. Used by the
|
||||
usage-metering hooks (Deck #67) to bill ``embeddings_queries`` by
|
||||
tokens rather than by operation count.
|
||||
"""
|
||||
embedding = await self.embed(text)
|
||||
return embedding, self._estimate_tokens([text])
|
||||
|
||||
async def embed_batch_with_usage(
|
||||
self, texts: list[str]
|
||||
) -> tuple[list[list[float]], int]:
|
||||
"""Embed multiple texts and report the total token count.
|
||||
|
||||
Returns ``(embeddings, token_count)``; the default estimates. See
|
||||
:meth:`embed_with_usage`.
|
||||
"""
|
||||
embeddings = await self.embed_batch(texts)
|
||||
return embeddings, self._estimate_tokens(texts)
|
||||
|
||||
@abstractmethod
|
||||
def get_dimension(self) -> int:
|
||||
"""
|
||||
|
||||
@@ -164,6 +164,16 @@ class BedrockProvider(Provider):
|
||||
NotImplementedError: If embeddings not enabled (no embedding_model)
|
||||
ClientError: If Bedrock API call fails
|
||||
"""
|
||||
embedding, _ = await self.embed_with_usage(text)
|
||||
return embedding
|
||||
|
||||
async def embed_with_usage(self, text: str) -> tuple[list[float], int]:
|
||||
"""Embed one text, reporting the request's token count.
|
||||
|
||||
Titan Embed responses carry ``inputTextTokenCount``; for Cohere /
|
||||
unknown models (no token field) this falls back to a char-based
|
||||
estimate. Used by the usage-metering hooks (Deck #67).
|
||||
"""
|
||||
if not self.supports_embeddings:
|
||||
raise NotImplementedError(
|
||||
"Embedding not supported - no embedding_model configured"
|
||||
@@ -182,7 +192,13 @@ class BedrockProvider(Provider):
|
||||
response_body = json.loads(response["body"].read())
|
||||
embedding = self._parse_embedding_response(response_body)
|
||||
|
||||
return embedding
|
||||
token_count = response_body.get("inputTextTokenCount")
|
||||
tokens = (
|
||||
int(token_count)
|
||||
if isinstance(token_count, (int, float))
|
||||
else self._estimate_tokens([text])
|
||||
)
|
||||
return embedding, tokens
|
||||
|
||||
except (BotoCoreError, ClientError) as e:
|
||||
logger.error("Bedrock embedding error: %s", e)
|
||||
@@ -205,16 +221,30 @@ class BedrockProvider(Provider):
|
||||
NotImplementedError: If embeddings not enabled (no embedding_model)
|
||||
ClientError: If Bedrock API call fails
|
||||
"""
|
||||
embeddings, _ = await self.embed_batch_with_usage(texts)
|
||||
return embeddings
|
||||
|
||||
async def embed_batch_with_usage(
|
||||
self, texts: list[str]
|
||||
) -> tuple[list[list[float]], int]:
|
||||
"""Embed multiple texts, summing the per-call token counts.
|
||||
|
||||
Bedrock has no batch embedding API, so requests run sequentially and
|
||||
the token total is the sum of each call's ``inputTextTokenCount``
|
||||
(Titan) or estimate (Cohere/unknown).
|
||||
"""
|
||||
if not self.supports_embeddings:
|
||||
raise NotImplementedError(
|
||||
"Embedding not supported - no embedding_model configured"
|
||||
)
|
||||
|
||||
embeddings = []
|
||||
embeddings: list[list[float]] = []
|
||||
total_tokens = 0
|
||||
for text in texts:
|
||||
embedding = await self.embed(text)
|
||||
embedding, tokens = await self.embed_with_usage(text)
|
||||
embeddings.append(embedding)
|
||||
return embeddings
|
||||
total_tokens += tokens
|
||||
return embeddings, total_tokens
|
||||
|
||||
async def _detect_dimension(self):
|
||||
"""
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -116,12 +116,44 @@ class OllamaProvider(Provider):
|
||||
Raises:
|
||||
NotImplementedError: If embeddings not enabled (no embedding_model)
|
||||
"""
|
||||
embeddings, _ = await self.embed_batch_with_usage(texts, batch_size=batch_size)
|
||||
return embeddings
|
||||
|
||||
async def embed_with_usage(self, text: str) -> tuple[list[float], int]:
|
||||
"""Embed one text, reporting the request's token count.
|
||||
|
||||
Routes through ``/api/embed`` (which carries ``prompt_eval_count``)
|
||||
rather than the legacy ``/api/embeddings`` so a token count is
|
||||
available; falls back to a char-based estimate when the field is
|
||||
absent. Used by the usage-metering hooks (Deck #67).
|
||||
"""
|
||||
embeddings, tokens = await self.embed_batch_with_usage([text])
|
||||
if not embeddings:
|
||||
raise RuntimeError(
|
||||
"Ollama 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], batch_size: int = 32
|
||||
) -> tuple[list[list[float]], int]:
|
||||
"""Embed multiple texts, summing ``prompt_eval_count`` token usage.
|
||||
|
||||
Returns ``(embeddings, total_tokens)``. Ollama's ``/api/embed`` may
|
||||
omit ``prompt_eval_count`` (older versions); a char-based estimate is
|
||||
used per batch when it does.
|
||||
"""
|
||||
if not self.supports_embeddings:
|
||||
raise NotImplementedError(
|
||||
"Embedding not supported - no embedding_model configured"
|
||||
)
|
||||
|
||||
all_embeddings = []
|
||||
if not texts:
|
||||
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]
|
||||
response = await self.client.post(
|
||||
@@ -129,9 +161,17 @@ class OllamaProvider(Provider):
|
||||
json={"model": self.embedding_model, "input": batch},
|
||||
)
|
||||
response.raise_for_status()
|
||||
all_embeddings.extend(response.json()["embeddings"])
|
||||
data = response.json()
|
||||
all_embeddings.extend(data["embeddings"])
|
||||
|
||||
return all_embeddings
|
||||
prompt_eval = data.get("prompt_eval_count")
|
||||
total_tokens += (
|
||||
int(prompt_eval)
|
||||
if isinstance(prompt_eval, (int, float))
|
||||
else self._estimate_tokens(batch)
|
||||
)
|
||||
|
||||
return all_embeddings, total_tokens
|
||||
|
||||
async def _detect_dimension(self):
|
||||
"""
|
||||
|
||||
@@ -153,19 +153,49 @@ class OpenAIProvider(Provider):
|
||||
"Embedding not supported - no embedding_model configured"
|
||||
)
|
||||
|
||||
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 request's token count."""
|
||||
embeddings, tokens = await self.embed_batch_with_usage([text])
|
||||
if not embeddings:
|
||||
raise RuntimeError(
|
||||
"OpenAI 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 API-reported token usage.
|
||||
|
||||
Returns ``(embeddings, total_tokens)`` where ``total_tokens`` sums
|
||||
``response.usage.total_tokens`` across the sub-requests (the unit the
|
||||
provider bills on). Used by the usage-metering hooks (Deck #67). Also
|
||||
serves the gateway path via :class:`GatewayProvider`.
|
||||
"""
|
||||
if not self.supports_embeddings:
|
||||
raise NotImplementedError(
|
||||
"Embedding not supported - no embedding_model configured"
|
||||
)
|
||||
|
||||
if not texts:
|
||||
return []
|
||||
return [], 0
|
||||
|
||||
# OpenAI supports batches up to 2048, but use smaller batches for safety
|
||||
batch_size = 100
|
||||
all_embeddings: list[list[float]] = []
|
||||
total_tokens = 0
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
|
||||
# Use helper method with retry logic for each batch
|
||||
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
|
||||
|
||||
# Update dimension if not set
|
||||
if self._dimension is None and batch_embeddings:
|
||||
@@ -176,11 +206,18 @@ class OpenAIProvider(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]]:
|
||||
"""Make a single batch embedding request with retry logic."""
|
||||
async def _embed_batch_request(
|
||||
self, batch: list[str]
|
||||
) -> tuple[list[list[float]], int]:
|
||||
"""Make a single batch embedding request with retry logic.
|
||||
|
||||
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 # Type narrowing
|
||||
response = await self.client.embeddings.create(
|
||||
input=batch,
|
||||
@@ -188,7 +225,19 @@ class OpenAIProvider(Provider):
|
||||
)
|
||||
# Sort by index to maintain order
|
||||
sorted_data = sorted(response.data, key=lambda x: x.index)
|
||||
return [item.embedding for item in sorted_data]
|
||||
embeddings = [item.embedding for item in sorted_data]
|
||||
|
||||
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 embeddings, tokens
|
||||
|
||||
def get_dimension(self) -> int:
|
||||
"""
|
||||
|
||||
@@ -285,9 +285,15 @@ class SearchAlgorithm(ABC):
|
||||
query_embedding: The query embedding generated during the last search.
|
||||
Available after search() completes for algorithms that use embeddings.
|
||||
Can be reused by callers to avoid redundant embedding generation.
|
||||
query_token_count: Token count of the query embedding request from the
|
||||
last search (provider-reported, or estimated). Set by algorithms
|
||||
that embed the query so the usage-metering hook can bill
|
||||
``embeddings_queries`` by tokens (Deck #67). The instance is
|
||||
per-request, so this side-channel is concurrency-safe.
|
||||
"""
|
||||
|
||||
query_embedding: list[float] | None = None
|
||||
query_token_count: int | None = None
|
||||
|
||||
@abstractmethod
|
||||
async def search(
|
||||
|
||||
@@ -132,9 +132,13 @@ class BM25HybridSearchAlgorithm(SearchAlgorithm):
|
||||
with trace_operation("search.get_embedding_service"):
|
||||
embedding_service = get_embedding_service()
|
||||
with trace_operation("search.dense_embedding"):
|
||||
dense_embedding = await embedding_service.embed(query)
|
||||
# Store for reuse by callers (e.g., viz_routes PCA visualization)
|
||||
dense_embedding, query_tokens = await embedding_service.embed_with_usage(
|
||||
query
|
||||
)
|
||||
# Store for reuse by callers (e.g., viz_routes PCA visualization) and
|
||||
# for the usage-metering hook in server/semantic.py (token count).
|
||||
self.query_embedding = dense_embedding
|
||||
self.query_token_count = query_tokens
|
||||
logger.debug("Generated dense embedding (dimension=%s)", len(dense_embedding))
|
||||
|
||||
# Generate sparse embedding for BM25 keyword search
|
||||
|
||||
@@ -528,10 +528,18 @@ def configure_semantic_tools(mcp: FastMCP):
|
||||
logger.info("Returning %d results from BM25 hybrid search", len(results))
|
||||
|
||||
# Usage metering (Deck #67): one billable 'embeddings_queries'
|
||||
# event per successful search (the query embedding is the metered
|
||||
# cost). Best-effort and gated on the flag so the off-path touches
|
||||
# no storage. nc_semantic_search_answer reuses this tool, so it
|
||||
# records here too — do not add a second hook there.
|
||||
# event per successful search. The value is the query embedding's
|
||||
# token count (provider-reported, or estimated) — the unit upstream
|
||||
# providers bill on, and the same metric the indexing path records
|
||||
# for chunk embeddings. Best-effort and gated on the flag so the
|
||||
# off-path touches no storage. nc_semantic_search_answer reuses this
|
||||
# tool, so it records here too — do not add a second hook there.
|
||||
#
|
||||
# query_token_count is set by BM25HybridSearchAlgorithm during the
|
||||
# search() above. The doc_types loop reuses one search_algo instance
|
||||
# for the same query string, so the final value is the single query
|
||||
# embedding's cost (matches the prior one-query semantics). Falls
|
||||
# back to 0 only if the embedding never ran (e.g. a pre-embed error).
|
||||
#
|
||||
# Privacy note: user_id stays tenant-local. The CP rollup
|
||||
# aggregates GROUP BY (day, metric) into usage_daily, which has no
|
||||
@@ -543,7 +551,7 @@ def configure_semantic_tools(mcp: FastMCP):
|
||||
store = await UsageEventStore.shared()
|
||||
await store.record_usage_event(
|
||||
metric="embeddings_queries",
|
||||
value=1,
|
||||
value=search_algo.query_token_count or 0,
|
||||
metadata={
|
||||
"user_id": username,
|
||||
"fusion": fusion,
|
||||
|
||||
@@ -809,7 +809,10 @@ async def _index_document(
|
||||
embedding_service = get_embedding_service()
|
||||
embed_start = time.time()
|
||||
try:
|
||||
dense_embeddings = await embedding_service.embed_batch(chunk_texts)
|
||||
(
|
||||
dense_embeddings,
|
||||
embed_tokens,
|
||||
) = await embedding_service.embed_batch_with_usage(chunk_texts)
|
||||
except Exception:
|
||||
record_embedding(
|
||||
"dense", provider, time.time() - embed_start, status="error"
|
||||
@@ -833,30 +836,43 @@ async def _index_document(
|
||||
# keep Deck #67's future per-user attribution derivable from the
|
||||
# app DB without a re-migration.
|
||||
if settings.usage_metering_enabled:
|
||||
# Two billable events per indexed document: 'pages_chunks' is
|
||||
# the volume (chunks embedded); 'embeddings_queries' is the
|
||||
# token count of the embedding request — the same metric search
|
||||
# records, so the meter bills embedding tokens whether they were
|
||||
# incurred indexing a document or embedding a query (Deck #67).
|
||||
metering_metadata = {
|
||||
"provider": provider,
|
||||
"model": settings.get_embedding_model_name(),
|
||||
"doc_type": doc_task.doc_type,
|
||||
"user_id": doc_task.user_id,
|
||||
"total_chars": total_chars,
|
||||
}
|
||||
try:
|
||||
store = await UsageEventStore.shared()
|
||||
await store.record_usage_event(
|
||||
metric="pages_chunks",
|
||||
value=len(chunk_texts),
|
||||
metadata={
|
||||
"provider": provider,
|
||||
"model": settings.get_embedding_model_name(),
|
||||
"doc_type": doc_task.doc_type,
|
||||
"user_id": doc_task.user_id,
|
||||
"total_chars": total_chars,
|
||||
},
|
||||
metadata=metering_metadata,
|
||||
# The outer guard already confirmed the flag, so pass
|
||||
# enabled=True directly — the store then skips a second
|
||||
# uncached Settings build here (ADR-024).
|
||||
enabled=True,
|
||||
)
|
||||
await store.record_usage_event(
|
||||
metric="embeddings_queries",
|
||||
value=embed_tokens,
|
||||
metadata=metering_metadata,
|
||||
enabled=True,
|
||||
)
|
||||
except Exception:
|
||||
# Reached only when shared()/store construction itself
|
||||
# raises (record_usage_event swallows its own write
|
||||
# failures). Metering is on, so warn rather than hide the
|
||||
# "enabled but no billing data" case in DEBUG logs.
|
||||
logger.warning(
|
||||
"usage metering hook (pages_chunks) skipped", exc_info=True
|
||||
"usage metering hook (indexing embeddings) skipped",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
async def generate_sparse_embeddings():
|
||||
|
||||
@@ -255,6 +255,73 @@ async def test_bedrock_dimension_detection(mock_bedrock_client):
|
||||
assert provider.get_dimension() == 1536
|
||||
|
||||
|
||||
def _titan_body(embedding, token_count=None):
|
||||
payload = {"embedding": embedding}
|
||||
if token_count is not None:
|
||||
payload["inputTextTokenCount"] = token_count
|
||||
return {
|
||||
"body": MagicMock(read=MagicMock(return_value=json.dumps(payload).encode()))
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_bedrock_embed_with_usage_reports_titan_tokens(mock_bedrock_client):
|
||||
"""Titan's inputTextTokenCount is surfaced as the token count."""
|
||||
mock_bedrock_client.invoke_model.return_value = _titan_body(
|
||||
[0.1, 0.2], token_count=6
|
||||
)
|
||||
|
||||
provider = BedrockProvider(
|
||||
region_name="us-east-1",
|
||||
embedding_model="amazon.titan-embed-text-v2:0",
|
||||
generation_model=None,
|
||||
)
|
||||
embedding, tokens = await provider.embed_with_usage("test text")
|
||||
|
||||
assert embedding == [0.1, 0.2]
|
||||
assert tokens == 6
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_bedrock_embed_batch_with_usage_sums_token_counts(mock_bedrock_client):
|
||||
"""Sequential per-text calls sum their inputTextTokenCount values."""
|
||||
mock_bedrock_client.invoke_model.return_value = _titan_body(
|
||||
[0.1, 0.2], token_count=4
|
||||
)
|
||||
|
||||
provider = BedrockProvider(
|
||||
region_name="us-east-1",
|
||||
embedding_model="amazon.titan-embed-text-v2:0",
|
||||
generation_model=None,
|
||||
)
|
||||
embeddings, tokens = await provider.embed_batch_with_usage(["t1", "t2", "t3"])
|
||||
|
||||
assert len(embeddings) == 3
|
||||
assert tokens == 12 # 4 tokens per call × 3 calls
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_bedrock_with_usage_estimates_when_token_count_absent(
|
||||
mock_bedrock_client,
|
||||
):
|
||||
"""Cohere returns no inputTextTokenCount → char-based estimate."""
|
||||
mock_bedrock_client.invoke_model.return_value = {
|
||||
"body": MagicMock(
|
||||
read=MagicMock(
|
||||
return_value=json.dumps({"embeddings": [[0.1, 0.2]]}).encode()
|
||||
)
|
||||
)
|
||||
}
|
||||
|
||||
provider = BedrockProvider(
|
||||
region_name="us-east-1",
|
||||
embedding_model="cohere.embed-english-v3",
|
||||
)
|
||||
_, tokens = await provider.embed_with_usage("abcdefgh") # 8 chars → 2 tokens
|
||||
|
||||
assert tokens == 2
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_bedrock_cohere_embedding(mock_bedrock_client):
|
||||
"""Test Bedrock with Cohere embedding model."""
|
||||
|
||||
@@ -6,6 +6,7 @@ tenant realm); creds are all-or-nothing.
|
||||
"""
|
||||
|
||||
import time
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
@@ -327,6 +328,69 @@ def test_trailing_slash_base_url_normalized():
|
||||
assert not base.endswith("/v1/v1")
|
||||
|
||||
|
||||
async def test_gateway_embed_with_usage_forwards_after_bearer(monkeypatch):
|
||||
"""embed_with_usage refreshes the bearer, then returns the (embedding,
|
||||
token_count) from the inherited OpenAI implementation."""
|
||||
provider = GatewayProvider(
|
||||
base_url="http://gw:8083/v1", embedding_model="mistral/mistral-embed"
|
||||
)
|
||||
|
||||
order: list[str] = []
|
||||
|
||||
async def _ensure_bearer():
|
||||
order.append("bearer")
|
||||
|
||||
monkeypatch.setattr(provider, "_ensure_bearer", _ensure_bearer)
|
||||
|
||||
item = MagicMock()
|
||||
item.embedding = [0.1, 0.2]
|
||||
item.index = 0
|
||||
response = MagicMock()
|
||||
response.data = [item]
|
||||
response.usage = MagicMock(total_tokens=8)
|
||||
|
||||
async def _create(**_kwargs):
|
||||
order.append("embed")
|
||||
return response
|
||||
|
||||
monkeypatch.setattr(provider.client.embeddings, "create", _create)
|
||||
|
||||
embedding, tokens = await provider.embed_with_usage("hello")
|
||||
|
||||
assert embedding == [0.1, 0.2]
|
||||
assert tokens == 8
|
||||
assert order == ["bearer", "embed"] # bearer refreshed before the embed call
|
||||
|
||||
|
||||
async def test_gateway_embed_batch_with_usage_forwards_after_bearer(monkeypatch):
|
||||
"""embed_batch_with_usage also refreshes the bearer before delegating."""
|
||||
provider = GatewayProvider(
|
||||
base_url="http://gw:8083/v1", embedding_model="mistral/mistral-embed"
|
||||
)
|
||||
ensured = {"n": 0}
|
||||
|
||||
async def _ensure_bearer():
|
||||
ensured["n"] += 1
|
||||
|
||||
monkeypatch.setattr(provider, "_ensure_bearer", _ensure_bearer)
|
||||
|
||||
item = MagicMock()
|
||||
item.embedding = [0.3, 0.4]
|
||||
item.index = 0
|
||||
response = MagicMock()
|
||||
response.data = [item]
|
||||
response.usage = MagicMock(total_tokens=5)
|
||||
monkeypatch.setattr(
|
||||
provider.client.embeddings, "create", AsyncMock(return_value=response)
|
||||
)
|
||||
|
||||
embeddings, tokens = await provider.embed_batch_with_usage(["x"])
|
||||
|
||||
assert embeddings == [[0.3, 0.4]]
|
||||
assert tokens == 5
|
||||
assert ensured["n"] == 1
|
||||
|
||||
|
||||
async def test_detect_dimension_with_bare_base_url_hits_v1_models(monkeypatch):
|
||||
"""End-to-end of the fix: a bare-origin base_url still resolves the
|
||||
dimension because discovery lands on /v1/models."""
|
||||
|
||||
@@ -255,6 +255,68 @@ async def test_mistral_batch_raises_on_count_mismatch(mock_mistral_client):
|
||||
await provider.embed_batch(["a", "b"])
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_embed_batch_with_usage_reports_tokens(mock_mistral_client):
|
||||
"""embed_batch_with_usage returns the provider-reported total_tokens."""
|
||||
response = _make_response([[0.1, 0.2], [0.3, 0.4]])
|
||||
response.usage = MagicMock(total_tokens=11)
|
||||
mock_mistral_client.embeddings.create_async = AsyncMock(return_value=response)
|
||||
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
embeddings, tokens = await provider.embed_batch_with_usage(["a", "b"])
|
||||
|
||||
assert embeddings == [[0.1, 0.2], [0.3, 0.4]]
|
||||
assert tokens == 11
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_embed_batch_with_usage_sums_across_chunks(mock_mistral_client):
|
||||
"""Token counts sum across the BATCH_SIZE sub-requests (1 token/input here)."""
|
||||
|
||||
def _side_effect(*, model, inputs, **_kwargs):
|
||||
resp = _make_response([[float(i)] for i in range(len(inputs))])
|
||||
resp.usage = MagicMock(total_tokens=len(inputs))
|
||||
return resp
|
||||
|
||||
mock_mistral_client.embeddings.create_async = AsyncMock(side_effect=_side_effect)
|
||||
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
total = BATCH_SIZE * 2 + 5 # three chunks
|
||||
embeddings, tokens = await provider.embed_batch_with_usage(
|
||||
[f"t-{i}" for i in range(total)]
|
||||
)
|
||||
|
||||
assert len(embeddings) == total
|
||||
assert tokens == total # summed across all three chunks
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_with_usage_estimates_when_usage_absent(mock_mistral_client):
|
||||
"""Missing usage falls back to the char-based estimate, not a crash."""
|
||||
response = _make_response([[0.1, 0.2]])
|
||||
response.usage = None
|
||||
mock_mistral_client.embeddings.create_async = AsyncMock(return_value=response)
|
||||
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
_, tokens = await provider.embed_batch_with_usage(["abcd"]) # 4 chars → 1 token
|
||||
|
||||
assert tokens == 1
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_mistral_embed_with_usage_single(mock_mistral_client):
|
||||
"""embed_with_usage returns the single embedding plus its token count."""
|
||||
response = _make_response([[0.5, 0.6]])
|
||||
response.usage = MagicMock(total_tokens=3)
|
||||
mock_mistral_client.embeddings.create_async = AsyncMock(return_value=response)
|
||||
|
||||
provider = MistralProvider(api_key="test-key", embedding_model="mistral-embed")
|
||||
embedding, tokens = await provider.embed_with_usage("hello")
|
||||
|
||||
assert embedding == [0.5, 0.6]
|
||||
assert tokens == 3
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_mistral_is_rate_limit_predicate():
|
||||
"""_is_rate_limit returns True only for SDKErrors with status_code == 429."""
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Unit tests for Ollama provider token-usage surfacing.
|
||||
|
||||
The provider has no other unit coverage; these focus on the ``*_with_usage``
|
||||
methods added for usage metering (Deck #67) — provider-reported
|
||||
``prompt_eval_count`` and the char-based estimate fallback when it's absent.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from nextcloud_mcp_server.providers.ollama import OllamaProvider
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ollama_provider():
|
||||
# Construct with no models so __init__ skips _check_model_is_loaded (no
|
||||
# network call), then enable embeddings post-construction.
|
||||
provider = OllamaProvider(base_url="http://ollama:11434")
|
||||
provider.embedding_model = "nomic-embed-text"
|
||||
return provider
|
||||
|
||||
|
||||
def _embed_response(embeddings, prompt_eval_count=None):
|
||||
payload = {"embeddings": embeddings}
|
||||
if prompt_eval_count is not None:
|
||||
payload["prompt_eval_count"] = prompt_eval_count
|
||||
resp = MagicMock()
|
||||
resp.json = MagicMock(return_value=payload)
|
||||
resp.raise_for_status = MagicMock()
|
||||
return resp
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_ollama_embed_batch_with_usage_reports_prompt_eval_count(ollama_provider):
|
||||
"""prompt_eval_count from /api/embed is surfaced as the token count."""
|
||||
ollama_provider.client.post = AsyncMock(
|
||||
return_value=_embed_response([[0.1, 0.2], [0.3, 0.4]], prompt_eval_count=7)
|
||||
)
|
||||
|
||||
embeddings, tokens = await ollama_provider.embed_batch_with_usage(["a", "b"])
|
||||
|
||||
assert embeddings == [[0.1, 0.2], [0.3, 0.4]]
|
||||
assert tokens == 7
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_ollama_with_usage_estimates_when_count_absent(ollama_provider):
|
||||
"""Older Ollama omits prompt_eval_count → char-based estimate."""
|
||||
ollama_provider.client.post = AsyncMock(
|
||||
return_value=_embed_response([[0.1]], prompt_eval_count=None)
|
||||
)
|
||||
|
||||
_, tokens = await ollama_provider.embed_with_usage("abcdefgh") # 8 chars → 2
|
||||
|
||||
assert tokens == 2
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_ollama_empty_batch_with_usage(ollama_provider):
|
||||
"""Empty batch returns no embeddings, zero tokens, and makes no request."""
|
||||
ollama_provider.client.post = AsyncMock()
|
||||
|
||||
embeddings, tokens = await ollama_provider.embed_batch_with_usage([])
|
||||
|
||||
assert embeddings == []
|
||||
assert tokens == 0
|
||||
ollama_provider.client.post.assert_not_called()
|
||||
@@ -280,6 +280,46 @@ async def test_openai_empty_batch():
|
||||
assert embeddings == []
|
||||
|
||||
|
||||
def _embed_item(embedding, index):
|
||||
item = MagicMock()
|
||||
item.embedding = embedding
|
||||
item.index = index
|
||||
return item
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_openai_embed_batch_with_usage_reports_tokens(mock_openai_client):
|
||||
"""embed_batch_with_usage returns the response's total_tokens."""
|
||||
response = MagicMock()
|
||||
response.data = [_embed_item([0.1, 0.2], 0), _embed_item([0.3, 0.4], 1)]
|
||||
response.usage = MagicMock(total_tokens=9)
|
||||
mock_openai_client.embeddings.create = AsyncMock(return_value=response)
|
||||
|
||||
provider = OpenAIProvider(
|
||||
api_key="test-key", embedding_model="text-embedding-3-small"
|
||||
)
|
||||
embeddings, tokens = await provider.embed_batch_with_usage(["a", "b"])
|
||||
|
||||
assert embeddings == [[0.1, 0.2], [0.3, 0.4]]
|
||||
assert tokens == 9
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_openai_with_usage_estimates_when_usage_absent(mock_openai_client):
|
||||
"""Missing usage falls back to the char-based estimate."""
|
||||
response = MagicMock()
|
||||
response.data = [_embed_item([0.1], 0)]
|
||||
response.usage = None
|
||||
mock_openai_client.embeddings.create = AsyncMock(return_value=response)
|
||||
|
||||
provider = OpenAIProvider(
|
||||
api_key="test-key", embedding_model="text-embedding-3-small"
|
||||
)
|
||||
_, tokens = await provider.embed_with_usage("abcdefgh") # 8 chars → 2 tokens
|
||||
|
||||
assert tokens == 2
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_openai_close(mock_openai_client):
|
||||
"""Test OpenAI client close."""
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
"""Token-usage surfacing: the Provider ABC estimate default + SimpleProvider.
|
||||
|
||||
The usage-metering hooks (Deck #67) bill ``embeddings_queries`` by tokens. Real
|
||||
providers report exact counts from their API response; providers without a token
|
||||
field (Simple, and the ABC default) fall back to a char-based estimate so the
|
||||
billable value stays non-zero and monotone with input size.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from nextcloud_mcp_server.providers.base import Provider
|
||||
from nextcloud_mcp_server.providers.simple import SimpleProvider
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_estimate_tokens_is_char_based():
|
||||
"""~4-chars-per-token, ceil-rounded, summed across inputs."""
|
||||
assert Provider._estimate_tokens(["abcd"]) == 1 # 4 chars
|
||||
assert Provider._estimate_tokens(["abcde"]) == 2 # 5 chars → ceil(5/4)
|
||||
assert Provider._estimate_tokens(["ab", "cd"]) == 1 # 4 chars total
|
||||
assert Provider._estimate_tokens([]) == 0
|
||||
assert Provider._estimate_tokens([""]) == 0
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_simple_provider_embed_with_usage_estimates():
|
||||
"""SimpleProvider has no real usage → estimate path via the ABC default."""
|
||||
provider = SimpleProvider(dimension=8)
|
||||
embedding, tokens = await provider.embed_with_usage("abcdefgh") # 8 chars → 2
|
||||
|
||||
assert len(embedding) == 8
|
||||
assert tokens == 2
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_simple_provider_embed_batch_with_usage_estimates():
|
||||
"""Batch estimate sums character counts across all inputs."""
|
||||
provider = SimpleProvider(dimension=8)
|
||||
embeddings, tokens = await provider.embed_batch_with_usage(["abcd", "efgh"])
|
||||
|
||||
assert len(embeddings) == 2
|
||||
assert tokens == 2 # 8 chars total → 2 tokens
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
async def test_simple_provider_empty_batch_with_usage():
|
||||
"""Empty batch returns no embeddings and zero tokens."""
|
||||
provider = SimpleProvider(dimension=8)
|
||||
embeddings, tokens = await provider.embed_batch_with_usage([])
|
||||
|
||||
assert embeddings == []
|
||||
assert tokens == 0
|
||||
Reference in New Issue
Block a user