diff --git a/src/fabledassistant/services/generation_task.py b/src/fabledassistant/services/generation_task.py index 93793a1..4edff59 100644 --- a/src/fabledassistant/services/generation_task.py +++ b/src/fabledassistant/services/generation_task.py @@ -22,7 +22,7 @@ from fabledassistant.services.generation_buffer import ( GenerationBuffer, GenerationState, ) -from fabledassistant.services.llm import ChatChunk, build_context, generate_completion, stream_chat, stream_chat_with_tools, summarize_history_for_context +from fabledassistant.services.llm import ChatChunk, build_context, generate_completion, stream_chat, stream_chat_with_tools, summarize_history_for_context, wait_for_model_loaded from fabledassistant.services.chat import update_conversation_title from fabledassistant.services.intent import IntentResult, classify_intent from fabledassistant.services.logging import log_generation @@ -210,9 +210,12 @@ async def run_generation( buf.append_event("status", {"status": "Summarizing conversation history..."}) history_to_use, history_summary = await summarize_history_for_context(history, intent_model) - # Phase 3: Build context and start intent classification in parallel. - # We block on context (need messages to stream) — intent is consumed - # after context is ready, at the start of round 0. + # Phase 3: Build context, classify intent, and wait for model — all in parallel. + # build_context is fast DB/search ops that don't need the main model. + # classify_intent uses the small intent model, not the main model. + # wait_for_model_loaded polls /api/ps so the main stream starts without 500 errors. + model_load_task = asyncio.create_task(wait_for_model_loaded(model, timeout=90.0)) + context_task = asyncio.create_task(build_context( user_id, history_to_use, context_note_id, user_content, history_summary=history_summary, @@ -235,6 +238,14 @@ async def run_generation( # Emit context event buf.append_event("context", {"context": context_meta}) + # Wait for main model to be loaded before starting any generation. + # If it's already loaded (common case), this returns immediately. + if not model_load_task.done(): + buf.append_event("status", {"status": "Loading model..."}) + loaded = await model_load_task + if not loaded: + logger.warning("Model %s did not load within 90s — proceeding anyway", model) + t_start = time.monotonic() timing: dict = { "intent_ms": None, diff --git a/src/fabledassistant/services/llm.py b/src/fabledassistant/services/llm.py index a9a7a84..1564aee 100644 --- a/src/fabledassistant/services/llm.py +++ b/src/fabledassistant/services/llm.py @@ -2,6 +2,7 @@ import asyncio import json import logging import re +import time from collections.abc import AsyncGenerator from dataclasses import dataclass, field from typing import Literal @@ -75,6 +76,30 @@ async def ensure_model(model: str) -> None: logger.warning("Failed to pull model '%s' — chat may not work", model, exc_info=True) +async def wait_for_model_loaded(model: str, timeout: float = 90.0) -> bool: + """Poll /api/ps every 2s until the model appears in Ollama's loaded-model list. + + Returns True when the model is loaded, False if timeout elapses first. + Used before generation to avoid streaming 500s during cold model loads. + """ + base = model.removesuffix(":latest") + deadline = time.monotonic() + timeout + while True: + try: + async with httpx.AsyncClient(timeout=5.0) as client: + resp = await client.get(f"{Config.OLLAMA_URL}/api/ps") + resp.raise_for_status() + loaded = {m["name"] for m in resp.json().get("models", [])} + if model in loaded or f"{base}:latest" in loaded or base in loaded: + return True + except Exception: + pass # Ollama may still be starting up + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + await asyncio.sleep(min(2.0, remaining)) + + async def stream_chat( messages: list[dict], model: str,