routstr-core/routstr/upstream/helpers.py

from __future__ import annotations

import asyncio
import os
import re
from typing import TYPE_CHECKING, Callable

if TYPE_CHECKING:
    from ..core.settings import Settings

from sqlmodel import select

from ..core import get_logger
from ..core.db import AsyncSession, ModelRow, UpstreamProviderRow, create_session
from ..core.provider_slugs import allocate_unique_provider_slug
from ..payment.models import Model
from .base import BaseUpstreamProvider

logger = get_logger(__name__)


def resolve_model_alias(
    model_id: str, canonical_slug: str | None = None, alias_ids: list[str] | None = None
) -> list[str]:
    """Resolve model ID to all possible aliases.

    Returns list of aliases including canonical slug and variations without provider prefix.

    Args:
        model_id: Model identifier (e.g., "gpt-5-mini" or "openai/gpt-5-mini")
        canonical_slug: Optional canonical slug from provider (e.g., "openai/gpt-5-pro-2025-10-06")

    Returns:
        List of possible model ID aliases
    """
    aliases = [model_id]

    base_model = model_id
    if "/" in model_id:
        without_prefix = model_id.split("/", 1)[1]
        aliases.append(without_prefix)
        base_model = without_prefix

    date_pattern = re.compile(r"-\d{4}-\d{2}-\d{2}$")
    if date_pattern.search(base_model):
        base_without_date = date_pattern.sub("", base_model)
        if base_without_date not in aliases:
            aliases.append(base_without_date)
        if "/" in model_id:
            prefix = model_id.split("/", 1)[0]
            prefixed_without_date = f"{prefix}/{base_without_date}"
            if prefixed_without_date not in aliases:
                aliases.append(prefixed_without_date)

    if canonical_slug and canonical_slug not in aliases:
        aliases.append(canonical_slug)
        if "/" in canonical_slug:
            canonical_without_prefix = canonical_slug.split("/", 1)[1]
            if canonical_without_prefix not in aliases:
                aliases.append(canonical_without_prefix)
            if date_pattern.search(canonical_without_prefix):
                canonical_base = date_pattern.sub("", canonical_without_prefix)
                if canonical_base not in aliases:
                    aliases.append(canonical_base)

    if alias_ids:
        aliases.extend(alias_ids)

    return aliases


async def get_all_models_with_overrides(
    upstreams: list[BaseUpstreamProvider],
) -> list[Model]:
    """Get all models from all providers with database overrides applied.

    Models in the database with upstream_provider_id set are treated as overrides
    that replace the provider's model with the same ID.

    Args:
        upstreams: List of upstream provider instances

    Returns:
        List of Model objects with overrides applied
    """
    from sqlmodel import select

    from ..payment.models import _row_to_model

    async with create_session() as session:
        result = await session.exec(select(ModelRow).where(ModelRow.enabled))
        override_rows = result.all()

        provider_result = await session.exec(select(UpstreamProviderRow))
        providers_by_id = {p.id: p for p in provider_result.all()}

        overrides_by_key: dict[tuple[str, int], tuple[ModelRow, float]] = {
            (row.id.lower(), row.upstream_provider_id): (
                row,
                providers_by_id[row.upstream_provider_id].provider_fee,
            )
            for row in override_rows
            if row.upstream_provider_id is not None
            and row.upstream_provider_id in providers_by_id
            and providers_by_id[row.upstream_provider_id].enabled
        }

    all_models: dict[tuple[str, str], Model] = {}

    for upstream in upstreams:
        upstream_db_id = getattr(upstream, "db_id", None)
        provider_key = (
            f"db:{upstream_db_id}"
            if isinstance(upstream_db_id, int)
            else f"{getattr(upstream, 'provider_type', '')}|{getattr(upstream, 'base_url', '')}"
        )
        for model in upstream.get_cached_models():
            model_key = (
                (model.id.lower(), upstream_db_id)
                if isinstance(upstream_db_id, int)
                else None
            )
            if model_key is not None and model_key in overrides_by_key:
                override_row, provider_fee = overrides_by_key[model_key]
                all_models[(model.id.lower(), provider_key)] = _row_to_model(
                    override_row, apply_provider_fee=True, provider_fee=provider_fee
                )
            elif model.enabled:
                all_models[(model.id.lower(), provider_key)] = model

    return list(all_models.values())


async def refresh_upstreams_models_periodically(
    upstreams_provider: (
        Callable[[], list[BaseUpstreamProvider]] | list[BaseUpstreamProvider]
    ),
) -> None:
    """Background task to periodically refresh models cache for all providers.

    Args:
        upstreams_provider: Either a callable returning the live upstream list
            (preferred — picks up providers added/changed via reinitialize_upstreams),
            or a static list (legacy, will go stale after reinitialize_upstreams).
    """
    import asyncio
    import random

    from ..core.settings import settings

    interval = getattr(settings, "models_refresh_interval_seconds", 0)
    if not interval or interval <= 0:
        logger.info("Provider models refresh disabled (interval <= 0)")
        return

    def _resolve_upstreams() -> list[BaseUpstreamProvider]:
        if callable(upstreams_provider):
            return upstreams_provider()
        return upstreams_provider

    while True:
        try:
            for upstream in _resolve_upstreams():
                try:
                    await upstream.refresh_models_cache()
                except Exception as e:
                    logger.error(
                        f"Error refreshing models for {upstream.base_url}",
                        extra={"error": str(e), "error_type": type(e).__name__},
                    )

            try:
                from ..payment.models import _update_sats_pricing_once

                await _update_sats_pricing_once()
            except Exception as e:
                logger.warning(f"Failed to update pricing after model refresh: {e}")
                from ..proxy import refresh_model_maps

                await refresh_model_maps()

        except asyncio.CancelledError:
            break
        except Exception as e:
            logger.error(
                "Error in provider models refresh loop",
                extra={"error": str(e), "error_type": type(e).__name__},
            )

        try:
            jitter = max(0.0, float(interval) * 0.1)
            await asyncio.sleep(interval + random.uniform(0, jitter))
        except asyncio.CancelledError:
            break


async def init_upstreams() -> list[BaseUpstreamProvider]:
    """Initialize upstream providers from database.

    Seeds database with providers from settings if empty, then loads and instantiates
    provider instances from database records, and refreshes their models cache.
    """
    from ..core.settings import settings

    async with create_session() as session:
        result = await session.exec(select(UpstreamProviderRow))
        existing_providers = result.all()

        if not existing_providers:
            await _seed_providers_from_settings(session, settings)
            await session.commit()
            result = await session.exec(select(UpstreamProviderRow))
            existing_providers = result.all()
            if existing_providers:
                logger.info(
                    f"Seeded {len(existing_providers)} upstream providers from settings"
                )

        async def _init_single_provider(
            provider_row: UpstreamProviderRow,
        ) -> BaseUpstreamProvider | None:
            if not provider_row.enabled:
                logger.debug(f"Skipping disabled provider: {provider_row.base_url}")
                return None

            provider = _instantiate_provider(provider_row)
            if provider:
                await provider.refresh_models_cache()
                logger.debug(
                    f"Initialized {provider_row.provider_type} provider",
                    extra={
                        "base_url": provider_row.base_url,
                        "models_cached": len(provider.get_cached_models()),
                    },
                )
                return provider
            return None

        tasks = [_init_single_provider(row) for row in existing_providers]
        results = await asyncio.gather(*tasks)
        upstreams = [p for p in results if p is not None]

        return upstreams


async def _seed_providers_from_settings(
    session: AsyncSession, settings: "Settings"
) -> None:
    """Seed database with upstream providers from environment variables.

    Args:
        session: Database session
    """
    from sqlmodel import select

    from . import upstream_provider_classes

    providers_to_add: list[UpstreamProviderRow] = []
    seeded_provider_keys: set[tuple[str, str]] = set()
    reserved_slugs: set[str] = set()

    provider_classes_by_type = {
        cls.provider_type: cls
        for cls in upstream_provider_classes  # type: ignore[attr-defined]
    }

    env_mappings: list[tuple[str, str, str | None, str | None]] = [
        ("OPENAI_API_KEY", "openai", None, None),
        ("ANTHROPIC_API_KEY", "anthropic", None, None),
        ("OPENROUTER_API_KEY", "openrouter", None, None),
        ("GROQ_API_KEY", "groq", None, None),
        ("PERPLEXITY_API_KEY", "perplexity", None, None),
        ("FIREWORKS_API_KEY", "fireworks", None, None),
        ("XAI_API_KEY", "xai", None, None),
        ("TINFOIL_API_KEY", "tinfoil", None, None),
    ]

    for env_key, provider_type, _, _ in env_mappings:
        api_key = os.environ.get(env_key)
        if api_key and provider_type in provider_classes_by_type:
            provider_class = provider_classes_by_type[provider_type]
            if provider_class.default_base_url:  # type: ignore[attr-defined]
                base_url = provider_class.default_base_url  # type: ignore[attr-defined]
                result = await session.exec(
                    select(UpstreamProviderRow).where(
                        UpstreamProviderRow.base_url == base_url,
                        UpstreamProviderRow.api_key == api_key,
                    )
                )
                if not result.first():
                    slug = await allocate_unique_provider_slug(
                        session, provider_type, reserved_slugs
                    )
                    reserved_slugs.add(slug)
                    providers_to_add.append(
                        UpstreamProviderRow(
                            slug=slug,
                            provider_type=provider_type,
                            base_url=base_url,
                            api_key=api_key,
                            enabled=True,
                        )
                    )
                    seeded_provider_keys.add((base_url, api_key))

    ollama_base_url = os.environ.get("OLLAMA_BASE_URL")
    if ollama_base_url:
        ollama_api_key = os.environ.get("OLLAMA_API_KEY", "")
        result = await session.exec(
            select(UpstreamProviderRow).where(
                UpstreamProviderRow.base_url == ollama_base_url,
                UpstreamProviderRow.api_key == ollama_api_key,
            )
        )
        if not result.first():
            slug = await allocate_unique_provider_slug(
                session, "ollama", reserved_slugs
            )
            reserved_slugs.add(slug)
            providers_to_add.append(
                UpstreamProviderRow(
                    slug=slug,
                    provider_type="ollama",
                    base_url=ollama_base_url,
                    api_key=ollama_api_key,
                    enabled=True,
                )
            )
            seeded_provider_keys.add((ollama_base_url, ollama_api_key))

    if settings.chat_completions_api_version and settings.upstream_base_url:
        base_url = settings.upstream_base_url
        api_key = settings.upstream_api_key
        if (base_url, api_key) not in seeded_provider_keys:
            result = await session.exec(
                select(UpstreamProviderRow).where(
                    UpstreamProviderRow.base_url == base_url,
                    UpstreamProviderRow.api_key == api_key,
                )
            )
            if not result.first():
                slug = await allocate_unique_provider_slug(
                    session, "azure", reserved_slugs
                )
                reserved_slugs.add(slug)
                providers_to_add.append(
                    UpstreamProviderRow(
                        slug=slug,
                        provider_type="azure",
                        base_url=base_url,
                        api_key=api_key,
                        api_version=settings.chat_completions_api_version,
                        enabled=True,
                    )
                )
                seeded_provider_keys.add((base_url, api_key))

    if settings.upstream_base_url and settings.upstream_api_key:
        base_url = settings.upstream_base_url
        api_key = settings.upstream_api_key
        if (base_url, api_key) not in seeded_provider_keys:
            result = await session.exec(
                select(UpstreamProviderRow).where(
                    UpstreamProviderRow.base_url == base_url,
                    UpstreamProviderRow.api_key == api_key,
                )
            )
            if not result.first():
                slug = await allocate_unique_provider_slug(
                    session, "custom", reserved_slugs
                )
                reserved_slugs.add(slug)
                providers_to_add.append(
                    UpstreamProviderRow(
                        slug=slug,
                        provider_type="custom",
                        base_url=base_url,
                        api_key=api_key,
                        enabled=True,
                    )
                )
                seeded_provider_keys.add((base_url, api_key))

    for provider in providers_to_add:
        session.add(provider)
        logger.info(
            f"Seeding {provider.provider_type} provider",  # type: ignore[str-format]
            extra={"base_url": provider.base_url, "slug": provider.slug},
        )


def _instantiate_provider(
    provider_row: UpstreamProviderRow,
) -> BaseUpstreamProvider | None:
    """Instantiate an UpstreamProvider from a database row.

    Args:
        provider_row: Database row containing provider configuration

    Returns:
        Instantiated provider or None if provider type is unknown
    """
    from . import upstream_provider_classes

    try:
        provider_classes_by_type = {
            cls.provider_type: cls
            for cls in upstream_provider_classes  # type: ignore[attr-defined]
        }

        provider_class = provider_classes_by_type.get(provider_row.provider_type)

        if provider_class:
            provider = provider_class.from_db_row(provider_row)  # type: ignore[attr-defined]
            if provider is None:
                logger.error(
                    f"Failed to instantiate {provider_row.provider_type} provider",
                    extra={"base_url": provider_row.base_url},
                )
            return provider

        if provider_row.provider_type == "custom":
            return BaseUpstreamProvider.from_db_row(provider_row)

        logger.error(
            f"Unknown provider type: {provider_row.provider_type}",
            extra={"base_url": provider_row.base_url},
        )
        return None
    except Exception as e:
        logger.error(
            f"Failed to instantiate provider: {e}",
            extra={
                "provider_type": provider_row.provider_type,
                "base_url": provider_row.base_url,
                "error": str(e),
            },
        )
        return None