emdash-core 0.1.7__py3-none-any.whl

emdash_core/embeddings/models.py

@@ -0,0 +1,192 @@
+"""Embedding models enum - single source of truth for all supported models."""
+
+from enum import Enum
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass(frozen=True)
+class ModelSpec:
+    """Specification for an embedding model."""
+
+    provider: str  # "openai", "fireworks"
+    model_id: str  # The actual model identifier for the API
+    dimensions: int  # Output embedding dimensions
+    max_tokens: int  # Max input tokens
+    batch_size: int  # Recommended batch size
+    description: str  # Human-readable description
+
+
+class EmbeddingModel(Enum):
+    """
+    All supported embedding models.
+
+    Format: PROVIDER_MODEL_NAME
+
+    Usage:
+        model = EmbeddingModel.OPENAI_TEXT_3_SMALL
+        print(model.spec.dimensions)  # 1536
+        print(model.spec.provider)    # "openai"
+    """
+
+    # ═══════════════════════════════════════════════════════════════════════════
+    # OpenAI Models
+    # ═══════════════════════════════════════════════════════════════════════════
+
+    OPENAI_TEXT_3_SMALL = ModelSpec(
+        provider="openai",
+        model_id="text-embedding-3-small",
+        dimensions=1536,
+        max_tokens=8191,
+        batch_size=100,
+        description="OpenAI's small, fast embedding model (best value)",
+    )
+
+    OPENAI_TEXT_3_LARGE = ModelSpec(
+        provider="openai",
+        model_id="text-embedding-3-large",
+        dimensions=3072,
+        max_tokens=8191,
+        batch_size=50,
+        description="OpenAI's large, high-quality embedding model",
+    )
+
+    OPENAI_ADA_002 = ModelSpec(
+        provider="openai",
+        model_id="text-embedding-ada-002",
+        dimensions=1536,
+        max_tokens=8191,
+        batch_size=100,
+        description="OpenAI's legacy Ada model (deprecated, use text-3-small)",
+    )
+
+    # ═══════════════════════════════════════════════════════════════════════════
+    # Fireworks AI Models
+    # ═══════════════════════════════════════════════════════════════════════════
+
+    FIREWORKS_NOMIC_EMBED_V1_5 = ModelSpec(
+        provider="fireworks",
+        model_id="nomic-ai/nomic-embed-text-v1.5",
+        dimensions=768,
+        max_tokens=8192,
+        batch_size=100,
+        description="Nomic's open-source embedding model (fast, good quality)",
+    )
+
+    FIREWORKS_E5_MISTRAL_7B = ModelSpec(
+        provider="fireworks",
+        model_id="intfloat/e5-mistral-7b-instruct",
+        dimensions=4096,
+        max_tokens=4096,
+        batch_size=20,
+        description="E5-Mistral 7B (highest quality, slower)",
+    )
+
+    FIREWORKS_UAE_LARGE_V1 = ModelSpec(
+        provider="fireworks",
+        model_id="WhereIsAI/UAE-Large-V1",
+        dimensions=1024,
+        max_tokens=512,
+        batch_size=50,
+        description="UAE-Large-V1 (good balance of speed/quality)",
+    )
+
+    FIREWORKS_GTE_LARGE = ModelSpec(
+        provider="fireworks",
+        model_id="thenlper/gte-large",
+        dimensions=1024,
+        max_tokens=512,
+        batch_size=50,
+        description="GTE-Large (Alibaba's efficient embedding model)",
+    )
+
+    FIREWORKS_BGE_LARGE_EN = ModelSpec(
+        provider="fireworks",
+        model_id="BAAI/bge-large-en-v1.5",
+        dimensions=1024,
+        max_tokens=512,
+        batch_size=50,
+        description="BGE-Large-EN (BAAI's high-quality English model)",
+    )
+
+    # ═══════════════════════════════════════════════════════════════════════════
+
+    @property
+    def spec(self) -> ModelSpec:
+        """Get the model specification."""
+        return self.value
+
+    @property
+    def provider(self) -> str:
+        """Shortcut to get provider name."""
+        return self.value.provider
+
+    @property
+    def model_id(self) -> str:
+        """Shortcut to get the API model ID."""
+        return self.value.model_id
+
+    @property
+    def dimensions(self) -> int:
+        """Shortcut to get embedding dimensions."""
+        return self.value.dimensions
+
+    @classmethod
+    def get_default(cls) -> "EmbeddingModel":
+        """Get the default embedding model."""
+        return cls.OPENAI_TEXT_3_SMALL
+
+    @classmethod
+    def from_string(cls, value: str) -> Optional["EmbeddingModel"]:
+        """
+        Parse model from string.
+
+        Accepts:
+            - Enum name: "OPENAI_TEXT_3_SMALL"
+            - Provider:model: "openai:text-embedding-3-small"
+            - Just model_id: "text-embedding-3-small"
+        """
+        value = value.strip()
+
+        # Try enum name first
+        try:
+            return cls[value.upper().replace("-", "_").replace(":", "_")]
+        except KeyError:
+            pass
+
+        # Try provider:model format
+        if ":" in value:
+            provider, model_id = value.split(":", 1)
+            for model in cls:
+                if model.provider == provider and model.model_id == model_id:
+                    return model
+
+        # Try just model_id
+        for model in cls:
+            if model.model_id == value:
+                return model
+
+        return None
+
+    @classmethod
+    def list_by_provider(cls, provider: str) -> list["EmbeddingModel"]:
+        """List all models for a specific provider."""
+        return [m for m in cls if m.provider == provider]
+
+    @classmethod
+    def list_all(cls) -> list[dict]:
+        """List all models with their specs for display."""
+        return [
+            {
+                "name": m.name,
+                "provider": m.provider,
+                "model_id": m.model_id,
+                "dimensions": m.dimensions,
+                "description": m.spec.description,
+            }
+            for m in cls
+        ]
+
+    def __str__(self) -> str:
+        """String representation as provider:model_id."""
+        return f"{self.provider}:{self.model_id}"

emdash_core/embeddings/providers/__init__.py

@@ -0,0 +1,7 @@
+"""Embedding providers package."""
+
+from .base import EmbeddingProvider
+from .openai import OpenAIProvider
+from .fireworks import FireworksProvider
+
+__all__ = ["EmbeddingProvider", "OpenAIProvider", "FireworksProvider"]

emdash_core/embeddings/providers/base.py

@@ -0,0 +1,112 @@
+"""Base class for embedding providers."""
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from ..models import EmbeddingModel
+
+
+class EmbeddingProvider(ABC):
+    """
+    Abstract base class for embedding providers.
+
+    Each provider (OpenAI, Fireworks, etc.) implements this interface.
+    The registry uses this to provide a unified embedding API.
+    """
+
+    def __init__(self, model: EmbeddingModel):
+        """
+        Initialize provider with a specific model.
+
+        Args:
+            model: The embedding model to use
+        """
+        self._model = model
+
+    @property
+    def model(self) -> EmbeddingModel:
+        """Get the embedding model."""
+        return self._model
+
+    @property
+    def dimensions(self) -> int:
+        """Get embedding dimensions for the current model."""
+        return self._model.dimensions
+
+    @property
+    @abstractmethod
+    def is_available(self) -> bool:
+        """Check if the provider is available (API key configured, etc.)."""
+        pass
+
+    @abstractmethod
+    def embed_texts(self, texts: list[str]) -> list[Optional[list[float]]]:
+        """
+        Generate embeddings for multiple texts.
+
+        Args:
+            texts: List of text strings to embed
+
+        Returns:
+            List of embedding vectors. None for failed embeddings.
+        """
+        pass
+
+    def embed_text(self, text: str) -> Optional[list[float]]:
+        """
+        Generate embedding for a single text.
+
+        Args:
+            text: Text string to embed
+
+        Returns:
+            Embedding vector or None if failed
+        """
+        if not text:
+            return None
+        embeddings = self.embed_texts([text])
+        return embeddings[0] if embeddings else None
+
+    def embed_query(self, query: str) -> Optional[list[float]]:
+        """
+        Generate embedding for a search query.
+
+        Some models treat queries differently from documents.
+        Override this method if the model requires special query handling.
+
+        Args:
+            query: Search query string
+
+        Returns:
+            Embedding vector or None if failed
+        """
+        return self.embed_text(query)
+
+    def _truncate_text(self, text: str, max_chars: int = 8000) -> str:
+        """
+        Truncate text to avoid token limits.
+
+        Args:
+            text: Text to truncate
+            max_chars: Maximum character length (roughly 4 chars per token)
+
+        Returns:
+            Truncated text
+        """
+        if text and len(text) > max_chars:
+            return text[:max_chars]
+        return text or ""
+
+    def _clean_batch(self, texts: list[str]) -> list[str]:
+        """
+        Clean and truncate a batch of texts.
+
+        Args:
+            texts: List of texts to clean
+
+        Returns:
+            Cleaned texts
+        """
+        # Calculate max chars based on model's max tokens (roughly 4 chars per token)
+        max_chars = min(self._model.spec.max_tokens * 4, 32000)
+        return [self._truncate_text(t, max_chars) for t in texts]

emdash_core/embeddings/providers/fireworks.py

@@ -0,0 +1,141 @@
+"""Fireworks AI embedding provider."""
+
+from typing import Optional
+
+from ..models import EmbeddingModel
+from .base import EmbeddingProvider
+from ...core.config import get_config
+from ...utils.logger import log
+
+
+class FireworksProvider(EmbeddingProvider):
+    """
+    Fireworks AI embedding provider.
+
+    Uses the Fireworks API (OpenAI-compatible) to generate embeddings.
+    Requires FIREWORKS_API_KEY environment variable.
+
+    API docs: https://docs.fireworks.ai/guides/querying-embeddings-models
+    """
+
+    # Fireworks API base URL
+    BASE_URL = "https://api.fireworks.ai/inference/v1"
+
+    # Models this provider handles
+    SUPPORTED_MODELS = {
+        EmbeddingModel.FIREWORKS_NOMIC_EMBED_V1_5,
+        EmbeddingModel.FIREWORKS_E5_MISTRAL_7B,
+        EmbeddingModel.FIREWORKS_UAE_LARGE_V1,
+        EmbeddingModel.FIREWORKS_GTE_LARGE,
+        EmbeddingModel.FIREWORKS_BGE_LARGE_EN,
+    }
+
+    def __init__(self, model: EmbeddingModel):
+        """
+        Initialize Fireworks provider.
+
+        Args:
+            model: The embedding model to use (must be a Fireworks model)
+        """
+        if model not in self.SUPPORTED_MODELS:
+            raise ValueError(f"Model {model} is not supported by FireworksProvider")
+        super().__init__(model)
+        self._client = None
+
+    @property
+    def _api_key(self) -> Optional[str]:
+        """Get Fireworks API key from config."""
+        config = get_config()
+        # Check if fireworks config exists
+        if hasattr(config, "fireworks") and config.fireworks.api_key:
+            return config.fireworks.api_key
+        return None
+
+    @property
+    def is_available(self) -> bool:
+        """Check if Fireworks API key is configured."""
+        return self._api_key is not None and len(self._api_key) > 0
+
+    @property
+    def _client_instance(self):
+        """Lazy-load OpenAI client configured for Fireworks."""
+        if self._client is None:
+            if not self.is_available:
+                raise RuntimeError(
+                    "Fireworks API key not configured. Set FIREWORKS_API_KEY environment variable."
+                )
+            try:
+                from openai import OpenAI
+
+                # Fireworks uses OpenAI-compatible API
+                self._client = OpenAI(
+                    api_key=self._api_key,
+                    base_url=self.BASE_URL,
+                )
+            except ImportError:
+                raise RuntimeError(
+                    "OpenAI library not installed. Run: pip install openai"
+                )
+        return self._client
+
+    def embed_texts(self, texts: list[str]) -> list[Optional[list[float]]]:
+        """
+        Generate embeddings using Fireworks API.
+
+        Args:
+            texts: List of text strings to embed
+
+        Returns:
+            List of embedding vectors. None for failed embeddings.
+        """
+        if not texts:
+            return []
+
+        all_embeddings = []
+        batch_size = self._model.spec.batch_size
+
+        for i in range(0, len(texts), batch_size):
+            batch = texts[i : i + batch_size]
+            cleaned_batch = self._clean_batch(batch)
+
+            try:
+                # Fireworks requires accounts/ prefix for model IDs
+                model_id = f"accounts/fireworks/models/{self._model.model_id}"
+
+                response = self._client_instance.embeddings.create(
+                    model=model_id,
+                    input=cleaned_batch,
+                )
+                batch_embeddings = [item.embedding for item in response.data]
+                all_embeddings.extend(batch_embeddings)
+
+            except Exception as e:
+                log.error(f"Fireworks embedding error: {e}")
+                all_embeddings.extend([None] * len(cleaned_batch))
+
+        return all_embeddings
+
+    def embed_query(self, query: str) -> Optional[list[float]]:
+        """
+        Generate embedding for a search query.
+
+        Some Fireworks models (like Nomic, E5) benefit from query prefixes.
+
+        Args:
+            query: Search query string
+
+        Returns:
+            Embedding vector or None if failed
+        """
+        if not query:
+            return None
+
+        # E5 models expect "query: " prefix for queries
+        if self._model == EmbeddingModel.FIREWORKS_E5_MISTRAL_7B:
+            query = f"query: {query}"
+
+        # Nomic models can optionally use "search_query: " prefix
+        elif self._model == EmbeddingModel.FIREWORKS_NOMIC_EMBED_V1_5:
+            query = f"search_query: {query}"
+
+        return self.embed_text(query)

emdash_core/embeddings/providers/openai.py

@@ -0,0 +1,104 @@
+"""OpenAI embedding provider."""
+
+from typing import Optional
+
+from ..models import EmbeddingModel
+from .base import EmbeddingProvider
+from ...core.config import get_config
+from ...utils.logger import log
+
+
+class OpenAIProvider(EmbeddingProvider):
+    """
+    OpenAI embedding provider.
+
+    Uses the OpenAI API to generate embeddings.
+    Requires OPENAI_API_KEY environment variable.
+    """
+
+    # Models this provider handles
+    SUPPORTED_MODELS = {
+        EmbeddingModel.OPENAI_TEXT_3_SMALL,
+        EmbeddingModel.OPENAI_TEXT_3_LARGE,
+        EmbeddingModel.OPENAI_ADA_002,
+    }
+
+    def __init__(self, model: EmbeddingModel):
+        """
+        Initialize OpenAI provider.
+
+        Args:
+            model: The embedding model to use (must be an OpenAI model)
+        """
+        if model not in self.SUPPORTED_MODELS:
+            raise ValueError(f"Model {model} is not supported by OpenAIProvider")
+        super().__init__(model)
+        self._client = None
+        self._config = get_config().openai
+
+    @property
+    def is_available(self) -> bool:
+        """Check if OpenAI API key is configured."""
+        return self._config.is_available
+
+    @property
+    def _client_instance(self):
+        """Lazy-load OpenAI client."""
+        if self._client is None:
+            if not self.is_available:
+                raise RuntimeError(
+                    "OpenAI API key not configured. Set OPENAI_API_KEY environment variable."
+                )
+            try:
+                from openai import OpenAI
+
+                self._client = OpenAI(api_key=self._config.api_key)
+            except ImportError:
+                raise RuntimeError(
+                    "OpenAI library not installed. Run: pip install openai"
+                )
+        return self._client
+
+    def embed_texts(self, texts: list[str]) -> list[Optional[list[float]]]:
+        """
+        Generate embeddings using OpenAI API.
+
+        Args:
+            texts: List of text strings to embed
+
+        Returns:
+            List of embedding vectors. None for failed embeddings.
+        """
+        if not texts:
+            return []
+
+        all_embeddings = []
+        batch_size = self._model.spec.batch_size
+
+        for i in range(0, len(texts), batch_size):
+            batch = texts[i : i + batch_size]
+            cleaned_batch = self._clean_batch(batch)
+
+            try:
+                # Use dimensions parameter for text-embedding-3 models
+                kwargs = {
+                    "model": self._model.model_id,
+                    "input": cleaned_batch,
+                }
+
+                # text-embedding-3 models support custom dimensions
+                if self._model in {
+                    EmbeddingModel.OPENAI_TEXT_3_SMALL,
+                    EmbeddingModel.OPENAI_TEXT_3_LARGE,
+                }:
+                    kwargs["dimensions"] = self._model.dimensions
+
+                response = self._client_instance.embeddings.create(**kwargs)
+                batch_embeddings = [item.embedding for item in response.data]
+                all_embeddings.extend(batch_embeddings)
+
+            except Exception as e:
+                log.error(f"OpenAI embedding error: {e}")
+                all_embeddings.extend([None] * len(cleaned_batch))
+
+        return all_embeddings

emdash_core/embeddings/registry.py

@@ -0,0 +1,146 @@
+"""Provider registry for embedding models."""
+
+from typing import Type, Optional
+
+from .models import EmbeddingModel
+from .providers.base import EmbeddingProvider
+from .providers.openai import OpenAIProvider
+from .providers.fireworks import FireworksProvider
+
+
+class ProviderRegistry:
+    """
+    Registry for embedding providers.
+
+    Maps provider names to provider classes. No if-else chains needed.
+    Just register your provider once and it's available everywhere.
+
+    Usage:
+        registry = ProviderRegistry()
+        provider = registry.get_provider(EmbeddingModel.OPENAI_TEXT_3_SMALL)
+        embeddings = provider.embed_texts(["hello world"])
+    """
+
+    # Provider class registry: provider_name -> provider_class
+    _providers: dict[str, Type[EmbeddingProvider]] = {}
+
+    @classmethod
+    def register(cls, provider_name: str, provider_class: Type[EmbeddingProvider]):
+        """
+        Register a provider class.
+
+        Args:
+            provider_name: Name of the provider (e.g., "openai", "fireworks")
+            provider_class: The provider class to register
+        """
+        cls._providers[provider_name] = provider_class
+
+    @classmethod
+    def get_provider_class(cls, provider_name: str) -> Optional[Type[EmbeddingProvider]]:
+        """
+        Get the provider class for a provider name.
+
+        Args:
+            provider_name: Name of the provider
+
+        Returns:
+            Provider class or None if not registered
+        """
+        return cls._providers.get(provider_name)
+
+    @classmethod
+    def get_provider(cls, model: EmbeddingModel) -> EmbeddingProvider:
+        """
+        Get an instantiated provider for a model.
+
+        Args:
+            model: The embedding model
+
+        Returns:
+            Instantiated provider for the model
+
+        Raises:
+            ValueError: If no provider is registered for the model's provider
+        """
+        provider_class = cls._providers.get(model.provider)
+        if provider_class is None:
+            raise ValueError(
+                f"No provider registered for '{model.provider}'. "
+                f"Available providers: {list(cls._providers.keys())}"
+            )
+        return provider_class(model)
+
+    @classmethod
+    def list_providers(cls) -> list[str]:
+        """List all registered provider names."""
+        return list(cls._providers.keys())
+
+    @classmethod
+    def is_provider_available(cls, provider_name: str) -> bool:
+        """
+        Check if a provider is available (registered and configured).
+
+        Args:
+            provider_name: Name of the provider
+
+        Returns:
+            True if provider is registered and has valid credentials
+        """
+        provider_class = cls._providers.get(provider_name)
+        if provider_class is None:
+            return False
+
+        # Get any model for this provider to check availability
+        models = EmbeddingModel.list_by_provider(provider_name)
+        if not models:
+            return False
+
+        try:
+            provider = provider_class(models[0])
+            return provider.is_available
+        except Exception:
+            return False
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Register all providers
+# ═══════════════════════════════════════════════════════════════════════════════
+
+ProviderRegistry.register("openai", OpenAIProvider)
+ProviderRegistry.register("fireworks", FireworksProvider)
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Convenience functions
+# ═══════════════════════════════════════════════════════════════════════════════
+
+
+def get_provider(model: EmbeddingModel) -> EmbeddingProvider:
+    """Get an instantiated provider for a model."""
+    return ProviderRegistry.get_provider(model)
+
+
+def get_default_provider() -> EmbeddingProvider:
+    """Get the default embedding provider (OpenAI text-embedding-3-small)."""
+    return ProviderRegistry.get_provider(EmbeddingModel.get_default())
+
+
+def get_available_model() -> Optional[EmbeddingModel]:
+    """
+    Get the first available model (has valid API credentials).
+
+    Checks OpenAI first, then Fireworks.
+
+    Returns:
+        First available model or None if no providers are configured
+    """
+    # Priority order
+    priority = ["openai", "fireworks"]
+
+    for provider_name in priority:
+        if ProviderRegistry.is_provider_available(provider_name):
+            models = EmbeddingModel.list_by_provider(provider_name)
+            if models:
+                return models[0]
+
+    return None