diff --git a/routstr/payment/models.py b/routstr/payment/models.py index 8d7788a1..f75b5620 100644 --- a/routstr/payment/models.py +++ b/routstr/payment/models.py @@ -3,7 +3,7 @@ import json import random import httpx -from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi import APIRouter, Depends, HTTPException, Request, Response from pydantic import BaseModel as V2BaseModel from pydantic.v1 import BaseModel, validator from sqlmodel.ext.asyncio.session import AsyncSession @@ -767,13 +767,19 @@ async def model_paths_for_model(model_id: str) -> dict: return result -@models_router.get("/v1/models") -@models_router.get("/v1/models/", include_in_schema=False) -@models_router.get("/models") -@models_router.get("/models/", include_in_schema=False) -async def models(session: AsyncSession = Depends(get_session)) -> dict: +@models_router.get("/v1/models", response_model=None) +@models_router.get("/v1/models/", include_in_schema=False, response_model=None) +@models_router.get("/models", response_model=None) +@models_router.get("/models/", include_in_schema=False, response_model=None) +async def models( + request: Request, session: AsyncSession = Depends(get_session) +) -> dict | Response: """Get all available models from all providers with database overrides applied.""" - from ..proxy import get_unique_models + from ..proxy import get_unique_models, is_warming_up, models_warming_up_response + + # An empty list here reads as "this node serves nothing"; tell clients to retry. + if is_warming_up(): + return models_warming_up_response(request) items = get_unique_models() data = [] diff --git a/routstr/proxy.py b/routstr/proxy.py index 7c0cf908..6f04d2be 100644 --- a/routstr/proxy.py +++ b/routstr/proxy.py @@ -106,6 +106,19 @@ def is_warming_up() -> bool: return _warming_up +def models_warming_up_response(request: Request) -> Response: + response = create_error_response( + "service_unavailable", + "Models are still loading, retry shortly", + 503, + request=request, + code=MODELS_WARMING_UP, + error_scope=ERROR_SCOPE_NODE, + ) + response.headers["Retry-After"] = str(WARMING_UP_RETRY_AFTER_SECONDS) + return response + + async def reinitialize_upstreams() -> None: """Re-initialize upstream providers from database (called after admin changes).""" global _upstreams @@ -671,16 +684,7 @@ async def _proxy( if not candidates: if is_warming_up(): - response = create_error_response( - "service_unavailable", - "Models are still loading, retry shortly", - 503, - request=request, - code=MODELS_WARMING_UP, - error_scope=ERROR_SCOPE_NODE, - ) - response.headers["Retry-After"] = str(WARMING_UP_RETRY_AFTER_SECONDS) - return response + return models_warming_up_response(request) return create_error_response( "invalid_model", f"Model '{model_id}' not found", 400, request=request ) diff --git a/tests/unit/test_models_warming_up.py b/tests/unit/test_models_warming_up.py index 99f286e7..ea3077bc 100644 --- a/tests/unit/test_models_warming_up.py +++ b/tests/unit/test_models_warming_up.py @@ -4,6 +4,7 @@ import json from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import Response from routstr import proxy as proxy_module from routstr.core.error_scope import ERROR_SCOPE_HEADER, ERROR_SCOPE_NODE @@ -56,3 +57,26 @@ async def test_failed_initialization_still_ends_warm_up() -> None: await proxy_module.initialize_upstreams() assert not proxy_module.is_warming_up() + + +@pytest.mark.asyncio +async def test_models_list_while_warming_up_is_retryable_503() -> None: + from routstr.payment.models import models + + proxy_module.mark_warming_up() + + response = await models(_chat_request(), MagicMock()) + + assert isinstance(response, Response) + assert response.status_code == 503 + assert response.headers["Retry-After"] == "2" + + +@pytest.mark.asyncio +async def test_models_list_after_warm_up_returns_data() -> None: + from routstr.payment.models import models + + with patch.object(proxy_module, "get_unique_models", return_value=[]): + response = await models(_chat_request(), MagicMock()) + + assert response == {"data": []}