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) <noreply@anthropic.com>
53 lines
1.9 KiB
Python
53 lines
1.9 KiB
Python
"""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
|