Source code for lalandre_embedding.base

"""
Base class for embedding providers
"""

from abc import ABC, abstractmethod
from typing import List, Optional




class EmbeddingProvider(ABC):
    """Abstract base class for embedding providers"""



    @abstractmethod
    def embed_text(self, text: str) -> List[float]:
        """Generate embedding for a single text"""
        pass




    @abstractmethod
    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for multiple texts"""
        pass




    @abstractmethod
    def get_vector_size(self) -> int:
        """Get the dimension of the embedding vectors"""
        pass




    def estimate_tokens(self, text: str) -> Optional[int]:
        """Return a provider-native token count when available."""
        return None




    def get_max_input_tokens(self) -> Optional[int]:
        """Return the provider-native max input length when known."""
        return None






class SupportsCacheSize(ABC):
    """Protocol-like mixin for providers exposing cache occupancy."""



    @abstractmethod
    def get_cache_size(self) -> int:
        """Return the current number of cached embeddings."""
        ...






class SupportsNumKeys(ABC):
    """Protocol-like mixin for providers exposing API-key fan-out."""



    @abstractmethod
    def get_num_keys(self) -> int:
        """Return the number of configured upstream API keys."""
        ...