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>
127 lines
3.7 KiB
Python
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
|