diff --git a/py/services/llm_service.py b/py/services/llm_service.py index 04ccb346..2aa1a98f 100644 --- a/py/services/llm_service.py +++ b/py/services/llm_service.py @@ -11,6 +11,7 @@ from __future__ import annotations import asyncio import json import logging +import time from typing import Any, Dict, List, Optional import aiohttp @@ -32,6 +33,16 @@ _catalog_cache: Optional[Dict[str, List[str]]] = None # ``{provider_id: {model_id: max_output_tokens}}``. _model_output_limits: Dict[str, Dict[str, int]] = {} +# Monotonic timestamp of the last failed catalog fetch (None = no failure +# yet). Failed fetches are negatively cached: further calls return the +# empty fallback without hitting the network until the cooldown elapses, +# so users on broken networks don't stall on every settings-modal open. +_catalog_last_failure: Optional[float] = None +_CATALOG_FAILURE_COOLDOWN = 600.0 # seconds + +# Serializes catalog fetches so concurrent callers don't duplicate requests. +_catalog_lock = asyncio.Lock() + _CATALOG_TIMEOUT = aiohttp.ClientTimeout(total=30) # Cloudflare serves brotli when the client advertises it, and brotli is a @@ -54,61 +65,85 @@ async def _load_model_catalog() -> Dict[str, List[str]]: value has a ``models`` sub-dict keyed by model ID. The result is cached in memory after the first successful fetch. Subsequent calls return the cached data immediately. + + Failed fetches are negatively cached: further calls return an empty + dict without hitting the network until ``_CATALOG_FAILURE_COOLDOWN`` + has elapsed, so a broken network does not stall every settings-modal + open. Concurrent callers are serialized behind :data:`_catalog_lock` + so only one request is ever in flight. """ - global _catalog_cache, _model_output_limits + global _catalog_cache, _model_output_limits, _catalog_last_failure if _catalog_cache is not None: return _catalog_cache - try: - async with aiohttp.ClientSession(timeout=_CATALOG_TIMEOUT) as session: - async with session.get(_MODEL_CATALOG_URL, headers=_NO_BROTLI_HEADERS) as resp: - if resp.status != 200: - logger.warning("Model catalog returned HTTP %s", resp.status) - return _catalog_cache or {} - data = await resp.json() - except (aiohttp.ClientError, asyncio.TimeoutError, json.JSONDecodeError, UnicodeDecodeError) as exc: - logger.warning("Failed to fetch model catalog: %s", exc) - return _catalog_cache or {} + async with _catalog_lock: + # Re-check under the lock: another caller may have fetched (or + # failed) while we were waiting. + if _catalog_cache is not None: + return _catalog_cache + if ( + _catalog_last_failure is not None + and time.monotonic() - _catalog_last_failure < _CATALOG_FAILURE_COOLDOWN + ): + logger.debug( + "Skipping model catalog fetch: last attempt failed %.0fs ago", + time.monotonic() - _catalog_last_failure, + ) + return {} - if not isinstance(data, dict): - logger.warning("Model catalog is not a dict, got %s", type(data).__name__) - return _catalog_cache or {} + try: + async with aiohttp.ClientSession(timeout=_CATALOG_TIMEOUT) as session: + async with session.get(_MODEL_CATALOG_URL, headers=_NO_BROTLI_HEADERS) as resp: + if resp.status != 200: + logger.warning("Model catalog returned HTTP %s", resp.status) + _catalog_last_failure = time.monotonic() + return {} + data = await resp.json() + except (aiohttp.ClientError, asyncio.TimeoutError, json.JSONDecodeError, UnicodeDecodeError) as exc: + logger.warning("Failed to fetch model catalog: %s", exc) + _catalog_last_failure = time.monotonic() + return {} - result: Dict[str, List[str]] = {} - output_limits: Dict[str, Dict[str, int]] = {} - for provider_id, provider_info in data.items(): - if not isinstance(provider_info, dict): - continue - models_dict = provider_info.get("models") - if not isinstance(models_dict, dict): - continue - model_ids: List[str] = [] - provider_limits: Dict[str, int] = {} - for mid, model_info in models_dict.items(): - if not isinstance(mid, str): + if not isinstance(data, dict): + logger.warning("Model catalog is not a dict, got %s", type(data).__name__) + _catalog_last_failure = time.monotonic() + return {} + + result: Dict[str, List[str]] = {} + output_limits: Dict[str, Dict[str, int]] = {} + for provider_id, provider_info in data.items(): + if not isinstance(provider_info, dict): continue - model_ids.append(mid) - if isinstance(model_info, dict): - limit = model_info.get("limit") - if isinstance(limit, dict): - output = limit.get("output") - if isinstance(output, (int, float)) and output > 0: - provider_limits[mid] = int(output) - if model_ids: - result[provider_id] = model_ids - if provider_limits: - output_limits[provider_id] = provider_limits + models_dict = provider_info.get("models") + if not isinstance(models_dict, dict): + continue + model_ids: List[str] = [] + provider_limits: Dict[str, int] = {} + for mid, model_info in models_dict.items(): + if not isinstance(mid, str): + continue + model_ids.append(mid) + if isinstance(model_info, dict): + limit = model_info.get("limit") + if isinstance(limit, dict): + output = limit.get("output") + if isinstance(output, (int, float)) and output > 0: + provider_limits[mid] = int(output) + if model_ids: + result[provider_id] = model_ids + if provider_limits: + output_limits[provider_id] = provider_limits - _catalog_cache = result - _model_output_limits = output_limits - logger.debug( - "Loaded model catalog: %d providers, %d total models " - "(%d providers have output limits)", - len(result), - sum(len(m) for m in result.values()), - len(output_limits), - ) - return result + _catalog_cache = result + _model_output_limits = output_limits + logger.debug( + "Loaded model catalog: %d providers, %d total models " + "(%d providers have output limits)", + len(result), + sum(len(m) for m in result.values()), + len(output_limits), + ) + return result def _get_model_max_output(provider: str, model: str) -> Optional[int]: diff --git a/tests/services/test_llm_service.py b/tests/services/test_llm_service.py index 9c00c5c1..4b195157 100644 --- a/tests/services/test_llm_service.py +++ b/tests/services/test_llm_service.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio import json +import time from unittest import mock import pytest @@ -344,6 +345,14 @@ class CorruptJsonResponse(MockResponse): raise UnicodeDecodeError("utf-8", b"\x9a", 0, 1, "invalid start byte") +class SlowResponse(MockResponse): + """Response whose body takes a moment to read, to force contention.""" + + async def json(self): + await asyncio.sleep(0.05) + return self._json_data + + class TestModelCatalog: """Tests for _load_model_catalog / fetch_ollama_models error handling.""" @@ -352,9 +361,11 @@ class TestModelCatalog: """Reset the module-level catalog cache around each test.""" llm_module._catalog_cache = None llm_module._model_output_limits = {} + llm_module._catalog_last_failure = None yield llm_module._catalog_cache = None llm_module._model_output_limits = {} + llm_module._catalog_last_failure = None @pytest.mark.asyncio async def test_load_model_catalog_falls_back_on_unicode_decode_error(self): @@ -401,3 +412,62 @@ class TestModelCatalog: assert models == ["llama3"] assert session.last_headers == {"Accept-Encoding": "gzip, deflate"} + + @pytest.mark.asyncio + async def test_failed_fetch_is_negatively_cached(self): + """A failed fetch is not retried until the cooldown elapses.""" + created = [] + + def factory(*args, **kwargs): + session = MockGetSession(MockResponse(500, text_data="error")) + created.append(session) + return session + + with mock.patch("aiohttp.ClientSession", side_effect=factory): + first = await llm_module._load_model_catalog() + second = await llm_module._load_model_catalog() + + assert first == {} + assert second == {} + assert len(created) == 1 + assert llm_module._catalog_last_failure is not None + + @pytest.mark.asyncio + async def test_fetch_retries_after_cooldown(self): + """Once the cooldown elapses, the next call fetches again.""" + bad = MockGetSession(MockResponse(500, text_data="error")) + with mock.patch("aiohttp.ClientSession", return_value=bad): + assert await llm_module._load_model_catalog() == {} + + # Simulate the cooldown having elapsed. + llm_module._catalog_last_failure = ( + time.monotonic() - llm_module._CATALOG_FAILURE_COOLDOWN - 1 + ) + + good = MockGetSession( + MockResponse(200, json_data={"openai": {"models": {"gpt-4o": {}}}}) + ) + with mock.patch("aiohttp.ClientSession", return_value=good): + catalog = await llm_module._load_model_catalog() + + assert catalog == {"openai": ["gpt-4o"]} + + @pytest.mark.asyncio + async def test_concurrent_fetches_are_deduplicated(self): + """Concurrent callers share a single in-flight fetch.""" + created = [] + + def factory(*args, **kwargs): + session = MockGetSession( + SlowResponse(200, json_data={"openai": {"models": {"gpt-4o": {}}}}) + ) + created.append(session) + return session + + with mock.patch("aiohttp.ClientSession", side_effect=factory): + results = await asyncio.gather( + *(llm_module._load_model_catalog() for _ in range(3)) + ) + + assert len(created) == 1 + assert all(r == {"openai": ["gpt-4o"]} for r in results)