diff --git a/src/comfydv/_llm/ollama_provider.py b/src/comfydv/_llm/ollama_provider.py index 09cb38e..b9ce828 100644 --- a/src/comfydv/_llm/ollama_provider.py +++ b/src/comfydv/_llm/ollama_provider.py @@ -280,6 +280,7 @@ class OllamaProvider: payload_messages = [m.model_dump() for m in messages] total_attempts = max(0, min(int(max_retries), 5)) + 1 response_text = "" + incomplete = False for attempt in range(1, total_attempts + 1): attempt_options = dict(options) if options else {} @@ -317,12 +318,30 @@ class OllamaProvider: _CHAT_RESPONSE_CACHE.set(cache_key, response_text) return response_text + # done: false alongside blank content is a distinct signal from + # an ordinary blank generation — it's Ollama answering before + # the model has actually finished loading/swapping in, observed + # live under model-swap load (issue #27), not the model having + # genuinely generated nothing. Tracked separately so it can be + # raised on below instead of silently returned like a real + # blank generation would be. + incomplete = result.get("done") is False + if attempt < total_attempts: await asyncio.sleep(RETRY_BACKOFF_SECS) - # Every attempt came back blank — never raises here (chat() has - # never validated its output, unlike chat_structured()); return the - # last (blank) attempt uncached so the next queue run tries fresh. + if incomplete: + raise RuntimeError( + f"Ollama returned an incomplete response after " + f"{total_attempts} attempt(s) for model '{model}' — it may " + "still be loading or swapping in memory. Try again in a " + "few seconds." + ) + + # Every attempt came back blank (and complete) — never raises here + # (chat() has never validated its output, unlike chat_structured()); + # return the last (blank) attempt uncached so the next queue run + # tries fresh. return response_text async def chat_structured( diff --git a/tests/test_ollama_provider.py b/tests/test_ollama_provider.py index 03a07b4..6b3d73c 100644 --- a/tests/test_ollama_provider.py +++ b/tests/test_ollama_provider.py @@ -386,6 +386,94 @@ def test_chat_no_retry_needed_does_not_sleep(monkeypatch): assert sleep_calls == [] +# --------------------------------------------------------------------------- +# chat — incomplete (done: false) response distinct from a genuine blank +# generation (issue #27) +# --------------------------------------------------------------------------- +# +# Live-observed against a real Ollama server under model-swap load: /api/chat +# can answer HTTP 200 with a `done: false` stub — model/created_at/message +# all blank — before the target model has actually finished loading. Unlike +# a genuine blank generation, chat() must not fold this into a silent "". + + +def test_chat_raises_when_every_attempt_is_an_incomplete_stub(monkeypatch): + calls = [] + + async def fake_post(url, payload, *, timeout=120.0, headers=None): + calls.append(payload) + return { + "model": "", + "created_at": "0001-01-01T00:00:00Z", + "message": {"role": "", "content": ""}, + "done": False, + } + + monkeypatch.setattr(provider_mod, "_post_json", fake_post) + monkeypatch.setattr(provider_mod.asyncio, "sleep", _fake_sleep) + + with pytest.raises(RuntimeError, match="incomplete response"): + _run_async( + OllamaProvider("http://localhost:11434").chat( + "llama3", [Message(role="user", content="hi")], max_retries=2 + ) + ) + + assert len(calls) == 3 # original + 2 retries, per max_retries=2 + + +def test_chat_recovers_after_incomplete_stub_then_real_response(monkeypatch): + calls = [] + + async def fake_post(url, payload, *, timeout=120.0, headers=None): + calls.append(payload) + if len(calls) == 1: + return { + "model": "", + "created_at": "0001-01-01T00:00:00Z", + "message": {"role": "", "content": ""}, + "done": False, + } + return {"message": {"content": "real answer"}} + + monkeypatch.setattr(provider_mod, "_post_json", fake_post) + monkeypatch.setattr(provider_mod.asyncio, "sleep", _fake_sleep) + + result = _run_async( + OllamaProvider("http://localhost:11434").chat( + "llama3", [Message(role="user", content="hi")] + ) + ) + + assert result == "real answer" + assert len(calls) == 2 + + +def test_chat_exhausted_retries_still_returns_blank_without_raising_when_done_true( + monkeypatch, +): + """Regression guard: a genuine blank generation (done: true, or no + `done` field at all) must keep the pre-existing silent-return behavior — + only the explicit done: false stub escalates to a raised error.""" + calls = [] + + async def fake_post(url, payload, *, timeout=120.0, headers=None): + calls.append(payload) + return {"message": {"content": ""}, "done": True} + + monkeypatch.setattr(provider_mod, "_post_json", fake_post) + monkeypatch.setattr(provider_mod.asyncio, "sleep", _fake_sleep) + + result = _run_async( + OllamaProvider("http://localhost:11434").chat( + "llama3", [Message(role="user", content="hi")], max_retries=2 + ) + ) + + assert result == "" + assert len(calls) == 3 + + # --------------------------------------------------------------------------- # chat_structured # ---------------------------------------------------------------------------