fix(llm): enhance fallback logic to handle empty completions and impr… (#5491)

* fix(llm): enhance fallback logic to handle empty completions and improve metadata handling

* fix(llm): stream tool call deltas immediately
This commit is contained in:
Abhishek Kumbhar
2026-07-19 02:36:14 +05:30
committed by GitHub
parent b3f8b77317
commit d05900cb90
2 changed files with 191 additions and 15 deletions
+50 -15
View File
@@ -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
+141
View File
@@ -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."""