fix(lint): use builtin generics in FallbackEntry and cover defensive branches

The FallbackEntry alias used typing.Dict/List, tripping the UP006 strict-rule
budget; switch to builtin generics. Drop the now-unreachable None guard in
_prefill_explicitly_unsupported (callers filter None before the lookup) and add
tests for the malformed-fallback and resolver-raises paths so the defensive
branches are exercised.
This commit is contained in:
mateo-berri 2026-06-18 13:45:32 +00:00
parent a7d4f7f131
commit 5fbae0a1e4
No known key found for this signature in database
2 changed files with 34 additions and 4 deletions

View file

@ -20,10 +20,10 @@ else:
MID_STREAM_CONTINUATION_SYSTEM_PROMPT = "You are a helpful assistant. You are given a message and you need to respond to it. You are also given a generated content. You need to respond to the message in continuation of the generated content. Do not repeat the same content. Your response should be in continuation of this text: "
FallbackEntry = Union[str, Dict[str, List[str]]]
FallbackEntry = str | dict[str, list[str]]
def _prefill_explicitly_unsupported(model: str | None) -> bool:
def _prefill_explicitly_unsupported(model: str) -> bool:
"""True only when the model registry explicitly marks the model as NOT
supporting assistant prefill (``supports_assistant_prefill: false``).
@ -31,8 +31,6 @@ def _prefill_explicitly_unsupported(model: str | None) -> bool:
the legacy prefill behavior only models that would reject the prefill with
a 400 anyway are routed to the user-message continuation.
"""
if model is None:
return False
try:
from litellm.utils import get_model_info

View file

@ -241,3 +241,35 @@ def test_resolver_returning_prefill_supporting_model_keeps_legacy():
resolve_underlying_model=aliases.get,
)
_assert_legacy_prefill(result)
def test_malformed_fallback_entry_does_not_crash():
"""A malformed fallback entry (e.g. an empty dict) makes the underlying
fallback-group resolution raise; the capability check must swallow it and
fall back to the legacy prefill rather than break the mid-stream retry."""
result = build_mid_stream_continuation_messages(
messages=MESSAGES,
generated_content=PARTIAL,
model_group="gpt-4",
fallbacks=[{}],
)
_assert_legacy_prefill(result)
def test_resolver_raising_is_swallowed():
"""The injected resolver wraps router lookups that can raise; an exception
must not propagate out of the capability check, and the raw group name is
still evaluated."""
def boom(_model_group):
raise RuntimeError("deployment lookup failed")
result = build_mid_stream_continuation_messages(
messages=MESSAGES,
generated_content=PARTIAL,
model_group="claude-sonnet-4-6",
resolve_underlying_model=boom,
)
# raw name still classified as prefill-rejecting -> user continuation
assert len(result) == 2
assert result[1]["role"] == "user"