model-selector/model_selector/main.py

from __future__ import annotations

import asyncio
from contextlib import asynccontextmanager, suppress
from pathlib import Path

from fastapi import FastAPI, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse
from pydantic import BaseModel
from starlette.routing import Router

from .cache import SnapshotCache
from .config import OpenClawConfigRepository
from .providers import ProviderAdapter
from .service import ModelDashboardService
from .settings import AppSettings
from .ui import render_page


def _patch_router_startup_compatibility() -> None:
    """Bridge older FastAPI startup arguments to newer Starlette routers.

    The environment ships a Starlette build that no longer accepts the legacy
    ``on_startup`` and ``on_shutdown`` keyword arguments FastAPI still sends
    during app construction.  Keep the wrapper narrow: ignore those legacy
    hooks because this app uses lifespan management instead.
    """

    original_init = Router.__init__
    if getattr(original_init, "__model_selector_patched__", False):
        return

    def patched_init(self, *args, on_startup=None, on_shutdown=None, **kwargs):
        return original_init(self, *args, **kwargs)

    patched_init.__model_selector_patched__ = True  # type: ignore[attr-defined]
    Router.__init__ = patched_init  # type: ignore[assignment]


_patch_router_startup_compatibility()


class AliasUpdate(BaseModel):
    full_key: str
    alias: str | None = None


class ModelDeleteRequest(BaseModel):
    full_key: str


class ModelAddRequest(BaseModel):
    provider_key: str
    model_id: str


def create_app(
    settings: AppSettings | None = None,
    config_path: Path | None = None,
    cache_dir: Path | None = None,
    adapters: dict[str, ProviderAdapter] | None = None,
) -> FastAPI:
    settings = settings or AppSettings.from_env()
    config_repo = OpenClawConfigRepository(config_path or settings.config_path)
    cache = SnapshotCache(cache_dir or settings.cache_dir)
    service = ModelDashboardService(config_repo, cache, adapters=adapters)

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        stop_event = asyncio.Event()
        await service.refresh_all()
        refresh_task = asyncio.create_task(
            service.refresh_periodically(stop_event, settings.refresh_interval_seconds)
        )
        app.state.service = service
        app.state.refresh_stop_event = stop_event
        app.state.refresh_task = refresh_task
        try:
            yield
        finally:
            stop_event.set()
            refresh_task.cancel()
            with suppress(asyncio.CancelledError):
                await refresh_task

    app = FastAPI(title="Model Selector", lifespan=lifespan)
    app.state.service = service
    app.state.refresh_stop_event = None
    app.state.refresh_task = None

    @app.get("/", response_class=HTMLResponse)
    async def index() -> HTMLResponse:
        return HTMLResponse(render_page(service.state_payload()))

    @app.get("/api/state")
    async def state() -> JSONResponse:
        return JSONResponse(service.state_payload())

    @app.get("/api/config")
    async def config() -> JSONResponse:
        return JSONResponse(
            {
                "config_path": str(config_repo.config_path),
                "configured_models": service.configured_models(),
            }
        )

    @app.post("/api/alias")
    async def alias(update: AliasUpdate) -> JSONResponse:
        if "/" not in update.full_key:
            raise HTTPException(status_code=400, detail="full_key must be provider/model")
        updated = service.update_alias(update.full_key, update.alias)
        return JSONResponse({"ok": True, "config": updated})

    @app.post("/api/config/models/delete")
    async def delete_model(request: ModelDeleteRequest) -> JSONResponse:
        try:
            updated = service.delete_configured_model(request.full_key)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc
        return JSONResponse({"ok": True, "config": updated})

    @app.post("/api/config/models/add")
    async def add_model(request: ModelAddRequest) -> JSONResponse:
        try:
            updated = await service.add_provider_model(request.provider_key, request.model_id)
        except ValueError as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc
        return JSONResponse({"ok": True, "config": updated})

    @app.post("/api/refresh")
    async def refresh() -> JSONResponse:
        states = await service.refresh_all()
        return JSONResponse(
            {
                "ok": True,
                "providers": {key: service._state_to_payload(value) for key, value in states.items()},
            }
        )

    return app


app = create_app()


def main() -> None:
    import uvicorn

    uvicorn.run("model_selector.main:app", host="0.0.0.0", port=3300, reload=False)


if __name__ == "__main__":
    main()