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