diff --git a/app/src/components/Ai.js b/app/src/components/Ai.js index 0da4c9fd..e7c697c0 100644 --- a/app/src/components/Ai.js +++ b/app/src/components/Ai.js @@ -8,7 +8,7 @@ import {StatusContext} from '../contexts/contexts'; import {AiChat} from './AiChat'; import {APIButton, humanFileSize, InfoMessage, PageContainer, TabLinks, useTitle, WROLModeMessage} from './Common'; import {SortableTable} from './SortableTable'; -import {Button, Divider, Group, Header, NumberInput, Select, Table, Text, Toggle, toast} from './ui'; +import {Button, Divider, Group, Header, Loading, Message, NumberInput, Select, Table, Text, Toggle, toast} from './ui'; import {Downloaders} from './Vars'; // The Controller's name for llama-server: the compose service in docker, the systemd unit natively. @@ -38,6 +38,7 @@ export function AiModelRow({model, onDownload}) { export function ManageAi() { const [catalog, setCatalog] = useState(null); + const [catalogError, setCatalogError] = useState(null); const [serviceStatus, setServiceStatus] = useState(null); const [pendingSave, setPendingSave] = useState(false); const {status} = useContext(StatusContext); @@ -48,9 +49,13 @@ export function ManageAi() { const data = await getAiCatalog(); if (data) { setCatalog(data); + setCatalogError(null); + } else { + setCatalogError('The API did not return the model catalog.'); } } catch (e) { console.error(e); + setCatalogError(e.message || 'Failed to fetch the model catalog.'); } }; @@ -108,7 +113,18 @@ export function ManageAi() { }; if (!catalog) { - return ; + // Waiting on the first catalog response. On a slow first start this can take a while, so + // say so rather than showing an empty tab. + return <> + + {catalogError + ? +

{catalogError}

+

The API may still be starting. This page retries automatically, or

+ +
+ : Loading the model catalog…} + ; } const {models, recommended_tier, total_ram, slow_hardware, disk_usage, enabled, active_model} = catalog; diff --git a/app/src/components/Ai.test.js b/app/src/components/Ai.test.js index 12775237..a7212c75 100644 --- a/app/src/components/Ai.test.js +++ b/app/src/components/Ai.test.js @@ -122,6 +122,30 @@ describe('ManageAi', () => { await waitFor(() => expect(postAiSettings).toHaveBeenCalledWith({active_model: 'My-Custom-Model-Q8_0.gguf'})); }); + it('shows a loading state, not a blank tab, until the catalog arrives', async () => { + let resolveCatalog; + getAiCatalog.mockReturnValue(new Promise(resolve => resolveCatalog = resolve)); + render(); + expect(screen.getByText('Loading the model catalog…')).toBeInTheDocument(); + expect(screen.queryByText(/tier is recommended/)).not.toBeInTheDocument(); + + resolveCatalog(catalogFixture()); + await waitFor(() => expect(screen.getByText(/tier is recommended/)).toBeInTheDocument()); + expect(screen.queryByText('Loading the model catalog…')).not.toBeInTheDocument(); + }); + + it('shows the error with a retry when the catalog request fails', async () => { + getAiCatalog.mockRejectedValueOnce(new Error('promise timeout')); + render(); + await waitFor(() => expect(screen.getByText('Could not load the model catalog')).toBeInTheDocument()); + expect(screen.getByText('promise timeout')).toBeInTheDocument(); + + getAiCatalog.mockResolvedValue(catalogFixture()); + fireEvent.click(screen.getByText('Retry now')); + await waitFor(() => expect(screen.getByText(/tier is recommended/)).toBeInTheDocument()); + expect(screen.queryByText('Could not load the model catalog')).not.toBeInTheDocument(); + }); + it('shows the slow-hardware warning on a Pi 4', async () => { getAiCatalog.mockResolvedValue(catalogFixture({slow_hardware: true})); render(); diff --git a/modules/ai/catalog.py b/modules/ai/catalog.py index 416729b4..e8fdae9b 100644 --- a/modules/ai/catalog.py +++ b/modules/ai/catalog.py @@ -6,17 +6,23 @@ committed WROLPi public key. The release process refreshes AI_MODELS from the p """ import json import tempfile +import time from pathlib import Path from typing import Optional import psutil +from wrolpi import flags from wrolpi.common import aiohttp_get, get_media_directory, logger, verify_gpg_signature logger = logger.getChild(__name__) CDN = 'https://wrolpi.nyc3.cdn.digitaloceanspaces.com' MODELS_MANIFEST_URL = f'{CDN}/ai/models.json' +# The manifest is ~2KB; a black-holed network must not cost 30s per Manage-page poll. +MANIFEST_TIMEOUT = 5 +# How long a fetched (or failed) manifest is reused before the CDN is asked again. +MANIFEST_CACHE_SECONDS = 60 * 60 # Sizes are bytes; sha256 values match the published manifest (wrolpi/scripts/publish_ai_models.py # refreshes both). The meta4 sidecar (GPG-verified at download time) is what enforces the hash. @@ -74,12 +80,12 @@ async def fetch_models_manifest(url: str = None) -> dict: url = url or MODELS_MANIFEST_URL sig_url = f'{url}.sig' - async with aiohttp_get(url, timeout=30) as response: + async with aiohttp_get(url, timeout=MANIFEST_TIMEOUT) as response: if response.status != 200: raise RuntimeError(f'Failed to fetch models manifest: HTTP {response.status}') manifest_bytes = await response.content.read() - async with aiohttp_get(sig_url, timeout=30) as response: + async with aiohttp_get(sig_url, timeout=MANIFEST_TIMEOUT) as response: if response.status != 200: raise RuntimeError(f'Failed to fetch models manifest signature: HTTP {response.status}') signature_bytes = await response.content.read() @@ -97,19 +103,40 @@ async def fetch_models_manifest(url: str = None) -> dict: return json.loads(manifest_bytes) +# (models, source, expires_at) from the last CDN attempt, success or failure. Per process. +_catalog_cache: Optional[tuple[list, str, float]] = None + + +def clear_models_catalog_cache(): + global _catalog_cache + _catalog_cache = None + + async def get_models_catalog() -> (list, str): """The list of downloadable models, and where it came from ('cdn' or 'bundled'). The CDN manifest wins when reachable; the bundled constant keeps the Manage tab working - offline.""" + offline. The CDN is not tried while the have_internet flag is clear, and the result of an + attempt (including a failure) is reused for MANIFEST_CACHE_SECONDS so the Manage tab's 30s + polling does not pay the fetch on every request.""" + global _catalog_cache + if _catalog_cache and _catalog_cache[2] > time.monotonic(): + return _catalog_cache[0], _catalog_cache[1] + + # `is_set()` is None under pytest without the flags lock: only an explicit False skips. + if flags.have_internet.is_set() is False: + logger.debug('No internet, using bundled models catalog') + return AI_MODELS, 'bundled' + + models, source = AI_MODELS, 'bundled' try: manifest = await fetch_models_manifest() - models = manifest.get('models') or [] - if models: - return models, 'cdn' + if cdn_models := manifest.get('models'): + models, source = cdn_models, 'cdn' except Exception as e: logger.debug(f'Could not fetch models manifest, using bundled catalog: {e}') - return AI_MODELS, 'bundled' + _catalog_cache = (models, source, time.monotonic() + MANIFEST_CACHE_SECONDS) + return models, source def get_total_ram_bytes() -> int: diff --git a/modules/ai/conftest.py b/modules/ai/conftest.py index 7fdd759b..54df141e 100644 --- a/modules/ai/conftest.py +++ b/modules/ai/conftest.py @@ -15,3 +15,12 @@ def test_ai_config(test_directory, async_client): with ai_config_context() as config: config.initialize(api_app.shared_ctx.ai_config) yield config_path + + +@pytest.fixture(autouse=True) +def _clear_models_catalog_cache(): + """Each test starts with an empty manifest cache.""" + from modules.ai import catalog + catalog.clear_models_catalog_cache() + yield + catalog.clear_models_catalog_cache() diff --git a/modules/ai/test/test_config_catalog.py b/modules/ai/test/test_config_catalog.py index 28369004..8ec51c4d 100644 --- a/modules/ai/test/test_config_catalog.py +++ b/modules/ai/test/test_config_catalog.py @@ -54,6 +54,7 @@ async def test_models_catalog_fallback(async_client): assert source == 'bundled' assert models == catalog.AI_MODELS + catalog.clear_models_catalog_cache() # the failed attempt is cached; see test_models_catalog_is_cached manifest = dict(version=1, models=[dict(name='new.gguf', tier='small')]) with mock.patch('modules.ai.catalog.fetch_models_manifest', return_value=manifest): models, source = await catalog.get_models_catalog() @@ -285,3 +286,53 @@ async def test_manage_catalog_skips_unloadable_custom_models(async_client, test_ content=json.dumps(dict(active_model=bad))) assert response.status_code == HTTPStatus.BAD_REQUEST, bad assert 'cannot be loaded' in response.json['error'] + + +@pytest.mark.asyncio +async def test_models_catalog_is_cached(async_client): + """The CDN manifest is fetched once per hour, and a failed attempt is cached too. + + The Manage tab polls the catalog every 30s; without a cache an offline device paid the fetch + timeout on every poll.""" + manifest = dict(version=1, models=[dict(name='new.gguf', tier='small')]) + with mock.patch('modules.ai.catalog.fetch_models_manifest', return_value=manifest) as fetch: + assert (await catalog.get_models_catalog())[1] == 'cdn' + assert (await catalog.get_models_catalog())[1] == 'cdn' + assert fetch.call_count == 1 + + # A failure is cached for the same window: the next call does not retry. + catalog.clear_models_catalog_cache() + with mock.patch('modules.ai.catalog.fetch_models_manifest', side_effect=RuntimeError('offline')) as fetch: + assert (await catalog.get_models_catalog())[1] == 'bundled' + assert (await catalog.get_models_catalog())[1] == 'bundled' + assert fetch.call_count == 1 + + # Once the window passes, the CDN is tried again. + with mock.patch('modules.ai.catalog.fetch_models_manifest', return_value=manifest) as fetch, \ + mock.patch('modules.ai.catalog.time.monotonic', return_value=1e9): + assert (await catalog.get_models_catalog())[1] == 'cdn' + assert fetch.call_count == 1 + + +@pytest.mark.asyncio +async def test_models_catalog_skips_cdn_without_internet(async_client, flags_lock): + """When the have_internet flag is clear the bundled catalog is served without touching the + network, so an offline Manage page renders immediately. A skipped attempt is not cached: the + CDN is tried as soon as internet comes up.""" + from wrolpi import flags + flags.have_internet.clear() + with mock.patch('modules.ai.catalog.fetch_models_manifest') as fetch: + models, source = await catalog.get_models_catalog() + assert source == 'bundled' and models == catalog.AI_MODELS + fetch.assert_not_called() + + flags.have_internet.set() + manifest = dict(version=1, models=[dict(name='new.gguf', tier='small')]) + with mock.patch('modules.ai.catalog.fetch_models_manifest', return_value=manifest) as fetch: + assert (await catalog.get_models_catalog())[1] == 'cdn' + fetch.assert_called_once() + + +def test_manifest_fetch_timeout_is_short(): + """A 2KB manifest must not wait 30s per request on a black-holed network.""" + assert catalog.MANIFEST_TIMEOUT <= 5