Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 50 additions & 15 deletions src/llm_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -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 changes: 141 additions & 0 deletions tests/test_llm_core_fallback.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
import json
import asyncio

import pytest

from src import llm_core


Expand Down Expand Up @@ -55,6 +57,145 @@ def per_model(model):
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."""
Expand Down