Files
mcp-nextcloud/nextcloud_mcp_server/providers/base.py
T
Chris CoutinhoandClaude Opus 4.8 64318f0b25 feat(usage): meter embedding tokens as embeddings_queries on both paths
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>
2026-06-08 00:53:58 +02:00

127 lines
3.7 KiB
Python

"""Unified provider interface for embeddings and text generation."""
import math
from abc import ABC, abstractmethod
class Provider(ABC):
"""
Unified base class for LLM providers.
Providers can support embeddings, text generation, or both.
Use capability properties to determine what features are available.
"""
@property
@abstractmethod
def supports_embeddings(self) -> bool:
"""Whether this provider supports embedding generation."""
pass
@property
@abstractmethod
def supports_generation(self) -> bool:
"""Whether this provider supports text generation."""
pass
@abstractmethod
async def embed(self, text: str) -> list[float]:
"""
Generate embedding vector for text.
Args:
text: Input text to embed
Returns:
Vector embedding as list of floats
Raises:
NotImplementedError: If provider doesn't support embeddings
"""
pass
@abstractmethod
async def embed_batch(self, texts: list[str]) -> list[list[float]]:
"""
Generate embeddings for multiple texts (optimized).
Args:
texts: List of texts to embed
Returns:
List of vector embeddings
Raises:
NotImplementedError: If provider doesn't support embeddings
"""
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:
"""
Get embedding dimension for this provider.
Returns:
Vector dimension (e.g., 768 for nomic-embed-text)
Raises:
NotImplementedError: If provider doesn't support embeddings
"""
pass
@abstractmethod
async def generate(self, prompt: str, max_tokens: int = 500) -> str:
"""
Generate text from a prompt.
Args:
prompt: The prompt to generate from
max_tokens: Maximum tokens to generate
Returns:
Generated text
Raises:
NotImplementedError: If provider doesn't support generation
"""
pass
@abstractmethod
async def close(self) -> None:
"""Close the provider and release resources."""
pass