From a0bb5642cb57f97c818cfff5ea5a73af11bcae36 Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Mon, 8 Jun 2026 01:07:22 +0200 Subject: [PATCH] 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