diff --git a/src/fabledassistant/app.py b/src/fabledassistant/app.py index 8bcdf69..286b928 100644 --- a/src/fabledassistant/app.py +++ b/src/fabledassistant/app.py @@ -117,33 +117,88 @@ def create_app() -> Quart: start_log_retention_loop() start_notification_loop() + async def _warm_model(model: str) -> None: + """Warm an already-installed model into VRAM (no pull).""" + try: + async with httpx.AsyncClient(timeout=300.0) as client: + await client.post( + f"{Config.OLLAMA_URL}/api/generate", + json={"model": model, "prompt": "", "keep_alive": "30m"}, + ) + logger.info("Warmed model '%s' into VRAM", model) + except Exception: + logger.warning("Failed to warm model '%s'", model, exc_info=True) + + async def _warm_user_models() -> None: + """ + Warm whichever chat model(s) users have selected in Settings. + + Only warms models that are already installed in Ollama — never auto-pulls. + Falls back silently if no user preferences exist or Ollama is unreachable. + """ + from sqlalchemy import select as sa_select, distinct + + from fabledassistant.models import async_session + from fabledassistant.models.setting import Setting + + # 1. Collect all distinct default_model values users have saved. + try: + async with async_session() as session: + rows = await session.execute( + sa_select(distinct(Setting.value)).where( + Setting.key == "default_model", + Setting.value.isnot(None), + Setting.value != "", + ) + ) + user_models: set[str] = {r for (r,) in rows} + except Exception: + logger.debug("Could not read user model preferences from DB", exc_info=True) + return + + if not user_models: + logger.debug("No user model preferences found; skipping warm-up") + return + + # 2. Ask Ollama which models are currently installed. + try: + async with httpx.AsyncClient(timeout=10.0) as client: + resp = await client.get(f"{Config.OLLAMA_URL}/api/tags") + resp.raise_for_status() + installed: set[str] = {m["name"] for m in resp.json().get("models", [])} + except Exception: + logger.debug("Could not reach Ollama to check installed models", exc_info=True) + return + + # 3. Warm only the intersection (installed AND user-preferred). + for model in user_models: + base = model.removesuffix(":latest") + if model in installed or f"{base}:latest" in installed or base in installed: + await _warm_model(model) + else: + logger.info( + "User-preferred model '%s' is not installed; skipping warm-up " + "(install it via Settings → Models to enable auto-warm)", + model, + ) + async def _pull_model(model: str, warm: bool = False) -> None: try: await ensure_model(model) except Exception: logger.warning( - "Failed to ensure model '%s' — chat may not work until model is available", + "Failed to ensure model '%s'", model, exc_info=True, ) return if warm: - try: - async with httpx.AsyncClient(timeout=300.0) as client: - await client.post( - f"{Config.OLLAMA_URL}/api/generate", - json={"model": model, "prompt": "", "keep_alive": "30m"}, - ) - logger.info("Warmed model '%s' into VRAM", model) - except Exception: - logger.warning("Failed to warm model '%s'", model, exc_info=True) + await _warm_model(model) - # Pull and warm the main model into VRAM at startup so the first request is fast. - asyncio.create_task(_pull_model(Config.OLLAMA_MODEL, warm=True)) - models_to_warm = {Config.OLLAMA_MODEL} - # Also pull the embedding model (nomic-embed-text by default), but no need to warm it. - if Config.EMBEDDING_MODEL not in models_to_warm: - asyncio.create_task(_pull_model(Config.EMBEDDING_MODEL, warm=False)) + # Warm user-preferred chat models that are already installed. + # Also ensure the embedding model is pulled (no warm needed). + asyncio.create_task(_warm_user_models()) + asyncio.create_task(_pull_model(Config.EMBEDDING_MODEL, warm=False)) # After models are pulled, backfill embeddings for existing notes. # Runs in the background so it never blocks the server from accepting requests.