From d05900cb90fcc6d31faef22405b7929952900ac7 Mon Sep 17 00:00:00 2001 From: Abhishek Kumbhar <159713562+Abhiiishek44@users.noreply.github.com> Date: Sun, 19 Jul 2026 02:36:14 +0530 Subject: [PATCH] =?UTF-8?q?fix(llm):=20enhance=20fallback=20logic=20to=20h?= =?UTF-8?q?andle=20empty=20completions=20and=20impr=E2=80=A6=20(#5491)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(llm): enhance fallback logic to handle empty completions and improve metadata handling * fix(llm): stream tool call deltas immediately --- src/llm_core.py | 65 +++++++++++---- tests/test_llm_core_fallback.py | 141 ++++++++++++++++++++++++++++++++ 2 files changed, 191 insertions(+), 15 deletions(-) diff --git a/src/llm_core.py b/src/llm_core.py index af1958f16..4694df7c0 100644 --- a/src/llm_core.py +++ b/src/llm_core.py @@ -2769,8 +2769,9 @@ async def stream_llm_with_fallback(candidates, messages, **kwargs): """Wrap stream_llm with an ordered fallback chain. `candidates` is a list of (url, model, headers). Each is tried in order, - but only retried on a *pre-content* failure — i.e. an ``event: error`` - that arrives before any assistant text / tool-call data has been yielded. + but only retried on a *pre-content* failure — an ``event: error`` or an + empty completion before any assistant text / completed tool call is yielded. + Metadata is held until substantive output commits the candidate. Once a candidate has emitted real output we never switch (that would duplicate streamed tokens); a later error from that candidate passes through unchanged. The dead-host cooldown in stream_llm makes repeat @@ -2789,6 +2790,7 @@ async def stream_llm_with_fallback(candidates, messages, **kwargs): is_last = (i == len(cands) - 1) emitted = False retried = False + pending_metadata = [] async for chunk in stream_llm(url, model, messages, headers=headers, **kwargs): if chunk.startswith("event: error"): if not emitted and not is_last: @@ -2801,34 +2803,67 @@ async def stream_llm_with_fallback(candidates, messages, **kwargs): else: logger.warning(f"[fallback] candidate {model} failed; trying next") break + if not emitted: + # A last-candidate error is already the clearest terminal + # result; do not append an empty-completion error as well. + yield chunk + return yield chunk continue - # Any data chunk other than the terminal [DONE] means real output. - if chunk.startswith("data: ") and not chunk.startswith("data: [DONE]"): + + event_data = {} + is_done = chunk.startswith("data: [DONE]") + if chunk.startswith("data: ") and not is_done: try: event_data = json.loads(chunk[6:]) except Exception: - event_data = {} - if event_data.get("type") == "model_actual": - yield chunk - continue + pass + + delta = event_data.get("delta") + event_type = event_data.get("type") + substantive = ( + isinstance(delta, str) and bool(delta) + ) or ( + event_type == "tool_call_delta" + ) or ( + event_type == "tool_calls" + and bool(event_data.get("calls")) + ) + + if substantive and not emitted: # First real output from a NON-primary candidate: tell the client # the selected model failed and another answered. Without this the # fallback is invisible — a misconfigured provider looks like it # works because the reply is shown under the originally selected # model's name (e.g. a Bedrock/Claude endpoint that 400s every # request but appears fine because another model silently answered). - if not emitted and i > 0: + if i > 0: yield ('data: ' + json.dumps({ "type": "fallback", "selected_model": primary_model, "answered_by": model, "reason": _summarize_stream_error(last_error), }) + '\n\n') + # Metadata must not commit a candidate. Once real output arrives, + # flush it after any fallback notice and before the output itself. + for metadata_chunk in pending_metadata: + yield metadata_chunk + pending_metadata.clear() emitted = True - yield chunk - if not retried: - return # candidate finished (success, or terminal error already sent) - # Every candidate failed pre-content — surface the last error. - if last_error: - yield last_error + + if substantive or emitted: + yield chunk + elif not is_done: + pending_metadata.append(chunk) + + if emitted: + return + if retried: + continue + if not is_last: + last_error = f'event: error\ndata: {json.dumps({"error": f"Model {model} returned no substantive output", "status": 502})}\n\n' + tag = "primary" if i == 0 else "candidate" + logger.warning(f"[fallback] {tag} {model} returned no substantive output; trying next") + continue + yield f'event: error\ndata: {json.dumps({"error": "All model candidates returned no substantive output", "status": 502})}\n\n' + return diff --git a/tests/test_llm_core_fallback.py b/tests/test_llm_core_fallback.py index f1c4e6fef..af99b73b7 100644 --- a/tests/test_llm_core_fallback.py +++ b/tests/test_llm_core_fallback.py @@ -8,6 +8,8 @@ works while a different model silently answers). import json import asyncio +import pytest + from src import llm_core @@ -55,6 +57,145 @@ def test_no_fallback_event_when_primary_succeeds(monkeypatch): assert not any('"fallback"' in c for c in chunks) +def test_done_only_primary_invokes_fallback(monkeypatch): + calls = [] + + def per_model(model): + calls.append(model) + if model == "primary": + return ["data: [DONE]\n\n"] + return [ + 'data: {"type": "model_actual", "requested_model": "backup", "model": "backup-v2"}\n\n', + 'data: {"delta": "backup answer"}\n\n', + "data: [DONE]\n\n", + ] + + chunks = _run_fallback(monkeypatch, per_model) + assert calls == ["primary", "backup"] + assert any('"delta": "backup answer"' in c for c in chunks) + model_idx = next(i for i, c in enumerate(chunks) if '"model_actual"' in c) + fallback_idx = next(i for i, c in enumerate(chunks) if '"fallback"' in c) + answer_idx = next(i for i, c in enumerate(chunks) if '"delta": "backup answer"' in c) + assert fallback_idx < model_idx < answer_idx + + +def test_usage_then_done_primary_invokes_fallback_and_discards_usage(monkeypatch): + calls = [] + + def per_model(model): + calls.append(model) + if model == "primary": + return [ + 'data: {"type": "usage", "data": {"input_tokens": 4, "output_tokens": 0}}\n\n', + "data: [DONE]\n\n", + ] + return ['data: {"delta": "backup answer"}\n\n', "data: [DONE]\n\n"] + + chunks = _run_fallback(monkeypatch, per_model) + assert calls == ["primary", "backup"] + assert not any('"type": "usage"' in c for c in chunks) + + +@pytest.mark.parametrize( + "output_chunk", + [ + 'data: {"delta": "visible text"}\n\n', + 'data: {"delta": "reasoning", "thinking": true}\n\n', + ], +) +def test_text_or_reasoning_output_prevents_fallback(monkeypatch, output_chunk): + calls = [] + + def per_model(model): + calls.append(model) + return [output_chunk, "data: [DONE]\n\n"] + + chunks = _run_fallback(monkeypatch, per_model) + assert calls == ["primary"] + assert output_chunk in chunks + assert not any('"fallback"' in c for c in chunks) + + +def test_whitespace_only_delta_prevents_fallback(monkeypatch): + calls = [] + whitespace = 'data: {"delta": " "}\n\n' + + def per_model(model): + calls.append(model) + return [whitespace, "data: [DONE]\n\n"] + + chunks = _run_fallback(monkeypatch, per_model) + assert calls == ["primary"] + assert whitespace in chunks + assert not any('"fallback"' in c for c in chunks) + + +def test_completed_tool_call_output_prevents_fallback(monkeypatch): + calls = [] + tool_calls = 'data: {"type": "tool_calls", "calls": [{"id": "c1", "name": "bash", "arguments": "{}"}]}\n\n' + + def per_model(model): + calls.append(model) + return [tool_calls, "data: [DONE]\n\n"] + + chunks = _run_fallback(monkeypatch, per_model) + assert calls == ["primary"] + assert tool_calls in chunks + assert not any('"fallback"' in c for c in chunks) + + +def test_tool_call_delta_is_forwarded_immediately_and_prevents_fallback(monkeypatch): + calls = [] + advanced_past_delta = False + tool_delta = 'data: {"type": "tool_call_delta", "index": 0, "arg_delta": "{\\"path\\":"}\n\n' + tool_calls = 'data: {"type": "tool_calls", "calls": [{"id": "c1", "name": "write_file", "arguments": "{\\"path\\":\\"x\\"}"}]}\n\n' + + async def fake_stream(url, model, messages, **kw): + nonlocal advanced_past_delta + calls.append(model) + yield tool_delta + advanced_past_delta = True + yield tool_calls + yield "data: [DONE]\n\n" + + monkeypatch.setattr(llm_core, "stream_llm", fake_stream) + + async def run(): + stream = llm_core.stream_llm_with_fallback( + [("u1", "primary", {}), ("u2", "backup", {})], + [{"role": "user", "content": "hi"}], + ) + first = await anext(stream) + assert first == tool_delta + assert not advanced_past_delta + chunks = [first] + async for chunk in stream: + chunks.append(chunk) + return chunks + + chunks = asyncio.run(run()) + assert calls == ["primary"] + assert tool_calls in chunks + assert not any('"type": "fallback"' in c for c in chunks) + + +def test_empty_final_candidate_surfaces_terminal_error(monkeypatch): + calls = [] + + def per_model(model): + calls.append(model) + if model == "primary": + return [] # clean EOF without substantive output + return ["data: [DONE]\n\n"] + + chunks = _run_fallback(monkeypatch, per_model) + assert calls == ["primary", "backup"] + errors = [c for c in chunks if c.startswith("event: error")] + assert len(errors) == 1 + assert "All model candidates returned no substantive output" in errors[0] + assert '"status": 502' in errors[0] + + def test_dedupe_candidates_keeps_first_of_each_route(): """(url, model) is the route key; later repeats are dropped, order preserved, the first tuple (with its headers) kept, malformed entries filtered."""