From 64318f0b255169c767dd2fde658ed14ba9d8eaac Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Mon, 8 Jun 2026 00:53:58 +0200 Subject: [PATCH 01/10] feat(usage): meter embedding tokens as embeddings_queries on both paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../embedding/gateway_client.py | 12 ++++ nextcloud_mcp_server/embedding/service.py | 16 +++++ nextcloud_mcp_server/providers/base.py | 35 ++++++++++ nextcloud_mcp_server/providers/bedrock.py | 38 +++++++++-- nextcloud_mcp_server/providers/mistral.py | 55 +++++++++++++-- nextcloud_mcp_server/providers/ollama.py | 46 ++++++++++++- nextcloud_mcp_server/providers/openai.py | 61 +++++++++++++++-- nextcloud_mcp_server/search/algorithms.py | 6 ++ nextcloud_mcp_server/search/bm25_hybrid.py | 8 ++- nextcloud_mcp_server/server/semantic.py | 18 +++-- nextcloud_mcp_server/vector/processor.py | 34 +++++++--- tests/unit/providers/test_bedrock.py | 67 ++++++++++++++++++ tests/unit/providers/test_gateway_provider.py | 64 +++++++++++++++++ tests/unit/providers/test_mistral.py | 62 +++++++++++++++++ tests/unit/providers/test_ollama.py | 68 +++++++++++++++++++ tests/unit/providers/test_openai.py | 40 +++++++++++ tests/unit/providers/test_provider_usage.py | 52 ++++++++++++++ 17 files changed, 647 insertions(+), 35 deletions(-) create mode 100644 tests/unit/providers/test_ollama.py create mode 100644 tests/unit/providers/test_provider_usage.py diff --git a/nextcloud_mcp_server/embedding/gateway_client.py b/nextcloud_mcp_server/embedding/gateway_client.py index 82f37257..cc624486 100644 --- a/nextcloud_mcp_server/embedding/gateway_client.py +++ b/nextcloud_mcp_server/embedding/gateway_client.py @@ -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) 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/providers/base.py b/nextcloud_mcp_server/providers/base.py index 3ea95827..9dad34c4 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,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: """ diff --git a/nextcloud_mcp_server/providers/bedrock.py b/nextcloud_mcp_server/providers/bedrock.py index a1b9c179..62161dca 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 = ( + 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): """ diff --git a/nextcloud_mcp_server/providers/mistral.py b/nextcloud_mcp_server/providers/mistral.py index b5f61584..f24283e7 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 ``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: diff --git a/nextcloud_mcp_server/providers/ollama.py b/nextcloud_mcp_server/providers/ollama.py index 1d5705c4..962c34a8 100644 --- a/nextcloud_mcp_server/providers/ollama.py +++ b/nextcloud_mcp_server/providers/ollama.py @@ -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): """ diff --git a/nextcloud_mcp_server/providers/openai.py b/nextcloud_mcp_server/providers/openai.py index 536697bd..72001587 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 = ( + int(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..06c5e991 100644 --- a/nextcloud_mcp_server/search/algorithms.py +++ b/nextcloud_mcp_server/search/algorithms.py @@ -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( diff --git a/nextcloud_mcp_server/search/bm25_hybrid.py b/nextcloud_mcp_server/search/bm25_hybrid.py index fc1ce41e..2c2e2203 100644 --- a/nextcloud_mcp_server/search/bm25_hybrid.py +++ b/nextcloud_mcp_server/search/bm25_hybrid.py @@ -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 diff --git a/nextcloud_mcp_server/server/semantic.py b/nextcloud_mcp_server/server/semantic.py index 8d78849f..89db92de 100644 --- a/nextcloud_mcp_server/server/semantic.py +++ b/nextcloud_mcp_server/server/semantic.py @@ -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, diff --git a/nextcloud_mcp_server/vector/processor.py b/nextcloud_mcp_server/vector/processor.py index b9bd2472..6dd7cac2 100644 --- a/nextcloud_mcp_server/vector/processor.py +++ b/nextcloud_mcp_server/vector/processor.py @@ -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(): 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..24b7fbc3 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,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.""" 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..0a457107 --- /dev/null +++ b/tests/unit/providers/test_ollama.py @@ -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() 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..66fe91d8 --- /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 ``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 From a0bb5642cb57f97c818cfff5ea5a73af11bcae36 Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Mon, 8 Jun 2026 01:07:22 +0200 Subject: [PATCH 02/10] fix(usage): embed query once across doc_types; address review round 1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-1 claude-review findings: - 🔴 Multi-doc_type search billed N embedding calls as 1. nc_semantic_search loops search() once per doc_type on one BM25HybridSearchAlgorithm instance, and each call re-embedded the query, so only the last query_token_count was recorded. Cache the dense embedding per query on the (per-request) instance so the query is embedded — and metered — exactly once regardless of how many doc_types are searched. This also removes the redundant per-type embed work and avoids billing a user N× for one logical query. - 🟡 Ollama embed() now delegates to embed_with_usage() so single and batch embeds use the same /api/embed endpoint (was the legacy /api/embeddings), keeping _detect_dimension and other embed() callers consistent. - 🟢 round() instead of truncating int() when coercing provider-reported token counts (forward-compatible if a provider ever returns a float). Tests: per-instance query-embedding cache (embedded once across 3 doc_types; re-embeds on a different query). Deferred (stated on the PR): mistral/openai single-embed dual path (changes tested error/request semantics on the cloud-critical path — separate refactor), bedrock boto3 sync-in-async (pre-existing; no new invoke_model calls per doc). Deck #67. Co-Authored-By: Claude Opus 4.8 (1M context) --- nextcloud_mcp_server/providers/bedrock.py | 2 +- nextcloud_mcp_server/providers/mistral.py | 2 +- nextcloud_mcp_server/providers/ollama.py | 19 +++--- nextcloud_mcp_server/providers/openai.py | 2 +- nextcloud_mcp_server/search/bm25_hybrid.py | 31 +++++++--- nextcloud_mcp_server/server/semantic.py | 7 ++- tests/unit/search/test_bm25_hybrid.py | 71 ++++++++++++++++++++++ 7 files changed, 108 insertions(+), 26 deletions(-) diff --git a/nextcloud_mcp_server/providers/bedrock.py b/nextcloud_mcp_server/providers/bedrock.py index 62161dca..d8afa40f 100644 --- a/nextcloud_mcp_server/providers/bedrock.py +++ b/nextcloud_mcp_server/providers/bedrock.py @@ -194,7 +194,7 @@ class BedrockProvider(Provider): token_count = response_body.get("inputTextTokenCount") tokens = ( - int(token_count) + round(token_count) if isinstance(token_count, (int, float)) else self._estimate_tokens([text]) ) diff --git a/nextcloud_mcp_server/providers/mistral.py b/nextcloud_mcp_server/providers/mistral.py index f24283e7..1ba22800 100644 --- a/nextcloud_mcp_server/providers/mistral.py +++ b/nextcloud_mcp_server/providers/mistral.py @@ -209,7 +209,7 @@ class MistralProvider(Provider): # gives an int, but test doubles / partial responses can surface a # non-numeric attribute — fall back to the estimate there. tokens = ( - int(total_tokens) + round(total_tokens) if isinstance(total_tokens, (int, float)) else self._estimate_tokens(batch) ) diff --git a/nextcloud_mcp_server/providers/ollama.py b/nextcloud_mcp_server/providers/ollama.py index 962c34a8..57f8388c 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 @@ -166,7 +161,7 @@ class OllamaProvider(Provider): prompt_eval = data.get("prompt_eval_count") total_tokens += ( - int(prompt_eval) + round(prompt_eval) if isinstance(prompt_eval, (int, float)) else self._estimate_tokens(batch) ) diff --git a/nextcloud_mcp_server/providers/openai.py b/nextcloud_mcp_server/providers/openai.py index 72001587..dfa7a8a1 100644 --- a/nextcloud_mcp_server/providers/openai.py +++ b/nextcloud_mcp_server/providers/openai.py @@ -233,7 +233,7 @@ class OpenAIProvider(Provider): # gives an int, but test doubles / partial responses can surface a # non-numeric attribute — fall back to the estimate there. tokens = ( - int(total_tokens) + round(total_tokens) if isinstance(total_tokens, (int, float)) else self._estimate_tokens(batch) ) diff --git a/nextcloud_mcp_server/search/bm25_hybrid.py b/nextcloud_mcp_server/search/bm25_hybrid.py index 2c2e2203..08fb64ba 100644 --- a/nextcloud_mcp_server/search/bm25_hybrid.py +++ b/nextcloud_mcp_server/search/bm25_hybrid.py @@ -56,6 +56,10 @@ class BM25HybridSearchAlgorithm(SearchAlgorithm): self.score_threshold = score_threshold self.fusion = models.Fusion.RRF if fusion == "rrf" else models.Fusion.DBSF self.fusion_name = fusion + # The query string whose dense embedding is cached in + # ``self.query_embedding`` — lets repeated search() calls on this + # per-request instance (the doc_types loop) reuse one embedding. + self._embedded_query: str | None = None @property def name(self) -> str: @@ -128,17 +132,28 @@ 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, 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 + 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 logger.debug("Generated dense embedding (dimension=%s)", len(dense_embedding)) # Generate sparse embedding for BM25 keyword search diff --git a/nextcloud_mcp_server/server/semantic.py b/nextcloud_mcp_server/server/semantic.py index 89db92de..97fb8338 100644 --- a/nextcloud_mcp_server/server/semantic.py +++ b/nextcloud_mcp_server/server/semantic.py @@ -537,9 +537,10 @@ def configure_semantic_tools(mcp: FastMCP): # # 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). + # 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. 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 diff --git a/tests/unit/search/test_bm25_hybrid.py b/tests/unit/search/test_bm25_hybrid.py index a80b57bc..90fccba8 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,72 @@ 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" + 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 From d15ce627ab8b5faf9804d41657bd9ceb14cea921 Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Mon, 8 Jun 2026 01:22:03 +0200 Subject: [PATCH 03/10] refactor(usage): extract indexing metering helper; address review round 2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-2 claude-review findings: - 🟡 Base-class recursion invariant: documented on embed_with_usage / embed_batch_with_usage that a provider overriding embed()/embed_batch() to delegate to the *_with_usage variant MUST also override that variant, or the two recurse. (No recursion today; the shipped providers pair the overrides.) - 🟡 Processor metering had no unit test: extracted the two-event recording into a module-level record_indexing_usage() helper and added tests/unit/test_processor_metering.py (value mapping, flag/zero-chunk no-ops, best-effort failure swallowed). - 🟡 SonarQube hotspots (python:S5332) were 3 http:// URLs in the new test fixtures (mock hosts, never contacted) blocking the quality gate (new_security_hotspots_reviewed). Switched them to https:// so no hotspot is raised. - 🟢 Zero-chunk guard: record_indexing_usage() no-ops when chunk_count == 0, so an empty document no longer writes zero-value billing rows. Deferred (stated on the PR): Mistral x.index-or-0 sort key (pre-existing, equivalent), CHANGELOG note for the Ollama /api/embed switch (CHANGELOG is commitizen-generated from commit bodies, which document it), class-var query_token_count (safe under the per-request instance pattern). Deck #67. Co-Authored-By: Claude Opus 4.8 (1M context) --- nextcloud_mcp_server/providers/base.py | 11 ++ nextcloud_mcp_server/vector/processor.py | 123 +++++++++++------- tests/unit/providers/test_gateway_provider.py | 6 +- tests/unit/providers/test_ollama.py | 5 +- tests/unit/test_processor_metering.py | 104 +++++++++++++++ 5 files changed, 196 insertions(+), 53 deletions(-) create mode 100644 tests/unit/test_processor_metering.py diff --git a/nextcloud_mcp_server/providers/base.py b/nextcloud_mcp_server/providers/base.py index 9dad34c4..7149b13a 100644 --- a/nextcloud_mcp_server/providers/base.py +++ b/nextcloud_mcp_server/providers/base.py @@ -75,6 +75,12 @@ class Provider(ABC): 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. + + 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]) @@ -86,6 +92,11 @@ class Provider(ABC): 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) diff --git a/nextcloud_mcp_server/vector/processor.py b/nextcloud_mcp_server/vector/processor.py index 6dd7cac2..4a14fd73 100644 --- a/nextcloud_mcp_server/vector/processor.py +++ b/nextcloud_mcp_server/vector/processor.py @@ -112,6 +112,65 @@ 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_chunks`` is the volume (chunks embedded); ``embeddings_queries`` 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). + + 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). + await store.record_usage_event( + metric="pages_chunks", value=chunk_count, metadata=metadata, enabled=True + ) + await store.record_usage_event( + metric="embeddings_queries", + 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], @@ -825,55 +884,21 @@ 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: - # 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=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 (indexing embeddings) skipped", - exc_info=True, - ) + # 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_gateway_provider.py b/tests/unit/providers/test_gateway_provider.py index 24b7fbc3..3a834358 100644 --- a/tests/unit/providers/test_gateway_provider.py +++ b/tests/unit/providers/test_gateway_provider.py @@ -331,8 +331,9 @@ def test_trailing_slash_base_url_normalized(): 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="http://gw:8083/v1", embedding_model="mistral/mistral-embed" + base_url="https://gw:8083/v1", embedding_model="mistral/mistral-embed" ) order: list[str] = [] @@ -364,8 +365,9 @@ async def test_gateway_embed_with_usage_forwards_after_bearer(monkeypatch): 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="http://gw:8083/v1", embedding_model="mistral/mistral-embed" + base_url="https://gw:8083/v1", embedding_model="mistral/mistral-embed" ) ensured = {"n": 0} diff --git a/tests/unit/providers/test_ollama.py b/tests/unit/providers/test_ollama.py index 0a457107..94505dde 100644 --- a/tests/unit/providers/test_ollama.py +++ b/tests/unit/providers/test_ollama.py @@ -15,8 +15,9 @@ 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") + # 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 diff --git a/tests/unit/test_processor_metering.py b/tests/unit/test_processor_metering.py new file mode 100644 index 00000000..9a9630ef --- /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_chunks`` + +``embeddings_queries``) 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_chunks_and_token_count(store_spy): + """Both events fire: pages_chunks = chunk count, embeddings_queries = 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_chunks": 110, "embeddings_queries": 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, + ) From 9ac9e1ab097279d1d0e9674d5f10f95b13d5f9ba Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Mon, 8 Jun 2026 01:29:36 +0200 Subject: [PATCH 04/10] fix(usage): drop redundant GatewayProvider.embed_batch override (round 3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-3 claude-review finding: - 🟡 Double _ensure_bearer() on gateway.embed_batch(). Round 1 made OpenAIProvider.embed_batch() delegate to embed_batch_with_usage(); because GatewayProvider overrode both embed_batch() and embed_batch_with_usage() (each calling _ensure_bearer), gateway.embed_batch() refreshed the bearer twice (the second a cache-hit no-op). Remove the now-redundant embed_batch() override: OpenAI's embed_batch() routes through embed_batch_with_usage(), which the gateway still overrides, so the bearer refreshes exactly once on every path. The remaining two overrides (embed + embed_batch_with_usage) cover all four entrypoints; documented the topology. - 🟢 Added test_gateway_embed_batch_ensures_bearer_once locking in the single refresh. Cohere token-fallback (🟢 nit) is already covered by test_bedrock_with_usage_estimates_when_token_count_absent. Deck #67. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../embedding/gateway_client.py | 19 ++++++------ tests/unit/providers/test_gateway_provider.py | 31 +++++++++++++++++++ 2 files changed, 40 insertions(+), 10 deletions(-) diff --git a/nextcloud_mcp_server/embedding/gateway_client.py b/nextcloud_mcp_server/embedding/gateway_client.py index cc624486..f372c7f8 100644 --- a/nextcloud_mcp_server/embedding/gateway_client.py +++ b/nextcloud_mcp_server/embedding/gateway_client.py @@ -217,22 +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]]: - 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) diff --git a/tests/unit/providers/test_gateway_provider.py b/tests/unit/providers/test_gateway_provider.py index 3a834358..e7f5ad95 100644 --- a/tests/unit/providers/test_gateway_provider.py +++ b/tests/unit/providers/test_gateway_provider.py @@ -393,6 +393,37 @@ async def test_gateway_embed_batch_with_usage_forwards_after_bearer(monkeypatch) assert ensured["n"] == 1 +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.""" From df03d33fd49629735b68e598c07a64fe578579d8 Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Mon, 8 Jun 2026 01:43:49 +0200 Subject: [PATCH 05/10] test(usage): cover search metering hook; log dedup metering skip (round 4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-4 claude-review findings (no blockers): - 🟡 Untested server-layer metering hook (raised across rounds): extracted the nc_semantic_search embeddings_queries recording into a module-level record_search_usage() helper (mirroring record_indexing_usage) and added tests/unit/server/test_semantic_metering.py — value = query token count, flag-off no-op, None token → 0, doc_types metadata bounding, best-effort failure swallowed. - 🟡 Dedup-hit skipped metering invisibly: the existing dedup info log now states "no embedding/usage recorded" so a "fewer embeddings_queries rows than expected" audit lands on the dedup path directly. Deferred 🟢 nits (stated on the PR): search 0-token rows are recorded deliberately (the query embedding ran; zero is a sum no-op) — documented in the helper; embed_tokens closure locality and the OpenAI embed() dual path are unchanged (correct as-is / separate refactor). Deck #67. Co-Authored-By: Claude Opus 4.8 (1M context) --- nextcloud_mcp_server/server/semantic.py | 127 +++++++++++--------- nextcloud_mcp_server/vector/processor.py | 6 +- tests/unit/server/__init__.py | 1 + tests/unit/server/test_semantic_metering.py | 109 +++++++++++++++++ 4 files changed, 188 insertions(+), 55 deletions(-) create mode 100644 tests/unit/server/__init__.py create mode 100644 tests/unit/server/test_semantic_metering.py diff --git a/nextcloud_mcp_server/server/semantic.py b/nextcloud_mcp_server/server/semantic.py index 97fb8338..952a46e2 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 ``embeddings_queries`` 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. + + 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="embeddings_queries", + 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 (embeddings_queries) skipped", exc_info=True + ) + + def configure_semantic_tools(mcp: FastMCP): """Configure semantic search tools for MCP server.""" @@ -527,60 +585,21 @@ 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 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, 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. 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 - # 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=search_algo.query_token_count or 0, - 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, - ) + # Usage metering (Deck #67): record the query embedding's token + # count as a billable 'embeddings_queries' 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. + 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/vector/processor.py b/nextcloud_mcp_server/vector/processor.py index 4a14fd73..a44efbef 100644 --- a/nextcloud_mcp_server/vector/processor.py +++ b/nextcloud_mcp_server/vector/processor.py @@ -657,9 +657,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 embeddings_queries 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, 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..e93a79a0 --- /dev/null +++ b/tests/unit/server/test_semantic_metering.py @@ -0,0 +1,109 @@ +"""Unit tests for the search-path usage-metering helper (Deck #67). + +``record_search_usage`` records the billable ``embeddings_queries`` 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"] == "embeddings_queries" + 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_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, + ) From 141663bb0705e69e2358568c875299cac31b909e Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Mon, 8 Jun 2026 01:53:26 +0200 Subject: [PATCH 06/10] test(usage): close round-5 nits (empty doc_types, consistency tidy-ups) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-5 claude-review (merge-ready; all nits): - 🟡 Added test_empty_doc_types_normalizes_to_null pinning doc_types=[] → None in record_search_usage metadata (matches the None case). - 🟡 record_search_usage docstring now notes nc_semantic_search_answer always meters with doc_types=None (it exposes no doc_types parameter). - 🟢 BM25HybridSearchAlgorithm.__init__ now sets query_embedding / query_token_count alongside _embedded_query, so all three cache fields are instance attributes from construction (was relying on the class-level SearchAlgorithm defaults). - 🟢 Ollama embed_batch_with_usage caches _dimension inline (mirrors OpenAI/Mistral), so the dimension is set via any embed path. - 🟢 record_indexing_usage documents the independent-record / partial-failure semantics under SUM aggregation. Deck #67. Co-Authored-By: Claude Opus 4.8 (1M context) --- nextcloud_mcp_server/providers/ollama.py | 5 +++++ nextcloud_mcp_server/search/bm25_hybrid.py | 10 +++++++--- nextcloud_mcp_server/server/semantic.py | 4 +++- nextcloud_mcp_server/vector/processor.py | 4 ++++ tests/unit/server/test_semantic_metering.py | 15 +++++++++++++++ 5 files changed, 34 insertions(+), 4 deletions(-) diff --git a/nextcloud_mcp_server/providers/ollama.py b/nextcloud_mcp_server/providers/ollama.py index 57f8388c..b541750d 100644 --- a/nextcloud_mcp_server/providers/ollama.py +++ b/nextcloud_mcp_server/providers/ollama.py @@ -159,6 +159,11 @@ class OllamaProvider(Provider): data = response.json() all_embeddings.extend(data["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 = data.get("prompt_eval_count") total_tokens += ( round(prompt_eval) diff --git a/nextcloud_mcp_server/search/bm25_hybrid.py b/nextcloud_mcp_server/search/bm25_hybrid.py index 08fb64ba..7c6f8c51 100644 --- a/nextcloud_mcp_server/search/bm25_hybrid.py +++ b/nextcloud_mcp_server/search/bm25_hybrid.py @@ -56,9 +56,13 @@ class BM25HybridSearchAlgorithm(SearchAlgorithm): self.score_threshold = score_threshold self.fusion = models.Fusion.RRF if fusion == "rrf" else models.Fusion.DBSF self.fusion_name = fusion - # The query string whose dense embedding is cached in - # ``self.query_embedding`` — lets repeated search() calls on this - # per-request instance (the doc_types loop) reuse one embedding. + # Per-request query-embedding cache. ``_embedded_query`` is the query + # string whose dense embedding is held in ``query_embedding`` — repeated + # search() calls on this instance (the doc_types loop) reuse it. These + # shadow the class-level defaults on SearchAlgorithm; set here so all + # three cache fields are instance attributes from construction. + self.query_embedding: list[float] | None = None + self.query_token_count: int | None = None self._embedded_query: str | None = None @property diff --git a/nextcloud_mcp_server/server/semantic.py b/nextcloud_mcp_server/server/semantic.py index 952a46e2..c7524833 100644 --- a/nextcloud_mcp_server/server/semantic.py +++ b/nextcloud_mcp_server/server/semantic.py @@ -70,7 +70,9 @@ async def record_search_usage( 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. + 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 diff --git a/nextcloud_mcp_server/vector/processor.py b/nextcloud_mcp_server/vector/processor.py index a44efbef..adeeddbb 100644 --- a/nextcloud_mcp_server/vector/processor.py +++ b/nextcloud_mcp_server/vector/processor.py @@ -153,6 +153,10 @@ async def record_indexing_usage( 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_chunks somehow raised mid-way, embeddings_- + # queries would be skipped, leaving an unmatched pages_chunks row — + # acceptable under the (day, metric) SUM-aggregation billing model. await store.record_usage_event( metric="pages_chunks", value=chunk_count, metadata=metadata, enabled=True ) diff --git a/tests/unit/server/test_semantic_metering.py b/tests/unit/server/test_semantic_metering.py index e93a79a0..0b581767 100644 --- a/tests/unit/server/test_semantic_metering.py +++ b/tests/unit/server/test_semantic_metering.py @@ -75,6 +75,21 @@ async def test_none_token_count_records_zero(store_spy): 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.""" From ddefb037019ff3e1a9f967c87c606bff490c36dd Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Mon, 8 Jun 2026 02:00:24 +0200 Subject: [PATCH 07/10] test(usage): mark new gateway usage tests with @pytest.mark.unit (round 6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-6 claude-review (ready to merge): the three new GatewayProvider usage/bearer tests lacked @pytest.mark.unit, so `pytest -m unit` skipped them even though every other new test in this PR is marked. Add the marker to the three new tests (leaving the pre-existing unmarked tests in the file alone). Remaining 🟢 items (OpenAI embed() dual path, recursion-invariant runtime enforcement, Bedrock sync-in-async) are acknowledged deferrals — separate refactors, unchanged. Deck #67. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/unit/providers/test_gateway_provider.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit/providers/test_gateway_provider.py b/tests/unit/providers/test_gateway_provider.py index e7f5ad95..45f798a8 100644 --- a/tests/unit/providers/test_gateway_provider.py +++ b/tests/unit/providers/test_gateway_provider.py @@ -328,6 +328,7 @@ 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.""" @@ -363,6 +364,7 @@ async def test_gateway_embed_with_usage_forwards_after_bearer(monkeypatch): 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). @@ -393,6 +395,7 @@ async def test_gateway_embed_batch_with_usage_forwards_after_bearer(monkeypatch) 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 From 973f80e7b98b9ca89287877da672a01526106284 Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Mon, 8 Jun 2026 13:17:53 +0200 Subject: [PATCH 08/10] =?UTF-8?q?feat(usage):=20rename=20metrics=20?= =?UTF-8?q?=E2=86=92=20tokens=5Fembedded/pages=5Fembedded=20+=20export=20t?= =?UTF-8?q?oken=20cost=20to=20Prometheus?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Billing product model finalized (Deck #281): bill pages externally, record tokens internally. Rename the data-plane metric literals to match the now- canonical contract (Deck #284) — the control plane's METRIC_EVENT_NAMES is already renamed, so the old names would be unmapped and never sync to Stripe. Rename (values unchanged): - embeddings_queries → tokens_embedded (value = real token count, already emitted by this PR; the unit upstream providers bill on). - pages_chunks → pages_embedded (value kept as len(chunk_texts) interim; TODO(#282): real normalized "pages indexed" count — real pages for paginated types, chars/tokens-per-page constant otherwise — is deferred to the instrumentation card, this only lands the name/contract). - All literals, log strings, docstrings, comments, the migration comment, and tests renamed; grep confirms zero old strings remain. Observability (new): export embedding token cost to Prometheus as astrolabe_embedding_tokens_total{provider,operation} (operation = index|query) so the billed cost unit is visible in Grafana, not just the per-tenant billing DB. Dedicated counter (doesn't inflate the existing chunk/request metrics) and always-on (independent of USAGE_METERING_ENABLED, so OSS/self-host gets it). Wired on both the indexing batch embed and the search query embed (query inside the per-request cache-miss branch, so reused embeddings aren't double-counted). Note: the rename orphans any pre-existing embeddings_queries/pages_chunks rows in tenant app DBs (CP no longer maps them) — acceptable; pipeline is inert with throwaway dev/sandbox data. Deck #284 (folded into PR #875). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../20260607_1200_007_add_usage_events.py | 2 +- nextcloud_mcp_server/observability/metrics.py | 29 +++++++++++++++ nextcloud_mcp_server/providers/base.py | 2 +- nextcloud_mcp_server/providers/mistral.py | 2 +- nextcloud_mcp_server/search/algorithms.py | 2 +- nextcloud_mcp_server/search/bm25_hybrid.py | 12 ++++++- nextcloud_mcp_server/server/semantic.py | 10 +++--- nextcloud_mcp_server/usage/store.py | 4 +-- nextcloud_mcp_server/vector/processor.py | 35 +++++++++++++------ tests/unit/providers/test_provider_usage.py | 2 +- tests/unit/search/test_bm25_hybrid.py | 1 + tests/unit/server/test_semantic_metering.py | 4 +-- tests/unit/test_embedding_metrics.py | 33 ++++++++++++++++- tests/unit/test_processor_metering.py | 10 +++--- tests/unit/test_usage_store.py | 26 +++++++------- 15 files changed, 129 insertions(+), 45 deletions(-) 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/observability/metrics.py b/nextcloud_mcp_server/observability/metrics.py index 921617fa..99688f20 100644 --- a/nextcloud_mcp_server/observability/metrics.py +++ b/nextcloud_mcp_server/observability/metrics.py @@ -338,6 +338,16 @@ 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. +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 +737,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 7149b13a..bda908ae 100644 --- a/nextcloud_mcp_server/providers/base.py +++ b/nextcloud_mcp_server/providers/base.py @@ -73,7 +73,7 @@ class Provider(ABC): 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 + 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 diff --git a/nextcloud_mcp_server/providers/mistral.py b/nextcloud_mcp_server/providers/mistral.py index 1ba22800..bacdadf0 100644 --- a/nextcloud_mcp_server/providers/mistral.py +++ b/nextcloud_mcp_server/providers/mistral.py @@ -143,7 +143,7 @@ class MistralProvider(Provider): 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). + to record ``tokens_embedded`` by tokens (Deck #67). """ if not self.supports_embeddings: raise NotImplementedError(_NO_EMBEDDING_MODEL_MSG) diff --git a/nextcloud_mcp_server/search/algorithms.py b/nextcloud_mcp_server/search/algorithms.py index 06c5e991..9c94da63 100644 --- a/nextcloud_mcp_server/search/algorithms.py +++ b/nextcloud_mcp_server/search/algorithms.py @@ -288,7 +288,7 @@ class SearchAlgorithm(ABC): 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 + ``tokens_embedded`` by tokens (Deck #67). The instance is per-request, so this side-channel is concurrency-safe. """ diff --git a/nextcloud_mcp_server/search/bm25_hybrid.py b/nextcloud_mcp_server/search/bm25_hybrid.py index 7c6f8c51..bdbd32dc 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 ( @@ -158,6 +161,13 @@ class BM25HybridSearchAlgorithm(SearchAlgorithm): 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/server/semantic.py b/nextcloud_mcp_server/server/semantic.py index c7524833..57274ef9 100644 --- a/nextcloud_mcp_server/server/semantic.py +++ b/nextcloud_mcp_server/server/semantic.py @@ -64,7 +64,7 @@ async def record_search_usage( doc_types: list[str] | None, token_count: int | None, ) -> None: - """Record the billable ``embeddings_queries`` event for one semantic search. + """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 @@ -89,7 +89,7 @@ async def record_search_usage( try: store = await UsageEventStore.shared() await store.record_usage_event( - metric="embeddings_queries", + metric="tokens_embedded", value=token_count or 0, metadata={ "user_id": user_id, @@ -111,9 +111,7 @@ async def record_search_usage( # (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 - ) + logger.warning("usage metering hook (tokens_embedded) skipped", exc_info=True) def configure_semantic_tools(mcp: FastMCP): @@ -588,7 +586,7 @@ def configure_semantic_tools(mcp: FastMCP): logger.info("Returning %d results from BM25 hybrid search", len(results)) # Usage metering (Deck #67): record the query embedding's token - # count as a billable 'embeddings_queries' event. query_token_count + # 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 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 adeeddbb..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, @@ -125,10 +126,16 @@ async def record_indexing_usage( ) -> None: """Record the two billable usage events for one embedded document. - ``pages_chunks`` is the volume (chunks embedded); ``embeddings_queries`` 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). + ``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 @@ -154,14 +161,19 @@ async def record_indexing_usage( # 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_chunks somehow raised mid-way, embeddings_- - # queries would be skipped, leaving an unmatched pages_chunks row — - # acceptable under the (day, metric) SUM-aggregation billing model. + # 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( - metric="pages_chunks", value=chunk_count, metadata=metadata, enabled=True + # 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="embeddings_queries", + metric="tokens_embedded", value=token_count, metadata=metadata, enabled=True, @@ -662,7 +674,7 @@ async def _index_document( user_id=doc_task.user_id, ) # No embedding ran, so no usage is recorded here — stated - # explicitly so a "fewer embeddings_queries rows than expected" + # 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( @@ -892,6 +904,9 @@ async def _index_document( chunks=len(chunk_texts), chars=total_chars, ) + # 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 diff --git a/tests/unit/providers/test_provider_usage.py b/tests/unit/providers/test_provider_usage.py index 66fe91d8..54772cf6 100644 --- a/tests/unit/providers/test_provider_usage.py +++ b/tests/unit/providers/test_provider_usage.py @@ -1,6 +1,6 @@ """Token-usage surfacing: the Provider ABC estimate default + SimpleProvider. -The usage-metering hooks (Deck #67) bill ``embeddings_queries`` by tokens. Real +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. diff --git a/tests/unit/search/test_bm25_hybrid.py b/tests/unit/search/test_bm25_hybrid.py index 90fccba8..835adc41 100644 --- a/tests/unit/search/test_bm25_hybrid.py +++ b/tests/unit/search/test_bm25_hybrid.py @@ -85,6 +85,7 @@ def patched_search(monkeypatch): 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 ) diff --git a/tests/unit/server/test_semantic_metering.py b/tests/unit/server/test_semantic_metering.py index 0b581767..b0d6790a 100644 --- a/tests/unit/server/test_semantic_metering.py +++ b/tests/unit/server/test_semantic_metering.py @@ -1,6 +1,6 @@ """Unit tests for the search-path usage-metering helper (Deck #67). -``record_search_usage`` records the billable ``embeddings_queries`` event for a +``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 @@ -38,7 +38,7 @@ async def test_records_query_token_count(store_spy): store_spy.record_usage_event.assert_awaited_once() kwargs = store_spy.record_usage_event.await_args.kwargs - assert kwargs["metric"] == "embeddings_queries" + assert kwargs["metric"] == "tokens_embedded" assert kwargs["value"] == 42 assert kwargs["enabled"] is True assert kwargs["metadata"]["user_id"] == "alice" 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 index 9a9630ef..26838c29 100644 --- a/tests/unit/test_processor_metering.py +++ b/tests/unit/test_processor_metering.py @@ -1,7 +1,7 @@ """Unit tests for the indexing-path usage-metering helper (Deck #67). -``record_indexing_usage`` records the two billable events (``pages_chunks`` + -``embeddings_queries``) after a document's chunks are embedded. These cover the +``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. """ @@ -25,8 +25,8 @@ def store_spy(monkeypatch): @pytest.mark.unit -async def test_records_pages_chunks_and_token_count(store_spy): - """Both events fire: pages_chunks = chunk count, embeddings_queries = tokens.""" +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", @@ -40,7 +40,7 @@ async def test_records_pages_chunks_and_token_count(store_spy): 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_chunks": 110, "embeddings_queries": 4242} + 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 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. From 93698329774356d7721f590c1255354a948f9c6c Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Mon, 8 Jun 2026 13:52:51 +0200 Subject: [PATCH 09/10] refactor(search): structural per-instance query side-channel; doc search billing gap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-7 claude-review (no blockers): - 🟡 query_token_count/query_embedding were class-level defaults on SearchAlgorithm, relying on each subclass's __init__ to shadow them. Added SearchAlgorithm.__init__ that sets both as instance attributes and had BM25HybridSearchAlgorithm + SemanticSearchAlgorithm call super().__init__(), so per-request concurrency isolation is structural, not by convention. - 🟡 Documented the v1 search-path billing gap: record_search_usage fires only on a fully successful search, so if the query embed succeeded (provider billed + Prometheus recorded) but a later step (Qdrant/verify) raised, no tokens_embedded billing row is written. Added a NOTE at the call site. Left as-is (reasons in PR reply): deployment sequencing (CP METRIC_EVENT_NAMES already renamed; pipeline inert); Ollama _detect_dimension double dimension-set (idempotent, same value); SonarQube issues — 1 is the deliberate TODO(#282) (INFO), 4 are S7503 false positives on async test stubs that must be awaitable (gate green). Deck #284. Co-Authored-By: Claude Opus 4.8 (1M context) --- nextcloud_mcp_server/search/algorithms.py | 10 ++++++++++ nextcloud_mcp_server/search/bm25_hybrid.py | 13 ++++++------- nextcloud_mcp_server/search/semantic.py | 1 + nextcloud_mcp_server/server/semantic.py | 8 ++++++++ 4 files changed, 25 insertions(+), 7 deletions(-) diff --git a/nextcloud_mcp_server/search/algorithms.py b/nextcloud_mcp_server/search/algorithms.py index 9c94da63..b4990b0c 100644 --- a/nextcloud_mcp_server/search/algorithms.py +++ b/nextcloud_mcp_server/search/algorithms.py @@ -292,9 +292,19 @@ class SearchAlgorithm(ABC): 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( self, diff --git a/nextcloud_mcp_server/search/bm25_hybrid.py b/nextcloud_mcp_server/search/bm25_hybrid.py index bdbd32dc..9b19cb40 100644 --- a/nextcloud_mcp_server/search/bm25_hybrid.py +++ b/nextcloud_mcp_server/search/bm25_hybrid.py @@ -56,16 +56,15 @@ 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 - # Per-request query-embedding cache. ``_embedded_query`` is the query - # string whose dense embedding is held in ``query_embedding`` — repeated - # search() calls on this instance (the doc_types loop) reuse it. These - # shadow the class-level defaults on SearchAlgorithm; set here so all - # three cache fields are instance attributes from construction. - self.query_embedding: list[float] | None = None - self.query_token_count: int | None = None + # ``_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 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 57274ef9..382236a4 100644 --- a/nextcloud_mcp_server/server/semantic.py +++ b/nextcloud_mcp_server/server/semantic.py @@ -593,6 +593,14 @@ def configure_semantic_tools(mcp: FastMCP): # query is embedded — and metered — exactly once regardless of how # many doc_types were searched. See record_search_usage for the # metric/privacy details. + # + # 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, From 28de737c264aba22cad80f72e15441586f7a58ff Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Mon, 8 Jun 2026 14:09:18 +0200 Subject: [PATCH 10/10] docs(usage): note Ollama batch-token semantics + Prometheus/billing query divergence MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-8 claude-review (no blockers; comment-only): - 🟡 Documented that Ollama's /api/embed prompt_eval_count is assumed batch-level total and is unverified against a live instance (Ollama isn't the Cloud billing provider); if it proves last-item-only, switch to per-item summing. The char estimate already covers versions that omit the field. - 🟡 Noted on the astrolabe_embedding_tokens_total counter that operation="query" is recorded pre-Qdrant, so it can legitimately exceed the billing-store tokens_embedded aggregate when a search fails post-embed — dashboards shouldn't alert on that healthy gap. Deferred (reviewer: "minor nit, acceptable"): record_indexing_usage awaited in the task group — the group awaits all child tasks regardless, the write is best-effort + fast, and start_soon would need the tg threaded into the closure for marginal gain. Deck #284. Co-Authored-By: Claude Opus 4.8 (1M context) --- nextcloud_mcp_server/observability/metrics.py | 6 ++++++ nextcloud_mcp_server/providers/ollama.py | 7 +++++++ 2 files changed, 13 insertions(+) diff --git a/nextcloud_mcp_server/observability/metrics.py b/nextcloud_mcp_server/observability/metrics.py index 99688f20..8e18adb8 100644 --- a/nextcloud_mcp_server/observability/metrics.py +++ b/nextcloud_mcp_server/observability/metrics.py @@ -342,6 +342,12 @@ embedding_chars_total = Counter( # 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)", diff --git a/nextcloud_mcp_server/providers/ollama.py b/nextcloud_mcp_server/providers/ollama.py index b541750d..ebd3b6bf 100644 --- a/nextcloud_mcp_server/providers/ollama.py +++ b/nextcloud_mcp_server/providers/ollama.py @@ -164,6 +164,13 @@ class OllamaProvider(Provider): 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)