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