diff --git a/src/fabledassistant/app.py b/src/fabledassistant/app.py index 592805b..985d56e 100644 --- a/src/fabledassistant/app.py +++ b/src/fabledassistant/app.py @@ -179,34 +179,64 @@ def create_app() -> Quart: except Exception: logger.warning("Failed to warm model '%s'", model, exc_info=True) + async def _prime_kv_cache(user_id: int, model: str) -> None: + """Send a minimal chat request to prime Ollama's KV cache with the user's system prompt. + + This ensures the next real user message only needs to process its own tokens + rather than the full ~5,600-token system prompt, cutting TTFT from ~25s to <1s. + """ + try: + from fabledassistant.services.llm import build_context + messages, _ = await build_context( + user_id=user_id, + history=[], + current_note_id=None, + user_message=" ", + ) + async with httpx.AsyncClient(timeout=120.0) as client: + await client.post( + f"{Config.OLLAMA_URL}/api/chat", + json={ + "model": model, + "messages": messages, + "stream": False, + "options": {"num_predict": 1}, + "keep_alive": "2h", + }, + ) + logger.info("Primed KV cache for user %d with model '%s'", user_id, model) + except Exception: + logger.warning("Failed to prime KV cache for user %d", user_id, exc_info=True) + async def _warm_user_models() -> None: """ - Warm whichever chat model(s) users have selected in Settings. + Warm whichever chat model(s) users have selected in Settings, then prime + the KV cache with each user's system prompt so the first real message is fast. 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 sqlalchemy import select as sa_select from fabledassistant.models import async_session from fabledassistant.models.setting import Setting - # 1. Collect all distinct default_model values users have saved. + # 1. Collect (user_id, model) pairs for all users with a saved default_model. try: async with async_session() as session: rows = await session.execute( - sa_select(distinct(Setting.value)).where( + sa_select(Setting.user_id, Setting.value).where( Setting.key == "default_model", Setting.value.isnot(None), Setting.value != "", ) ) - user_models: set[str] = {r for (r,) in rows} + user_model_pairs: list[tuple[int, str]] = list(rows) except Exception: logger.debug("Could not read user model preferences from DB", exc_info=True) return - if not user_models: + if not user_model_pairs: logger.debug("No user model preferences found; skipping warm-up") return @@ -220,11 +250,15 @@ def create_app() -> Quart: 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: + # 3. Warm each unique model, then prime KV cache per user. + warmed: set[str] = set() + for user_id_val, model in user_model_pairs: base = model.removesuffix(":latest") if model in installed or f"{base}:latest" in installed or base in installed: - await _warm_model(model) + if model not in warmed: + await _warm_model(model) + warmed.add(model) + await _prime_kv_cache(user_id_val, model) else: logger.info( "User-preferred model '%s' is not installed; skipping warm-up "