model-selector/model_selector/service.py

from __future__ import annotations

import asyncio
import json
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from .cache import SnapshotCache, now_iso
from .config import OpenClawConfigRepository
from .providers import (
    ProviderModel,
    ProviderAdapter,
    build_provider_adapters,
    fetch_with_timeout,
    provider_model_to_config_entry,
)


@dataclass
class ProviderState:
    provider_key: str
    provider_label: str
    updated_at: str | None
    source: str
    error: str | None
    models: list[ProviderModel]


class ModelDashboardService:
    def __init__(
        self,
        config_repo: OpenClawConfigRepository,
        cache: SnapshotCache,
        adapters: dict[str, ProviderAdapter] | None = None,
    ) -> None:
        self.config_repo = config_repo
        self.cache = cache
        self.adapters = adapters or build_provider_adapters()
        self._lock = asyncio.Lock()

    def configured_models(self) -> list[dict[str, Any]]:
        config = self.config_repo.load()
        items = []
        for model in self.config_repo.configured_models(config):
            items.append(
                {
                    "provider_key": model.provider_key,
                    "provider_label": model.provider_label,
                    "model_key": model.model_key,
                    "model_id": model.model_id,
                    "name": model.name,
                    "reasoning": model.reasoning,
                    "inputs": list(model.inputs),
                    "context_window": model.context_window,
                    "max_tokens": model.max_tokens,
                    "alias": model.alias,
                }
            )
        return items

    def configured_model_keys(self) -> list[str]:
        return [model["model_key"] for model in self.configured_models()]

    def load_provider_state_from_cache(self, provider_key: str) -> ProviderState | None:
        cached = self.cache.read(provider_key)
        if not cached:
            return None
        provider_models = []
        for item in cached.data["models"]:
            provider_models.append(
                ProviderModel(
                    provider_key=item["provider_key"],
                    provider_label=item["provider_label"],
                    model_id=item["model_id"],
                    name=item["name"],
                    source=item["source"],
                    reasoning=item.get("reasoning"),
                    context_window=item.get("context_window"),
                    max_tokens=item.get("max_tokens"),
                    pricing_prompt=item.get("pricing_prompt"),
                    pricing_completion=item.get("pricing_completion"),
                    pricing_request=item.get("pricing_request"),
                    input_modalities=tuple(item.get("input_modalities", ())),
                    output_modalities=tuple(item.get("output_modalities", ())),
                    supported_parameters=tuple(item.get("supported_parameters", ())),
                    active=item.get("active"),
                    features=tuple(item.get("features", ())),
                    raw=item.get("raw", {}),
                )
            )
        return ProviderState(
            provider_key=provider_key,
            provider_label=provider_models[0].provider_label if provider_models else provider_key,
            updated_at=cached.updated_at,
            source=cached.data.get("source", "cache"),
            error=cached.data.get("error"),
            models=provider_models,
        )

    def _configured_provider_models(self, provider_key: str) -> list[ProviderModel]:
        models: list[ProviderModel] = []
        for model in self.config_repo.configured_models():
            if model.provider_key != provider_key:
                continue
            raw = {
                "id": model.model_id,
                "name": model.name,
                "context_window": model.context_window,
                "max_tokens": model.max_tokens,
            }
            models.append(
                ProviderModel(
                    provider_key=provider_key,
                    provider_label=model.provider_label,
                    model_id=model.model_id,
                    name=model.name,
                    source="config",
                    reasoning=model.reasoning,
                    context_window=model.context_window,
                    max_tokens=model.max_tokens,
                    pricing_prompt=None,
                    pricing_completion=None,
                    pricing_request=None,
                    input_modalities=model.inputs or ("text",),
                    output_modalities=("text",),
                    supported_parameters=(),
                    active=True,
                    features=("configured",),
                    raw=raw,
                )
            )
        return models

    async def refresh_provider(self, provider_key: str) -> ProviderState:
        adapter = self.adapters[provider_key]
        error: str | None = None
        source = "live"
        updated_at = now_iso()
        models: list[ProviderModel] = []
        try:
            models = await fetch_with_timeout(adapter)
        except Exception as exc:  # pragma: no cover - network failure fallback
            error = str(exc)
            source = "config"
            models = self._configured_provider_models(provider_key)
            cached = self.load_provider_state_from_cache(provider_key)
            if cached:
                return cached
        payload = {
            "source": source,
            "error": error,
            "models": [asdict(model) for model in models],
        }
        self.cache.write(provider_key, payload)
        return ProviderState(
            provider_key=provider_key,
            provider_label=models[0].provider_label if models else provider_key,
            updated_at=updated_at,
            source=source,
            error=error,
            models=models,
        )

    async def add_provider_model(self, provider_key: str, model_id: str) -> dict[str, Any]:
        state = self.load_provider_state_from_cache(provider_key)
        if state is None or not any(model.model_id == model_id for model in state.models):
            state = await self.refresh_provider(provider_key)

        model = next((item for item in state.models if item.model_id == model_id), None)
        if model is None:
            raise ValueError(f"model {provider_key}/{model_id} is not available")

        config_entry = provider_model_to_config_entry(model)
        return self.config_repo.add_model(provider_key, config_entry)

    def delete_configured_model(self, full_key: str) -> dict[str, Any]:
        return self.config_repo.remove_model(full_key)

    async def refresh_all(self) -> dict[str, ProviderState]:
        states: dict[str, ProviderState] = {}
        for provider_key in self.adapters:
            states[provider_key] = await self.refresh_provider(provider_key)
        return states

    async def refresh_periodically(self, stop_event: asyncio.Event, interval_seconds: int) -> None:
        while not stop_event.is_set():
            async with self._lock:
                await self.refresh_all()
            try:
                await asyncio.wait_for(stop_event.wait(), timeout=interval_seconds)
            except asyncio.TimeoutError:
                continue

    def state_payload(self) -> dict[str, Any]:
        providers = {}
        for provider_key in self.adapters:
            cached = self.load_provider_state_from_cache(provider_key)
            providers[provider_key] = self._state_to_payload(cached)

        return {
            "updated_at": now_iso(),
            "configured_models": self.configured_models(),
            "configured_model_keys": self.configured_model_keys(),
            "providers": providers,
        }

    def _state_to_payload(self, state: ProviderState | None) -> dict[str, Any]:
        if not state:
            return {
                "provider_key": None,
                "provider_label": None,
                "updated_at": None,
                "source": "missing",
                "error": "No cached data yet.",
                "models": [],
            }
        return {
            "provider_key": state.provider_key,
            "provider_label": state.provider_label,
            "updated_at": state.updated_at,
            "source": state.source,
            "error": state.error,
            "models": [asdict(model) for model in state.models],
        }

    def update_alias(self, full_key: str, alias: str | None) -> dict[str, Any]:
        return self.config_repo.update_alias(full_key, alias)

    def dump_state(self) -> str:
        return json.dumps(self.state_payload(), indent=2, ensure_ascii=False)