Merge pull request #875 from cbcoutinho/feat/meter-embedding-tokens

feat(usage): meter embedding tokens (tokens_embedded/pages_embedded) on both paths + Prometheus export
This commit is contained in:
Chris Coutinho
2026-06-08 15:18:36 +02:00
committed by GitHub
27 changed files with 1284 additions and 136 deletions
+80 -44
View File
@@ -56,6 +56,64 @@ logger = logging.getLogger(__name__)
_USAGE_METADATA_MAX_DOC_TYPES = 16
async def record_search_usage(
*,
enabled: bool,
user_id: str,
fusion: str,
doc_types: list[str] | None,
token_count: int | None,
) -> None:
"""Record the billable ``tokens_embedded`` event for one semantic 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 (Deck #67). ``nc_semantic_search``
and ``nc_semantic_search_answer`` (which reuses it) both flow through here —
do not add a second hook. ``nc_semantic_search_answer`` exposes no
``doc_types`` parameter, so its searches always meter with
``doc_types=None``.
Best-effort and flag-gated: a metering failure is logged and never breaks
the search. Unlike the indexing path's chunk-count guard, a 0-token query is
still recorded (the query embedding ran); a zero-value row is a no-op at the
Stripe ``sum`` aggregation.
Privacy note: ``user_id`` stays tenant-local — the CP rollup aggregates
GROUP BY (day, metric) into ``usage_daily`` (no metadata column), so nothing
here propagates to Stripe; it is retained only to keep Deck #67's future
per-user attribution derivable from app-DB metadata without a re-migration.
"""
if not enabled:
return
try:
store = await UsageEventStore.shared()
await store.record_usage_event(
metric="tokens_embedded",
value=token_count or 0,
metadata={
"user_id": user_id,
"fusion": fusion,
# Bounded copy — see _USAGE_METADATA_MAX_DOC_TYPES. Both None and
# [] normalize to null so a future metadata->'doc_types' IS NULL
# query counts the all-types case consistently.
"doc_types": (
doc_types[:_USAGE_METADATA_MAX_DOC_TYPES] if doc_types else None
),
},
# The caller already confirmed the flag, so pass enabled=True
# directly — the store then skips a second uncached Settings build on
# this hot query path (ADR-024).
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 — a silent DEBUG line would hide "operator enabled metering
# but gets no data".
logger.warning("usage metering hook (tokens_embedded) skipped", exc_info=True)
def configure_semantic_tools(mcp: FastMCP):
"""Configure semantic search tools for MCP server."""
@@ -527,51 +585,29 @@ 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.
# Usage metering (Deck #67): record the query embedding's token
# count as a billable 'tokens_embedded' event. query_token_count
# is set by BM25HybridSearchAlgorithm during the search() above; the
# doc_types loop reuses one search_algo instance for the same query
# and the algorithm caches the dense embedding per query, so the
# query is embedded — and metered — exactly once regardless of how
# many doc_types were searched. See record_search_usage for the
# metric/privacy details.
#
# Privacy note: user_id stays tenant-local. The CP rollup
# aggregates GROUP BY (day, metric) into usage_daily, which has no
# metadata column, so nothing here propagates to Stripe; the value
# is retained only so Deck #67's "per-user attribution derivable
# from app-DB metadata later" stays possible without a re-migration.
if settings.usage_metering_enabled:
try:
store = await UsageEventStore.shared()
await store.record_usage_event(
metric="embeddings_queries",
value=1,
metadata={
"user_id": username,
"fusion": fusion,
# Bounded copy — see _USAGE_METADATA_MAX_DOC_TYPES.
# Both None and [] normalize to null so a future
# metadata->'doc_types' IS NULL query counts the
# all-types case consistently.
"doc_types": (
doc_types[:_USAGE_METADATA_MAX_DOC_TYPES]
if doc_types
else None
),
},
# The outer guard already confirmed the flag, so pass
# enabled=True directly — the store then skips a second
# uncached Settings build on this hot query path
# (ADR-024).
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 — a silent DEBUG line
# would hide "operator enabled metering but gets no data".
logger.warning(
"usage metering hook (embeddings_queries) skipped",
exc_info=True,
)
# NOTE (v1 billing gap): this fires only on a fully successful
# search. If the query embed succeeded (provider billed the tokens,
# and Prometheus recorded them via record_embedding_tokens) but a
# later step (Qdrant/verify) raised, no tokens_embedded row is
# written — the embed cost is real but absent from the billing
# ledger. Acceptable for v1 (search failures are rare and the meter
# is not billed today); revisit if billing accuracy needs it.
await record_search_usage(
enabled=settings.usage_metering_enabled,
user_id=username,
fusion=fusion,
doc_types=doc_types,
token_count=search_algo.query_token_count,
)
return SemanticSearchResponse(
results=results,