smart-org/apps/api/app/domain/ai_control/models.py

"""AI Control Plane domain: providers, credentials, models, routing, rate cards.

Everything here is platform-owner-only. A normal customer never receives these
records through the API — the application layer strips them.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from enum import StrEnum

from app.domain.base import new_id, utcnow


class ProviderType(StrEnum):
    MOCK = "mock"
    OPENAI_COMPATIBLE = "openai_compatible"
    ANTHROPIC = "anthropic"
    GEMINI = "gemini"
    OPENROUTER = "openrouter"
    #: Google Antigravity (``agy``) — OAuth accounts billed through a Google AI
    #: subscription. Credentials of this type are *accounts*, not API keys.
    AGY = "agy"


class CredentialStatus(StrEnum):
    """Lifecycle of one agy account in the rotation pool."""

    ACTIVE = "active"
    #: Rate limited by Google; skipped until ``cooldown_until``.
    COOLDOWN = "cooldown"
    #: Refresh token rejected or another hard error; needs re-login.
    ERROR = "error"


class RoutingMode(StrEnum):
    FAST = "fast"
    CHEAP = "cheap"
    BALANCED = "balanced"
    REASONING = "reasoning"
    VISION = "vision"
    CODING = "coding"


class Capability(StrEnum):
    TEXT = "text"
    VISION = "vision"
    TOOLS = "tools"
    REASONING = "reasoning"
    LONG_CONTEXT = "long_context"
    CODE = "code"


class RateUnit(StrEnum):
    MODEL_INPUT_TOKENS = "model_input_tokens"
    MODEL_OUTPUT_TOKENS = "model_output_tokens"
    SANDBOX_CPU_SECOND = "sandbox_cpu_second"
    SANDBOX_MEMORY_GB_SECOND = "sandbox_memory_gb_second"
    STORAGE_GB_DAY = "storage_gb_day"
    TOOL_CALL = "tool_call"
    PREMIUM_EMPLOYEE_RUN = "premium_employee_run"


@dataclass(slots=True, kw_only=True)
class Provider:
    name: str
    type: ProviderType = ProviderType.MOCK
    base_url: str = ""
    id: str = field(default_factory=new_id)
    is_active: bool = True
    created_at: datetime = field(default_factory=utcnow)


@dataclass(slots=True, kw_only=True)
class ModelCredential:
    """One credential for a provider.

    For key-based providers ``encrypted_api_key`` holds the encrypted API key and
    the agy fields stay empty. For the agy provider the credential *is an
    account*: ``encrypted_api_key`` holds the encrypted JSON blob
    ``{"refresh_token": ..., "project_id": ..., "email": ...}`` while ``email`` /
    ``project_id`` duplicate the identity columns so the admin console can show
    and rotate accounts without touching the ciphertext.
    """

    provider_id: str
    name: str
    encrypted_api_key: str
    id: str = field(default_factory=new_id)
    is_active: bool = True
    created_at: datetime = field(default_factory=utcnow)
    #: agy account identity (empty for API-key credentials).
    email: str = ""
    project_id: str = ""
    status: str = CredentialStatus.ACTIVE.value
    cooldown_until: datetime | None = None
    last_error: str = ""
    last_used_at: datetime | None = None


@dataclass(slots=True, kw_only=True)
class AiModel:
    provider_id: str
    internal_name: str
    display_name_for_admin: str
    capabilities: list[str] = field(default_factory=list)
    input_price_per_1m: float = 0.0
    output_price_per_1m: float = 0.0
    id: str = field(default_factory=new_id)
    is_active: bool = True
    created_at: datetime = field(default_factory=utcnow)

    def supports(self, capability: str | Capability) -> bool:
        return str(capability) in set(self.capabilities)


@dataclass(slots=True, kw_only=True)
class RoutingPolicy:
    name: str
    mode: RoutingMode = RoutingMode.BALANCED
    description: str = ""
    fallback_policy_json: dict = field(default_factory=dict)
    budget_limits_json: dict = field(default_factory=dict)
    id: str = field(default_factory=new_id)
    is_default: bool = False
    is_active: bool = True
    created_at: datetime = field(default_factory=utcnow)


@dataclass(slots=True, kw_only=True)
class RateCard:
    name: str
    unit: RateUnit
    customer_price: float
    internal_cost_estimate: float = 0.0
    id: str = field(default_factory=new_id)
    is_active: bool = True
    created_at: datetime = field(default_factory=utcnow)