yera 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

yera/config2/dataclasses.py

@@ -0,0 +1,534 @@
+"""Data classes for Yera's config."""
+
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+MODEL_TYPES = ["llm", "embedding", "reranker", "image", "moderation", "tts", "stt"]
+
+# ============================================================================
+# Shared Components
+# ============================================================================
+
+
+class ModelProperties(BaseModel):
+    """Technical properties of a model."""
+
+    remote: bool = True
+    context_window: int | None = None  # LLMs only
+
+    model_config = ConfigDict(extra="forbid")
+
+
+# ============================================================================
+# Global Defaults
+# ============================================================================
+
+
+class GlobalDefaults(BaseModel):
+    """Global default settings for Yera."""
+
+    max_retries: int = 3
+    timeout_seconds: int = 60
+    log_level: Literal["debug", "info", "warn", "error"] = "info"
+    enable_telemetry: bool = False
+
+    model_config = ConfigDict(extra="forbid")
+
+
+# ============================================================================
+# LLM Models
+# ============================================================================
+
+
+class LLMCapabilities(BaseModel):
+    """Capabilities specific to LLM models."""
+
+    function_calling: bool = False
+    structured_generation: bool = False
+    vision: bool = False
+    reasoning: bool = False
+    multilingual: bool = False
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class LLMCost(BaseModel):
+    """Cost structure for LLM models."""
+
+    input_tokens_per_million: float
+    output_tokens_per_million: float
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class LLMInference(BaseModel):
+    """Runtime inference parameters for LLM models."""
+
+    temperature: float | None = None
+    max_tokens: int | None = None
+    top_p: float | None = None
+    top_k: int | None = None
+    stop_sequences: list[str] = Field(default_factory=list)
+
+    model_config = ConfigDict(extra="allow")
+
+
+class LLMConfig(BaseModel):
+    """Configuration for a Language Model."""
+
+    id: str  # Full identifier like "anthropic.claude-sonnet-4"
+    display_name: str
+    credentials: str
+    provider: str
+    interface: str
+    model_id: str
+    properties: ModelProperties = Field(default_factory=ModelProperties)
+    cost: LLMCost | None = None
+    capabilities: LLMCapabilities = Field(default_factory=LLMCapabilities)
+    inference: LLMInference = Field(default_factory=LLMInference)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+# ============================================================================
+# Embedding Models
+# ============================================================================
+
+
+class EmbeddingCapabilities(BaseModel):
+    """Capabilities specific to embedding models."""
+
+    cosine_similarity: bool = True
+    dot_product: bool = True
+    euclidean_distance: bool = False
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class EmbeddingCost(BaseModel):
+    """Cost structure for embedding models."""
+
+    tokens_per_million: float
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class EmbeddingInference(BaseModel):
+    """Runtime inference parameters for embedding models."""
+
+    batch_size: int | None = None
+    dimensions: int | None = None  # For models that support configurable dimensions
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class EmbeddingModelConfig(BaseModel):
+    """Configuration for an embedding model."""
+
+    id: str  # Full identifier like "openai.text-embedding-3-small"
+    credentials: str
+    provider: str
+    interface: str
+    model_id: str
+    dimensions: int
+    max_input_tokens: int
+    status: Literal["stable", "beta", "deprecated"] = "stable"
+    properties: ModelProperties
+    cost: EmbeddingCost
+    capabilities: EmbeddingCapabilities = Field(default_factory=EmbeddingCapabilities)
+    inference: EmbeddingInference = Field(default_factory=EmbeddingInference)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+# ============================================================================
+# Reranker Models
+# ============================================================================
+
+
+class RerankerCapabilities(BaseModel):
+    """Capabilities specific to reranker models."""
+
+    multilingual: bool = False
+    context_aware: bool = True
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class RerankerCost(BaseModel):
+    """Cost structure for reranker models."""
+
+    per_thousand_searches: float
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class RerankerInference(BaseModel):
+    """Runtime inference parameters for reranker models."""
+
+    top_n: int | None = None
+    return_documents: bool = True
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class RerankerModelConfig(BaseModel):
+    """Configuration for a reranker model."""
+
+    id: str  # Full identifier like "cohere.rerank-english-v3"
+    credentials: str
+    provider: str
+    interface: str
+    model_id: str
+    max_documents: int
+    max_query_length: int
+    status: Literal["stable", "beta", "deprecated"] = "stable"
+    properties: ModelProperties
+    cost: RerankerCost
+    capabilities: RerankerCapabilities = Field(default_factory=RerankerCapabilities)
+    inference: RerankerInference = Field(default_factory=RerankerInference)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+# ============================================================================
+# Image Generation Models
+# ============================================================================
+
+
+class ImageCapabilities(BaseModel):
+    """Capabilities specific to image generation models."""
+
+    supports_variations: bool = False
+    supports_edits: bool = False
+    supports_transparent_background: bool = False
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class ImageCost(BaseModel):
+    """Cost structure for image generation models."""
+
+    per_image_standard: float | None = None
+    per_image_hd: float | None = None
+    per_image_1024: float | None = None  # DALL-E 2 pricing
+    per_image_512: float | None = None
+    per_image_256: float | None = None
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class ImageInference(BaseModel):
+    """Runtime inference parameters for image generation models."""
+
+    size: str | None = None
+    quality: str | None = None
+    style: str | None = None
+    n: int = 1  # Number of images to generate
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class ImageModelConfig(BaseModel):
+    """Configuration for an image generation model."""
+
+    id: str  # Full identifier like "openai.dall-e-3"
+    credentials: str
+    provider: str
+    interface: str
+    model_id: str
+    supported_sizes: list[str]
+    supported_qualities: list[str]
+    supported_styles: list[str] = Field(default_factory=list)
+    status: Literal["stable", "beta", "deprecated"] = "stable"
+    properties: ModelProperties
+    cost: ImageCost
+    capabilities: ImageCapabilities = Field(default_factory=ImageCapabilities)
+    inference: ImageInference = Field(default_factory=ImageInference)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+# ============================================================================
+# Moderation Models
+# ============================================================================
+
+
+class ModerationCapabilities(BaseModel):
+    """Capabilities specific to moderation models."""
+
+    categories: list[str] = Field(default_factory=list)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class ModerationCost(BaseModel):
+    """Cost structure for moderation models."""
+
+    free: bool = True
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class ModerationInference(BaseModel):
+    """Runtime inference parameters for moderation models."""
+
+    # Most moderation models have no configurable inference params
+    model_config = ConfigDict(extra="forbid")
+
+
+class ModerationModelConfig(BaseModel):
+    """Configuration for a moderation model."""
+
+    id: str  # Full identifier like "openai.text-moderation-latest"
+    credentials: str
+    provider: str
+    interface: str
+    model_id: str
+    status: Literal["stable", "beta", "deprecated"] = "stable"
+    properties: ModelProperties
+    cost: ModerationCost
+    capabilities: ModerationCapabilities = Field(default_factory=ModerationCapabilities)
+    inference: ModerationInference = Field(default_factory=ModerationInference)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+# ============================================================================
+# Text-to-Speech Models
+# ============================================================================
+
+
+class TTSCapabilities(BaseModel):
+    """Capabilities specific to TTS models."""
+
+    streaming: bool = False
+    speed_control: bool = False
+    voice_cloning: bool = False
+    emotion_control: bool = False
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class TTSCost(BaseModel):
+    """Cost structure for TTS models."""
+
+    chars_per_million: float
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class TTSInference(BaseModel):
+    """Runtime inference parameters for TTS models."""
+
+    voice: str | None = None
+    speed: float | None = None
+    format: str | None = None  # mp3, opus, etc.
+    sample_rate: int | None = None
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class TTSModelConfig(BaseModel):
+    """Configuration for a text-to-speech model."""
+
+    id: str  # Full identifier like "openai.tts-1"
+    credentials: str
+    provider: str
+    interface: str
+    model_id: str
+    voices: list[str] = Field(default_factory=list)
+    default_voice: str = ""
+    supported_formats: list[str] = Field(default_factory=list)
+    sample_rates: list[int] = Field(default_factory=list)
+    status: Literal["stable", "beta", "deprecated"] = "stable"
+    properties: ModelProperties
+    cost: TTSCost
+    capabilities: TTSCapabilities = Field(default_factory=TTSCapabilities)
+    inference: TTSInference = Field(default_factory=TTSInference)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+# ============================================================================
+# Speech-to-Text Models
+# ============================================================================
+
+
+class STTCapabilities(BaseModel):
+    """Capabilities specific to STT models."""
+
+    timestamps: bool = False
+    translation: bool = False
+    language_detection: bool = False
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class STTCost(BaseModel):
+    """Cost structure for STT models."""
+
+    per_minute: float
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class STTInference(BaseModel):
+    """Runtime inference parameters for STT models."""
+
+    language: str | None = None
+    enable_timestamps: bool = False
+    response_format: str | None = None  # json, text, srt, vtt
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class STTModelConfig(BaseModel):
+    """Configuration for a speech-to-text model."""
+
+    id: str  # Full identifier like "openai.whisper-1"
+    credentials: str
+    provider: str
+    interface: str
+    model_id: str
+    supported_formats: list[str] = Field(default_factory=list)
+    max_file_size_mb: float = 25.0
+    languages: list[str] = Field(default_factory=list)
+    status: Literal["stable", "beta", "deprecated"] = "stable"
+    properties: ModelProperties
+    cost: STTCost
+    capabilities: STTCapabilities = Field(default_factory=STTCapabilities)
+    inference: STTInference = Field(default_factory=STTInference)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+# ============================================================================
+# Model Registry
+# ============================================================================
+
+
+class ModelRegistry(BaseModel):
+    """Registry of all configured models by type.
+
+    Models are keyed by their full dot-path identifier (e.g.,
+    "anthropic.claude-sonnet-4").
+
+    """
+
+    llm: dict[str, LLMConfig] = Field(default_factory=dict)
+    embedding: dict[str, EmbeddingModelConfig] = Field(default_factory=dict)
+    reranker: dict[str, RerankerModelConfig] = Field(default_factory=dict)
+    image: dict[str, ImageModelConfig] = Field(default_factory=dict)
+    moderation: dict[str, ModerationModelConfig] = Field(default_factory=dict)
+    tts: dict[str, TTSModelConfig] = Field(default_factory=dict)
+    stt: dict[str, STTModelConfig] = Field(default_factory=dict)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+# ============================================================================
+# Top-Level Configuration
+# ============================================================================
+
+
+class DefaultModels(BaseModel):
+    """Default model selections by type."""
+
+    llm: str | None = None
+    embedding: str | None = None
+    reranker: str | None = None
+    image: str | None = None
+    tts: str | None = None
+    stt: str | None = None
+    moderation: str | None = None
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class Defaults(BaseModel):
+    """Yera configuration defaults section."""
+
+    max_retries: int = 3
+    timeout_seconds: int = 60
+    log_level: Literal["debug", "info", "warn", "error"] = "info"
+    enable_telemetry: bool = False
+
+    models: DefaultModels = Field(default_factory=DefaultModels)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class CredentialsMap(BaseModel):
+    """Credentials for Yera configuration."""
+
+    providers: dict[str, dict[str, str]] = Field(default_factory=dict)
+    tools: dict[str, dict[str, str]] = Field(default_factory=dict)
+
+
+class YeraConfig(BaseModel):
+    """Complete Yera configuration.
+
+    This is the root configuration object that contains all provider credentials,
+    model definitions, and default selections.
+    """
+
+    defaults: Defaults = Field(default_factory=Defaults)
+    credentials: CredentialsMap = Field(default_factory=CredentialsMap)
+    models: ModelRegistry = Field(default_factory=ModelRegistry)
+
+    model_config = ConfigDict(extra="forbid")
+
+    def merge(self, other: "YeraConfig") -> "YeraConfig":
+        """Merge this configuration with another, creating a new config.
+
+        Creates a new YeraConfig by combining this configuration with another,
+        where values from `other` take precedence in case of conflicts. This is
+        for extending and updating yera config.
+
+        Args:
+            other: Another YeraConfig to merge with this one. Values from this
+                config will override values from self when keys conflict.
+
+        Returns:
+            A new YeraConfig instance containing the merged configuration. The
+            original configs (self and other) are not modified.
+
+        """
+        merged_defaults = {
+            **self.defaults.model_dump(),
+            **other.defaults.model_dump(),
+        }
+        merged_default_models = {
+            **self.defaults.models.model_dump(),
+            **other.defaults.models.model_dump(),
+        }
+        defaults = self.defaults.model_validate(
+            {
+                **merged_defaults,
+                "models": merged_default_models,
+            }
+        )
+
+        providers_merge = {**self.credentials.providers, **other.credentials.providers}
+        tools_merge = {**self.credentials.tools, **other.credentials.tools}
+        credentials = self.credentials.model_copy(
+            update={
+                "providers": providers_merge,
+                "tools": tools_merge,
+            }
+        )
+
+        model_reg1 = self.models.model_dump()
+        model_reg2 = other.models.model_dump()
+        models_merge = {}
+        for k in model_reg1:
+            models_merge[k] = {**model_reg1[k], **model_reg2[k]}
+        models = ModelRegistry(**models_merge)
+
+        return YeraConfig(
+            defaults=defaults,
+            credentials=credentials,
+            models=models,
+        )