diff --git a/nextcloud_mcp_server/alembic/versions/20260607_1200_007_add_usage_events.py b/nextcloud_mcp_server/alembic/versions/20260607_1200_007_add_usage_events.py index 6ce5bf32..e1e7712e 100644 --- a/nextcloud_mcp_server/alembic/versions/20260607_1200_007_add_usage_events.py +++ b/nextcloud_mcp_server/alembic/versions/20260607_1200_007_add_usage_events.py @@ -55,7 +55,7 @@ def upgrade() -> None: postgresql.TIMESTAMP(timezone=True) if is_pg else sa.TIMESTAMP(), nullable=False, ), - # Catalog metric: 'embeddings_queries' or 'pages_chunks'. Deliberately + # Catalog metric: 'tokens_embedded' or 'pages_embedded'. Deliberately # an unconstrained Text (no CHECK/enum) — the metric catalog lives in # control-plane config, not the app-DB schema. If a third metric is # ever added, the CP-side catalog must learn it too, or its rollup will diff --git a/nextcloud_mcp_server/embedding/gateway_client.py b/nextcloud_mcp_server/embedding/gateway_client.py index 82f37257..f372c7f8 100644 --- a/nextcloud_mcp_server/embedding/gateway_client.py +++ b/nextcloud_mcp_server/embedding/gateway_client.py @@ -217,10 +217,21 @@ class GatewayProvider(OpenAIProvider): exc, ) + # Bearer-refresh override topology. OpenAIProvider routes embed_batch(), + # embed_with_usage() and embed_batch_with_usage() all through + # embed_batch_with_usage(); only embed() (single) is self-contained. So we + # override exactly two methods to refresh the bearer exactly once on every + # path: embed() (its own entrypoint) and embed_batch_with_usage() (the + # shared funnel for the other three). Overriding embed_batch() as well would + # double-call _ensure_bearer() (override → super().embed_batch() → + # self.embed_batch_with_usage() → override again). + async def embed(self, text: str) -> list[float]: await self._ensure_bearer() return await super().embed(text) - async def embed_batch(self, texts: list[str]) -> list[list[float]]: + async def embed_batch_with_usage( + self, texts: list[str] + ) -> tuple[list[list[float]], int]: await self._ensure_bearer() - return await super().embed_batch(texts) + return await super().embed_batch_with_usage(texts) diff --git a/nextcloud_mcp_server/embedding/service.py b/nextcloud_mcp_server/embedding/service.py index 32e53c9a..a33abf88 100644 --- a/nextcloud_mcp_server/embedding/service.py +++ b/nextcloud_mcp_server/embedding/service.py @@ -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. diff --git a/nextcloud_mcp_server/observability/metrics.py b/nextcloud_mcp_server/observability/metrics.py index 921617fa..8e18adb8 100644 --- a/nextcloud_mcp_server/observability/metrics.py +++ b/nextcloud_mcp_server/observability/metrics.py @@ -338,6 +338,22 @@ embedding_chars_total = Counter( ["kind", "provider"], ) +# Token consumption — the billed cost unit (mirrors the tokens_embedded billing +# measure, Deck #67). On a dedicated counter (not folded into the chunk/request +# metrics above) so query embeds don't inflate indexing dashboards; labelled by +# operation = index | query. Always emitted, independent of USAGE_METERING_ENABLED. +# +# Dashboard note: operation="query" is recorded at embed time (before Qdrant / +# verify-on-read), whereas the billing-store tokens_embedded row is written only +# after the search fully succeeds. So this counter can legitimately exceed the +# billing aggregate when a search fails post-embed — don't alert on that gap as +# a divergence bug. +embedding_tokens_total = Counter( + "astrolabe_embedding_tokens_total", + "Total embedding tokens consumed (provider-reported or estimated)", + ["provider", "operation"], # operation: index | query +) + # --- Chunking & indexed-by-type ----------------------------------------------- document_chunks_total = Counter( @@ -727,6 +743,25 @@ def record_embedding( embedding_chars_total.labels(kind=kind, provider=provider).inc(chars) +def record_embedding_tokens(provider: str, operation: str, tokens: int) -> None: + """Export embedding token consumption to Prometheus. + + Mirrors the ``tokens_embedded`` billing measure (Deck #67) as an always-on + observability signal — emitted regardless of ``USAGE_METERING_ENABLED`` so + OSS/self-host deployments still see token cost in Grafana. + + Args: + provider: Provider family (mistral | openai | bedrock | ollama | simple). + operation: ``"index"`` (chunk-batch embedding) or ``"query"`` (search + query embedding). + tokens: Token count for this embedding request (no-op when ``<= 0``). + """ + if tokens > 0: + embedding_tokens_total.labels(provider=provider, operation=operation).inc( + tokens + ) + + def record_document_chunks(doc_type: str, count: int) -> None: """ Record the number of chunks produced for a document. diff --git a/nextcloud_mcp_server/providers/base.py b/nextcloud_mcp_server/providers/base.py index 3ea95827..bda908ae 100644 --- a/nextcloud_mcp_server/providers/base.py +++ b/nextcloud_mcp_server/providers/base.py @@ -1,5 +1,6 @@ """Unified provider interface for embeddings and text generation.""" +import math from abc import ABC, abstractmethod @@ -55,6 +56,51 @@ 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 ``tokens_embedded`` by + tokens rather than by operation count. + + IMPORTANT (recursion invariant): this default calls ``self.embed``. A + provider that overrides ``embed()`` to delegate to ``embed_with_usage()`` + (to avoid duplicating request logic) MUST also override this method, or + the two will call each other forever. The shipped providers that use + that delegation (Bedrock) do override both — keep that pairing. + """ + 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`. + + IMPORTANT (recursion invariant): this default calls ``self.embed_batch``. + A provider that overrides ``embed_batch()`` to delegate to + ``embed_batch_with_usage()`` (Mistral, OpenAI, Ollama do) MUST also + override this method, or the two recurse infinitely. Keep the pairing. + """ + embeddings = await self.embed_batch(texts) + return embeddings, self._estimate_tokens(texts) + @abstractmethod def get_dimension(self) -> int: """ diff --git a/nextcloud_mcp_server/providers/bedrock.py b/nextcloud_mcp_server/providers/bedrock.py index a1b9c179..d8afa40f 100644 --- a/nextcloud_mcp_server/providers/bedrock.py +++ b/nextcloud_mcp_server/providers/bedrock.py @@ -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 = ( + round(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): """ diff --git a/nextcloud_mcp_server/providers/mistral.py b/nextcloud_mcp_server/providers/mistral.py index b5f61584..bacdadf0 100644 --- a/nextcloud_mcp_server/providers/mistral.py +++ b/nextcloud_mcp_server/providers/mistral.py @@ -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 ``tokens_embedded`` 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 = ( + round(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: diff --git a/nextcloud_mcp_server/providers/ollama.py b/nextcloud_mcp_server/providers/ollama.py index 1d5705c4..ebd3b6bf 100644 --- a/nextcloud_mcp_server/providers/ollama.py +++ b/nextcloud_mcp_server/providers/ollama.py @@ -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 @@ -116,12 +111,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 +156,29 @@ 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 + # Cache the dimension inline (mirrors OpenAI/Mistral) so it is set + # via any embed path, not only an explicit _detect_dimension() call. + if self._dimension is None and data["embeddings"]: + self._dimension = len(data["embeddings"][0]) + + # ``prompt_eval_count`` is assumed to be the batch-level total for a + # multi-input /api/embed call. Ollama's API docs aren't explicit + # about batch aggregation; if a version reports only the last + # input's tokens this understates the batch. Unverified against a + # live instance — Ollama isn't the Cloud billing provider (Mistral + # is). If it proves last-item-only, switch to per-item requests and + # sum. The char-based estimate covers versions that omit the field. + prompt_eval = data.get("prompt_eval_count") + total_tokens += ( + round(prompt_eval) + if isinstance(prompt_eval, (int, float)) + else self._estimate_tokens(batch) + ) + + return all_embeddings, total_tokens async def _detect_dimension(self): """ diff --git a/nextcloud_mcp_server/providers/openai.py b/nextcloud_mcp_server/providers/openai.py index 536697bd..dfa7a8a1 100644 --- a/nextcloud_mcp_server/providers/openai.py +++ b/nextcloud_mcp_server/providers/openai.py @@ -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 = ( + round(total_tokens) + if isinstance(total_tokens, (int, float)) + else self._estimate_tokens(batch) + ) + return embeddings, tokens def get_dimension(self) -> int: """ diff --git a/nextcloud_mcp_server/search/algorithms.py b/nextcloud_mcp_server/search/algorithms.py index 7833b3e8..b4990b0c 100644 --- a/nextcloud_mcp_server/search/algorithms.py +++ b/nextcloud_mcp_server/search/algorithms.py @@ -285,9 +285,25 @@ 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 + ``tokens_embedded`` by tokens (Deck #67). The instance is + per-request, so this side-channel is concurrency-safe. """ + # Class-level defaults are a safety net; __init__ shadows them per instance. query_embedding: list[float] | None = None + query_token_count: int | None = None + + def __init__(self) -> None: + # Set the query-embedding side-channel as instance attributes so + # concurrent SearchAlgorithm instances never share it through the + # class-level defaults above — per-request isolation by construction, + # not just by the convention that each subclass redeclares them. + # Subclasses with their own __init__ should call super().__init__(). + self.query_embedding: list[float] | None = None + self.query_token_count: int | None = None @abstractmethod async def search( diff --git a/nextcloud_mcp_server/search/bm25_hybrid.py b/nextcloud_mcp_server/search/bm25_hybrid.py index fc1ce41e..9b19cb40 100644 --- a/nextcloud_mcp_server/search/bm25_hybrid.py +++ b/nextcloud_mcp_server/search/bm25_hybrid.py @@ -9,7 +9,10 @@ from qdrant_client.models import Filter from nextcloud_mcp_server.config import get_settings from nextcloud_mcp_server.embedding import get_bm25_service, get_embedding_service -from nextcloud_mcp_server.observability.metrics import record_qdrant_operation +from nextcloud_mcp_server.observability.metrics import ( + record_embedding_tokens, + record_qdrant_operation, +) from nextcloud_mcp_server.observability.tracing import trace_operation from nextcloud_mcp_server.search.access_filter import build_base_filter_conditions from nextcloud_mcp_server.search.algorithms import ( @@ -53,9 +56,16 @@ class BM25HybridSearchAlgorithm(SearchAlgorithm): f"Invalid fusion algorithm '{fusion}'. Must be 'rrf' or 'dbsf'" ) + # super() sets the per-instance query_embedding / query_token_count + # side-channel; this adds the cache key for it. + super().__init__() self.score_threshold = score_threshold self.fusion = models.Fusion.RRF if fusion == "rrf" else models.Fusion.DBSF self.fusion_name = fusion + # ``_embedded_query`` is the query string whose dense embedding is held + # in ``query_embedding`` — repeated search() calls on this per-request + # instance (the doc_types loop) reuse it instead of re-embedding. + self._embedded_query: str | None = None @property def name(self) -> str: @@ -128,13 +138,35 @@ class BM25HybridSearchAlgorithm(SearchAlgorithm): self.fusion_name, ) - # Generate dense embedding for semantic search + # Generate dense embedding for semantic search. Cache it per query on + # this (per-request) instance: nc_semantic_search calls search() once + # per doc_type with the same query, so re-embedding each time would make + # N redundant API calls and bill the query's tokens N times (Deck #67). + # Reuse the first call's embedding + token count so the query is embedded + # — and metered — exactly once. 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) - self.query_embedding = dense_embedding + if self.query_embedding is not None and self._embedded_query == query: + dense_embedding = self.query_embedding + else: + ( + 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 + self._embedded_query = query + # Export query-embedding token cost to Prometheus + # (operation=query), mirroring the per-search billing record in + # server/semantic.py. Inside the cache-miss branch so a reused + # embedding isn't double-counted. + record_embedding_tokens( + settings.get_embedding_provider_family(), "query", query_tokens + ) logger.debug("Generated dense embedding (dimension=%s)", len(dense_embedding)) # Generate sparse embedding for BM25 keyword search diff --git a/nextcloud_mcp_server/search/semantic.py b/nextcloud_mcp_server/search/semantic.py index ff516688..c81695ef 100644 --- a/nextcloud_mcp_server/search/semantic.py +++ b/nextcloud_mcp_server/search/semantic.py @@ -35,6 +35,7 @@ class SemanticSearchAlgorithm(SearchAlgorithm): Args: score_threshold: Minimum similarity score (0-1, default: 0.7) """ + super().__init__() self.score_threshold = score_threshold @property diff --git a/nextcloud_mcp_server/server/semantic.py b/nextcloud_mcp_server/server/semantic.py index 8d78849f..382236a4 100644 --- a/nextcloud_mcp_server/server/semantic.py +++ b/nextcloud_mcp_server/server/semantic.py @@ -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, diff --git a/nextcloud_mcp_server/usage/store.py b/nextcloud_mcp_server/usage/store.py index 092b0c40..1cf4b52a 100644 --- a/nextcloud_mcp_server/usage/store.py +++ b/nextcloud_mcp_server/usage/store.py @@ -102,8 +102,8 @@ class UsageEventStore: logged and swallowed — this must never break the caller's operation. Args: - metric: Catalog metric, e.g. ``"embeddings_queries"`` or - ``"pages_chunks"``. + metric: Catalog metric, e.g. ``"tokens_embedded"`` or + ``"pages_embedded"``. value: Count/quantity for this event. occurred_at: Operation completion time; defaults to now (UTC). metadata: Optional rawest-unit context (provider, model, tokens, diff --git a/nextcloud_mcp_server/vector/processor.py b/nextcloud_mcp_server/vector/processor.py index b9bd2472..3dffc381 100644 --- a/nextcloud_mcp_server/vector/processor.py +++ b/nextcloud_mcp_server/vector/processor.py @@ -23,6 +23,7 @@ from nextcloud_mcp_server.observability.metrics import ( record_document_chunks, record_document_parse_failed, record_embedding, + record_embedding_tokens, record_qdrant_operation, record_vector_sync_processing, update_vector_sync_queue_size, @@ -112,6 +113,80 @@ def should_use_page_aware( return page_aware_enabled and doc_type == "file" and bool(page_boundaries) +async def record_indexing_usage( + *, + enabled: bool, + provider: str, + model: str, + doc_type: str, + user_id: str, + chunk_count: int, + token_count: int, + total_chars: int, +) -> None: + """Record the two billable usage events for one embedded document. + + ``pages_embedded`` is the buyer-facing "pages indexed" dimension; + ``tokens_embedded`` is the embedding request's token count — the same metric + search records, so the meter bills embedding tokens whether they were + incurred indexing a document or embedding a query (Deck #67). + + TODO(#282): ``pages_embedded`` currently carries the raw chunk count + (``len(chunk_texts)``) as an interim value. The real normalized "pages + indexed" count — real pages for paginated types (PDF/DOCX/PPT), a fixed + chars/tokens-per-page constant otherwise — is deferred to instrumentation + card #282; this code (card #284) only lands the metric name/contract. + + Best-effort and flag-gated: a metering failure is logged and never breaks + indexing. No-op when metering is disabled or the document produced no chunks + (an empty batch embeds nothing and would only write zero-value rows). + + Privacy note: ``user_id`` stays tenant-local — the CP rollup aggregates + GROUP BY (day, metric) into ``usage_daily`` (no metadata column), so nothing + here reaches Stripe; it is retained only to keep Deck #67's future per-user + attribution derivable from the app DB without a re-migration. + """ + if not enabled or chunk_count == 0: + return + + metadata = { + "provider": provider, + "model": model, + "doc_type": doc_type, + "user_id": user_id, + "total_chars": total_chars, + } + try: + store = await UsageEventStore.shared() + # enabled=True: the guard above already confirmed the flag, so the store + # skips a second uncached Settings build per record (ADR-024). + # record_usage_event swallows its own write failures, so the two records + # are independent; if pages_embedded somehow raised mid-way, + # tokens_embedded would be skipped, leaving an unmatched pages_embedded + # row — acceptable under the (day, metric) SUM-aggregation billing model. + await store.record_usage_event( + # TODO(#282): value is the interim chunk count; switch to normalized + # real-page count when the per-page constant lands. + metric="pages_embedded", + value=chunk_count, + metadata=metadata, + enabled=True, + ) + await store.record_usage_event( + metric="tokens_embedded", + value=token_count, + metadata=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. + logger.warning( + "usage metering hook (indexing embeddings) skipped", exc_info=True + ) + + async def processor_task( worker_id: int, receive_stream: MemoryObjectReceiveStream[DocumentTask], @@ -598,9 +673,13 @@ async def _index_document( doc_type="file", user_id=doc_task.user_id, ) + # No embedding ran, so no usage is recorded here — stated + # explicitly so a "fewer tokens_embedded rows than expected" + # audit lands on the dedup path rather than reconstructing it + # from Qdrant claim logs. logger.info( "Dedup hit for file %s (etag=%s); claimed for user %s " - "without reprocessing", + "without reprocessing (no embedding/usage recorded)", doc_task.doc_id, doc_task.etag, doc_task.user_id, @@ -809,7 +888,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" @@ -822,42 +904,24 @@ async def _index_document( chunks=len(chunk_texts), chars=total_chars, ) - # Usage metering (Deck #67): record chunks embedded as a billable - # 'pages_chunks' event. Best-effort and gated on the flag so the - # off-path (OSS default) touches no storage; placed after the - # embedding succeeds so it can never affect the indexing path. - # - # Privacy note: user_id stays tenant-local — the CP rollup - # aggregates GROUP BY (day, metric) into usage_daily (no metadata - # column), so nothing here reaches Stripe; it is retained only to - # keep Deck #67's future per-user attribution derivable from the - # app DB without a re-migration. - if settings.usage_metering_enabled: - 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, - }, - # 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, - ) - 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 - ) + # Export token consumption to Prometheus (always-on, independent of + # the billing flag) so Grafana sees indexing token cost. + record_embedding_tokens(provider, "index", embed_tokens) + # Usage metering (Deck #67): record the chunk volume + + # embedding-token count for this document. Best-effort and + # flag-gated; placed after the embedding succeeds so it can never + # affect the indexing path. See record_indexing_usage for the + # metric/privacy details. + await record_indexing_usage( + enabled=settings.usage_metering_enabled, + provider=provider, + model=settings.get_embedding_model_name(), + doc_type=doc_task.doc_type, + user_id=doc_task.user_id, + chunk_count=len(chunk_texts), + token_count=embed_tokens, + total_chars=total_chars, + ) async def generate_sparse_embeddings(): """Generate sparse embeddings (BM25 for keyword matching).""" diff --git a/tests/unit/providers/test_bedrock.py b/tests/unit/providers/test_bedrock.py index e8ac0df6..0e97e589 100644 --- a/tests/unit/providers/test_bedrock.py +++ b/tests/unit/providers/test_bedrock.py @@ -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.""" diff --git a/tests/unit/providers/test_gateway_provider.py b/tests/unit/providers/test_gateway_provider.py index 1e941b8e..45f798a8 100644 --- a/tests/unit/providers/test_gateway_provider.py +++ b/tests/unit/providers/test_gateway_provider.py @@ -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,105 @@ def test_trailing_slash_base_url_normalized(): assert not base.endswith("/v1/v1") +@pytest.mark.unit +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.""" + # https mock host (never contacted — the OpenAI client is patched below). + provider = GatewayProvider( + base_url="https://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 + + +@pytest.mark.unit +async def test_gateway_embed_batch_with_usage_forwards_after_bearer(monkeypatch): + """embed_batch_with_usage also refreshes the bearer before delegating.""" + # https mock host (never contacted — the OpenAI client is patched below). + provider = GatewayProvider( + base_url="https://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 + + +@pytest.mark.unit +async def test_gateway_embed_batch_ensures_bearer_once(monkeypatch): + """embed_batch() has no override: it routes through the inherited OpenAI + embed_batch() → embed_batch_with_usage() (overridden), so the bearer is + refreshed exactly once — not twice.""" + # https mock host (never contacted — the OpenAI client is patched below). + provider = GatewayProvider( + base_url="https://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.1, 0.2] + item.index = 0 + response = MagicMock() + response.data = [item] + response.usage = MagicMock(total_tokens=4) + monkeypatch.setattr( + provider.client.embeddings, "create", AsyncMock(return_value=response) + ) + + embeddings = await provider.embed_batch(["x"]) + + assert embeddings == [[0.1, 0.2]] + assert ensured["n"] == 1 # not 2 — embed_batch() must not double-refresh + + 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.""" diff --git a/tests/unit/providers/test_mistral.py b/tests/unit/providers/test_mistral.py index 0e38da15..20bbce9a 100644 --- a/tests/unit/providers/test_mistral.py +++ b/tests/unit/providers/test_mistral.py @@ -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.""" diff --git a/tests/unit/providers/test_ollama.py b/tests/unit/providers/test_ollama.py new file mode 100644 index 00000000..94505dde --- /dev/null +++ b/tests/unit/providers/test_ollama.py @@ -0,0 +1,69 @@ +"""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. https mock host + # (never contacted — client.post is patched in each test). + provider = OllamaProvider(base_url="https://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() diff --git a/tests/unit/providers/test_openai.py b/tests/unit/providers/test_openai.py index bd0f2ad2..fa70fa17 100644 --- a/tests/unit/providers/test_openai.py +++ b/tests/unit/providers/test_openai.py @@ -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.""" diff --git a/tests/unit/providers/test_provider_usage.py b/tests/unit/providers/test_provider_usage.py new file mode 100644 index 00000000..54772cf6 --- /dev/null +++ b/tests/unit/providers/test_provider_usage.py @@ -0,0 +1,52 @@ +"""Token-usage surfacing: the Provider ABC estimate default + SimpleProvider. + +The usage-metering hooks (Deck #67) bill ``tokens_embedded`` 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 diff --git a/tests/unit/search/test_bm25_hybrid.py b/tests/unit/search/test_bm25_hybrid.py index a80b57bc..835adc41 100644 --- a/tests/unit/search/test_bm25_hybrid.py +++ b/tests/unit/search/test_bm25_hybrid.py @@ -1,5 +1,7 @@ """Unit tests for BM25 hybrid search algorithm.""" +from unittest.mock import AsyncMock, MagicMock + import pytest from qdrant_client import models @@ -52,3 +54,73 @@ def test_bm25_hybrid_requires_vector_db(): """Test BM25HybridSearchAlgorithm reports it requires vector database.""" algo = BM25HybridSearchAlgorithm() assert algo.requires_vector_db is True + + +@pytest.fixture +def patched_search(monkeypatch): + """Stub the embedding / BM25 / Qdrant deps of search() and return the + embed_with_usage mock so tests can assert how often the query was embedded.""" + embed = AsyncMock(return_value=([0.1, 0.2, 0.3], 7)) + svc = MagicMock() + svc.embed_with_usage = embed + monkeypatch.setattr( + "nextcloud_mcp_server.search.bm25_hybrid.get_embedding_service", lambda: svc + ) + + bm25 = MagicMock() + bm25.encode_async = AsyncMock(return_value={"indices": [1], "values": [0.5]}) + monkeypatch.setattr( + "nextcloud_mcp_server.search.bm25_hybrid.get_bm25_service", + AsyncMock(return_value=bm25), + ) + + qdrant = MagicMock() + empty = MagicMock() + empty.points = [] + qdrant.query_points = AsyncMock(return_value=empty) + monkeypatch.setattr( + "nextcloud_mcp_server.search.bm25_hybrid.get_qdrant_client", + AsyncMock(return_value=qdrant), + ) + + settings = MagicMock() + settings.get_collection_name.return_value = "test_collection" + settings.get_embedding_provider_family.return_value = "mistral" + monkeypatch.setattr( + "nextcloud_mcp_server.search.bm25_hybrid.get_settings", lambda: settings + ) + monkeypatch.setattr( + "nextcloud_mcp_server.search.bm25_hybrid.build_base_filter_conditions", + lambda **kwargs: [], + ) + return embed + + +@pytest.mark.unit +async def test_query_embedded_and_metered_once_across_doc_types(patched_search): + """nc_semantic_search calls search() once per doc_type on one instance with + the same query; the dense embedding (and its billed token count) must be + computed exactly once, not once per type.""" + embed = patched_search + algo = BM25HybridSearchAlgorithm() + + for dtype in ("note", "file", "deck_card"): + await algo.search(query="hello", user_id="alice", doc_type=dtype) + + assert embed.await_count == 1 # embedded once, not 3× + assert ( + algo.query_token_count == 7 + ) # single query's token count, not summed/overwritten + assert algo.query_embedding == [0.1, 0.2, 0.3] + + +@pytest.mark.unit +async def test_different_query_invalidates_cache(patched_search): + """A different query string re-embeds (and re-meters).""" + embed = patched_search + algo = BM25HybridSearchAlgorithm() + + await algo.search(query="hello", user_id="alice") + await algo.search(query="world", user_id="alice") + + assert embed.await_count == 2 diff --git a/tests/unit/server/__init__.py b/tests/unit/server/__init__.py new file mode 100644 index 00000000..1d5caa11 --- /dev/null +++ b/tests/unit/server/__init__.py @@ -0,0 +1 @@ +"""Unit tests for server-layer MCP tools.""" diff --git a/tests/unit/server/test_semantic_metering.py b/tests/unit/server/test_semantic_metering.py new file mode 100644 index 00000000..b0d6790a --- /dev/null +++ b/tests/unit/server/test_semantic_metering.py @@ -0,0 +1,124 @@ +"""Unit tests for the search-path usage-metering helper (Deck #67). + +``record_search_usage`` records the billable ``tokens_embedded`` event for a +semantic search. These pin the value mapping (query token count), the flag-off +no-op, the doc_types metadata bounding, and the best-effort failure path — +covering the server-tool metering wiring without standing up the full +``nc_semantic_search`` tool. +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from nextcloud_mcp_server.server import semantic + + +@pytest.fixture +def store_spy(monkeypatch): + """Patch UsageEventStore.shared() to return a spy store.""" + store = MagicMock() + store.record_usage_event = AsyncMock() + monkeypatch.setattr( + semantic.UsageEventStore, "shared", AsyncMock(return_value=store) + ) + return store + + +@pytest.mark.unit +async def test_records_query_token_count(store_spy): + """The event value is the query embedding's token count.""" + await semantic.record_search_usage( + enabled=True, + user_id="alice", + fusion="rrf", + doc_types=["note", "file"], + token_count=42, + ) + + store_spy.record_usage_event.assert_awaited_once() + kwargs = store_spy.record_usage_event.await_args.kwargs + assert kwargs["metric"] == "tokens_embedded" + assert kwargs["value"] == 42 + assert kwargs["enabled"] is True + assert kwargs["metadata"]["user_id"] == "alice" + assert kwargs["metadata"]["fusion"] == "rrf" + assert kwargs["metadata"]["doc_types"] == ["note", "file"] + + +@pytest.mark.unit +async def test_disabled_is_noop(store_spy): + """Flag off → no store access, no event.""" + await semantic.record_search_usage( + enabled=False, + user_id="alice", + fusion="rrf", + doc_types=None, + token_count=10, + ) + store_spy.record_usage_event.assert_not_awaited() + + +@pytest.mark.unit +async def test_none_token_count_records_zero(store_spy): + """A missing token count (pre-embed error) records value 0, not None.""" + await semantic.record_search_usage( + enabled=True, + user_id="alice", + fusion="dbsf", + doc_types=None, + token_count=None, + ) + kwargs = store_spy.record_usage_event.await_args.kwargs + assert kwargs["value"] == 0 + # None and [] both normalize to null for consistent IS NULL counting. + assert kwargs["metadata"]["doc_types"] is None + + +@pytest.mark.unit +async def test_empty_doc_types_normalizes_to_null(store_spy): + """An empty doc_types list normalizes to None, same as a None input, so a + metadata->'doc_types' IS NULL query counts the all-types case consistently.""" + await semantic.record_search_usage( + enabled=True, + user_id="alice", + fusion="rrf", + doc_types=[], + token_count=5, + ) + kwargs = store_spy.record_usage_event.await_args.kwargs + assert kwargs["metadata"]["doc_types"] is None + + +@pytest.mark.unit +async def test_doc_types_metadata_is_bounded(store_spy): + """A large doc_types list is truncated to the metadata cap.""" + many = [f"type-{i}" for i in range(40)] + await semantic.record_search_usage( + enabled=True, + user_id="alice", + fusion="rrf", + doc_types=many, + token_count=5, + ) + recorded = store_spy.record_usage_event.await_args.kwargs["metadata"]["doc_types"] + assert recorded == many[: semantic._USAGE_METADATA_MAX_DOC_TYPES] + + +@pytest.mark.unit +async def test_store_failure_is_swallowed(monkeypatch): + """A store-construction failure is logged, never raised into the search.""" + monkeypatch.setattr( + semantic.UsageEventStore, + "shared", + AsyncMock(side_effect=RuntimeError("boom")), + ) + + # Must not raise. + await semantic.record_search_usage( + enabled=True, + user_id="alice", + fusion="rrf", + doc_types=None, + token_count=7, + ) diff --git a/tests/unit/test_embedding_metrics.py b/tests/unit/test_embedding_metrics.py index 1f165480..d91c9e72 100644 --- a/tests/unit/test_embedding_metrics.py +++ b/tests/unit/test_embedding_metrics.py @@ -12,7 +12,10 @@ from __future__ import annotations import pytest from nextcloud_mcp_server.config import Settings -from nextcloud_mcp_server.observability.metrics import record_embedding +from nextcloud_mcp_server.observability.metrics import ( + record_embedding, + record_embedding_tokens, +) pytestmark = pytest.mark.unit @@ -120,3 +123,31 @@ class TestRecordEmbedding: assert metric_sample( "astrolabe_embedding_requests_total", {**labels, "status": "error"} ) == pytest.approx(1.0) + + +class TestRecordEmbeddingTokens: + """astrolabe_embedding_tokens_total — token cost split by index/query.""" + + def test_index_increments_by_token_count(self, metric_sample): + labels = {"provider": "tok-prov", "operation": "index"} + before = metric_sample("astrolabe_embedding_tokens_total", labels) + record_embedding_tokens("tok-prov", "index", 4242) + assert metric_sample( + "astrolabe_embedding_tokens_total", labels + ) == pytest.approx(before + 4242) + + def test_query_operation_is_separate_series(self, metric_sample): + labels = {"provider": "tok-prov", "operation": "query"} + before = metric_sample("astrolabe_embedding_tokens_total", labels) + record_embedding_tokens("tok-prov", "query", 7) + assert metric_sample( + "astrolabe_embedding_tokens_total", labels + ) == pytest.approx(before + 7) + + def test_zero_or_negative_is_noop(self, metric_sample): + labels = {"provider": "tok-noop", "operation": "index"} + record_embedding_tokens("tok-noop", "index", 0) + record_embedding_tokens("tok-noop", "index", -3) + assert metric_sample( + "astrolabe_embedding_tokens_total", labels + ) == pytest.approx(0.0) diff --git a/tests/unit/test_processor_metering.py b/tests/unit/test_processor_metering.py new file mode 100644 index 00000000..26838c29 --- /dev/null +++ b/tests/unit/test_processor_metering.py @@ -0,0 +1,104 @@ +"""Unit tests for the indexing-path usage-metering helper (Deck #67). + +``record_indexing_usage`` records the two billable events (``pages_embedded`` + +``tokens_embedded``) after a document's chunks are embedded. These cover the +value mapping, the flag/zero-chunk no-ops, and the best-effort failure path +without standing up the full document pipeline. +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from nextcloud_mcp_server.vector import processor + + +@pytest.fixture +def store_spy(monkeypatch): + """Patch UsageEventStore.shared() to return a spy store.""" + store = MagicMock() + store.record_usage_event = AsyncMock() + monkeypatch.setattr( + processor.UsageEventStore, "shared", AsyncMock(return_value=store) + ) + return store + + +@pytest.mark.unit +async def test_records_pages_embedded_and_token_count(store_spy): + """Both events fire: pages_embedded = chunk count, tokens_embedded = tokens.""" + await processor.record_indexing_usage( + enabled=True, + provider="mistral", + model="mistral-embed", + doc_type="file", + user_id="alice", + chunk_count=110, + token_count=4242, + total_chars=170826, + ) + + calls = store_spy.record_usage_event.await_args_list + by_metric = {c.kwargs["metric"]: c.kwargs["value"] for c in calls} + assert by_metric == {"pages_embedded": 110, "tokens_embedded": 4242} + for c in calls: + # Hot-path fast-gate + tenant-local attribution metadata. + assert c.kwargs["enabled"] is True + assert c.kwargs["metadata"]["provider"] == "mistral" + assert c.kwargs["metadata"]["model"] == "mistral-embed" + assert c.kwargs["metadata"]["user_id"] == "alice" + assert c.kwargs["metadata"]["doc_type"] == "file" + + +@pytest.mark.unit +async def test_disabled_is_noop(store_spy): + """Flag off → no store access, no events.""" + await processor.record_indexing_usage( + enabled=False, + provider="mistral", + model="mistral-embed", + doc_type="file", + user_id="alice", + chunk_count=10, + token_count=20, + total_chars=5, + ) + store_spy.record_usage_event.assert_not_awaited() + + +@pytest.mark.unit +async def test_zero_chunks_is_noop(store_spy): + """A document with no chunks records nothing (no zero-value rows).""" + await processor.record_indexing_usage( + enabled=True, + provider="mistral", + model="mistral-embed", + doc_type="file", + user_id="alice", + chunk_count=0, + token_count=0, + total_chars=0, + ) + store_spy.record_usage_event.assert_not_awaited() + + +@pytest.mark.unit +async def test_store_failure_is_swallowed(monkeypatch): + """A store-construction failure is logged, never raised into indexing.""" + monkeypatch.setattr( + processor.UsageEventStore, + "shared", + AsyncMock(side_effect=RuntimeError("boom")), + ) + + # Must not raise. + await processor.record_indexing_usage( + enabled=True, + provider="mistral", + model="mistral-embed", + doc_type="file", + user_id="alice", + chunk_count=3, + token_count=7, + total_chars=9, + ) diff --git a/tests/unit/test_usage_store.py b/tests/unit/test_usage_store.py index aa4b6944..0dd53512 100644 --- a/tests/unit/test_usage_store.py +++ b/tests/unit/test_usage_store.py @@ -95,7 +95,7 @@ async def test_flag_off_is_noop(storage, monkeypatch): """With metering disabled, nothing is written (zero DB work).""" _set_metering(monkeypatch, False) store = UsageEventStore(storage) - await store.record_usage_event(metric="pages_chunks", value=5) + await store.record_usage_event(metric="pages_embedded", value=5) assert await _count(storage) == 0 @@ -115,12 +115,12 @@ async def test_enabled_param_short_circuits_without_reading_settings( monkeypatch.setattr(store_module, "get_settings", _boom) store = UsageEventStore(storage) - await store.record_usage_event(metric="pages_chunks", value=1, enabled=False) + await store.record_usage_event(metric="pages_embedded", value=1, enabled=False) assert await _count(storage) == 0 eid = str(uuid.uuid4()) await store.record_usage_event( - metric="pages_chunks", value=1, event_id=eid, enabled=True + metric="pages_embedded", value=1, event_id=eid, enabled=True ) assert await _count(storage) == 1 @@ -131,7 +131,7 @@ async def test_insert_roundtrip(storage, monkeypatch): store = UsageEventStore(storage) eid = str(uuid.uuid4()) await store.record_usage_event( - metric="pages_chunks", + metric="pages_embedded", value=7, event_id=eid, metadata={"provider": "gateway"}, @@ -140,7 +140,7 @@ async def test_insert_roundtrip(storage, monkeypatch): assert row is not None # Postgres returns event_id as a uuid.UUID; normalize to str for compare. assert str(row[0]) == eid - assert row[2] == "pages_chunks" + assert row[2] == "pages_embedded" assert row[3] == 7 @@ -149,11 +149,11 @@ async def test_on_conflict_dedup(storage, monkeypatch): _set_metering(monkeypatch, True) store = UsageEventStore(storage) eid = str(uuid.uuid4()) - await store.record_usage_event(metric="pages_chunks", value=1, event_id=eid) - await store.record_usage_event(metric="embeddings_queries", value=99, event_id=eid) + await store.record_usage_event(metric="pages_embedded", value=1, event_id=eid) + await store.record_usage_event(metric="tokens_embedded", value=99, event_id=eid) assert await _count(storage) == 1 row = await _fetch(storage, eid) - assert row[2] == "pages_chunks" # DO NOTHING, not DO UPDATE + assert row[2] == "pages_embedded" # DO NOTHING, not DO UPDATE assert row[3] == 1 @@ -164,7 +164,7 @@ async def test_metadata_json_roundtrip(storage, monkeypatch): eid = str(uuid.uuid4()) meta = {"provider": "gateway", "model": "titan", "nested": {"chunks": 3}} await store.record_usage_event( - metric="pages_chunks", value=3, event_id=eid, metadata=meta + metric="pages_embedded", value=3, event_id=eid, metadata=meta ) row = await _fetch(storage, eid) raw = row[4] @@ -187,7 +187,7 @@ async def test_occurred_at_roundtrip(storage, monkeypatch): eid = str(uuid.uuid4()) when = datetime(2026, 1, 15, 12, 0, 0, tzinfo=timezone.utc) await store.record_usage_event( - metric="pages_chunks", value=1, event_id=eid, occurred_at=when + metric="pages_embedded", value=1, event_id=eid, occurred_at=when ) row = await _fetch(storage, eid) stored = row[1] @@ -207,7 +207,7 @@ async def test_metadata_none_is_null(storage, monkeypatch): store = UsageEventStore(storage) eid = str(uuid.uuid4()) await store.record_usage_event( - metric="embeddings_queries", value=1, event_id=eid, metadata=None + metric="tokens_embedded", value=1, event_id=eid, metadata=None ) row = await _fetch(storage, eid) assert row[4] is None @@ -233,7 +233,7 @@ async def test_best_effort_swallows_db_errors(storage, monkeypatch, caplog): # Must not raise. with caplog.at_level(logging.WARNING, logger="nextcloud_mcp_server.usage.store"): - await store.record_usage_event(metric="pages_chunks", value=1) + await store.record_usage_event(metric="pages_embedded", value=1) assert recorded, "record_db_operation should be called on the error path" assert recorded[-1][3] == "error" @@ -262,7 +262,7 @@ async def test_best_effort_swallows_unserializable_metadata( # Must not raise. with caplog.at_level(logging.WARNING, logger="nextcloud_mcp_server.usage.store"): await store.record_usage_event( - metric="pages_chunks", value=1, metadata=bad_metadata + metric="pages_embedded", value=1, metadata=bad_metadata ) # Nothing was written — the encode failed before the insert.