feat: better fallback shim

application/llm/base.py

@@ -203,6 +203,103 @@ class BaseLLM(ABC):
             logger.debug("Fallback payload size estimation failed", exc_info=True)
         return True
 
+    @staticmethod
+    def _fallback_attachment_texts(attachments):
+        """Extracted attachment texts, in upload order, for file-part swaps."""
+        texts = []
+        for attachment in attachments or []:
+            if not isinstance(attachment, dict):
+                continue
+            if attachment.get("content"):
+                texts.append(attachment["content"])
+        return texts
+
+    def _prepare_fallback_messages(self, fallback, messages, attachments=None):
+        """Rebuild primary-prepared messages so the fallback can accept them.
+
+        ``prepare_messages_with_attachments`` ran against the *primary*
+        model, so by the time a fallback engages the array can carry content
+        parts the backup provider cannot take: ``file`` parts hold Files-API
+        ids only the primary's endpoint+credential can resolve, and
+        ``image_url`` parts 4xx on non-vision models. Handing them over
+        unchanged makes the fallback die exactly like the primary did
+        ("Fallback LLM also failed"). Swap what the fallback can't accept
+        for text — the attachment's extracted content when available — and
+        collapse all-text parts arrays to plain string content, the one
+        shape every chat endpoint accepts.
+        """
+        if not messages:
+            return messages
+        try:
+            supported = list(fallback.get_supported_attachment_types() or [])
+        except Exception:
+            supported = []
+        keeps_images = any(str(t).startswith("image/") for t in supported)
+        # A Files-API id resolves only against the endpoint + credential
+        # that minted it (the invariant ``_scoped_file_id`` enforces), so a
+        # pdf-capable fallback keeps file parts only on the same endpoint.
+        keeps_files = "application/pdf" in supported
+        if keeps_files:
+            try:
+                self_scope = getattr(self, "_endpoint_scope", None)
+                fallback_scope = getattr(fallback, "_endpoint_scope", None)
+                keeps_files = (
+                    callable(self_scope)
+                    and callable(fallback_scope)
+                    and self_scope() == fallback_scope()
+                )
+            except Exception:
+                keeps_files = False
+        file_texts = self._fallback_attachment_texts(attachments)
+        prepared = []
+        for message in messages:
+            content = message.get("content") if isinstance(message, dict) else None
+            if not isinstance(content, list):
+                prepared.append(message)
+                continue
+            parts = []
+            all_text = True
+            for part in content:
+                part_type = part.get("type") if isinstance(part, dict) else None
+                if part_type == "file" and not keeps_files:
+                    if file_texts:
+                        parts.append(
+                            {
+                                "type": "text",
+                                "text": f"File content:\n\n{file_texts.pop(0)}",
+                            }
+                        )
+                    else:
+                        filename = (part.get("file") or {}).get(
+                            "filename"
+                        ) or "attachment"
+                        parts.append(
+                            {
+                                "type": "text",
+                                "text": f"[File '{filename}' could not be included]",
+                            }
+                        )
+                elif part_type == "image_url" and not keeps_images:
+                    parts.append(
+                        {
+                            "type": "text",
+                            "text": "[Image attachment omitted: the responding "
+                            "model does not support image input]",
+                        }
+                    )
+                else:
+                    parts.append(part)
+                    if part_type != "text":
+                        all_text = False
+            if all_text:
+                text = "\n\n".join(
+                    p.get("text", "") for p in parts if isinstance(p, dict)
+                )
+                prepared.append({**message, "content": text})
+            else:
+                prepared.append({**message, "content": parts})
+        return prepared
+
     def _execute_with_fallback(
         self, method_name: str, decorators: list, *args, **kwargs
     ):
@@ -274,6 +371,12 @@ class BaseLLM(ABC):
             for decorator in decorators:
                 fallback_method = decorator(fallback_method)
             fallback_kwargs = {**kwargs, "model": fallback.model_id}
+            if fallback_kwargs.get("messages"):
+                fallback_kwargs["messages"] = self._prepare_fallback_messages(
+                    fallback,
+                    fallback_kwargs["messages"],
+                    kwargs.get("_usage_attachments") or kwargs.get("attachments"),
+                )
             try:
                 return fallback_method(fallback, *args, **fallback_kwargs)
             except Exception as e2:
@@ -399,6 +502,12 @@ class BaseLLM(ABC):
             for decorator in decorators:
                 fallback_method = decorator(fallback_method)
             fallback_kwargs = {**kwargs, "model": fallback.model_id}
+            if fallback_kwargs.get("messages"):
+                fallback_kwargs["messages"] = self._prepare_fallback_messages(
+                    fallback,
+                    fallback_kwargs["messages"],
+                    kwargs.get("_usage_attachments") or kwargs.get("attachments"),
+                )
             try:
                 yield from fallback_method(fallback, *args, **fallback_kwargs)
             except Exception as e2:

tests/llm/test_fallback.py

@@ -4,6 +4,7 @@ Verifies that when a primary model fails (immediately or mid-stream), the
 per-agent backup model is used before the global FALLBACK_* settings.
 """
 
+import copy
 from unittest.mock import MagicMock
 
 import httpx
@@ -39,6 +40,7 @@ class FakeLLM(BaseLLM):
         self.gen_called = False
         self.gen_stream_called = False
         self.last_model_received = None  # tracks the model kwarg passed to gen/gen_stream
+        self.last_messages_received = None  # tracks the messages kwarg at the raw level
 
     # Track at the raw-method level. _execute_with_fallback applies
     # decorators to the fallback's raw method directly and
@@ -47,6 +49,7 @@ class FakeLLM(BaseLLM):
     def _raw_gen(self, baseself, model, messages, stream, tools=None, **kwargs):
         self.gen_called = True
         self.last_model_received = model
+        self.last_messages_received = messages
         if self.fail_at is not None:
             raise RuntimeError("primary model unavailable")
         return self.responses[0]
@@ -55,6 +58,7 @@ class FakeLLM(BaseLLM):
         self.gen_stream_called = True
         self.stream_calls = getattr(self, "stream_calls", 0) + 1
         self.last_model_received = model
+        self.last_messages_received = messages
         yielded = 0
         # Per-attempt failure schedule: `fail_schedule[n]` = fail_at value
         # for the n-th call (0-indexed). Falls back to the constant
@@ -1146,3 +1150,149 @@ class TestNoRestreamAfterFinish:
         out = list(primary.gen_stream(**CALL_ARGS))
         assert out == ["fallback chunk"]
         assert backup.gen_stream_called is True
+
+
+# Tests — fallback message reshaping (parts arrays prepared for the primary)
+#
+# ``prepare_messages_with_attachments`` runs against the *primary* model, so
+# by the time a fallback engages the messages can carry ``file`` parts (whose
+# Files-API ids only the primary's endpoint+credential can resolve) and
+# ``image_url`` parts a non-vision fallback 4xxes on. These tests pin the
+# handoff contract: the fallback must receive content it can accept.
+
+
+def _parts_messages():
+    return [
+        {"role": "system", "content": "You are helpful."},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "summarize the attached report"},
+                {"type": "file", "file": {"file_id": "assistant-abc123"}},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": "data:image/png;base64,AAAA"},
+                },
+            ],
+        },
+    ]
+
+
+ATTACHMENTS = [
+    {"id": "att-1", "filename": "report.pdf", "content": "EXTRACTED REPORT TEXT"}
+]
+
+
+class VisionFakeLLM(FakeLLM):
+    def get_supported_attachment_types(self):
+        return ["image/png", "image/jpeg"]
+
+
+class SharedEndpointPdfFakeLLM(FakeLLM):
+    def get_supported_attachment_types(self):
+        return ["application/pdf", "image/png"]
+
+    def _endpoint_scope(self):
+        return "scope-shared"
+
+
+@pytest.mark.integration
+class TestFallbackMessageReshaping:
+
+    def _run_stream(self, primary, messages, attachments):
+        return list(
+            primary.gen_stream(
+                model="test-model",
+                messages=messages,
+                _usage_attachments=attachments,
+            )
+        )
+
+    def test_stream_fallback_gets_flattened_string_content(self):
+        fallback = FakeLLM(stream_chunks=["fb"])
+        primary = FakeLLM(fail_at=0)
+        primary._fallback_llm = fallback
+        original = _parts_messages()
+        snapshot = copy.deepcopy(original)
+
+        chunks = self._run_stream(primary, original, ATTACHMENTS)
+
+        assert chunks == ["fb"]
+        received = fallback.last_messages_received
+        assert received[0] == {"role": "system", "content": "You are helpful."}
+        user_content = received[1]["content"]
+        assert isinstance(user_content, str)
+        assert "summarize the attached report" in user_content
+        assert "EXTRACTED REPORT TEXT" in user_content
+        assert "Image attachment omitted" in user_content
+        # The primary-scoped Files-API id must never reach another endpoint.
+        assert "assistant-abc123" not in user_content
+        # The primary's own message array is not mutated.
+        assert original == snapshot
+
+    def test_stream_vision_fallback_keeps_image_parts(self):
+        fallback = VisionFakeLLM(stream_chunks=["fb"])
+        primary = FakeLLM(fail_at=0)
+        primary._fallback_llm = fallback
+
+        self._run_stream(primary, _parts_messages(), ATTACHMENTS)
+
+        user_content = fallback.last_messages_received[1]["content"]
+        assert isinstance(user_content, list)
+        types = [part["type"] for part in user_content]
+        assert "image_url" in types
+        assert "file" not in types
+        joined = " ".join(
+            part.get("text", "") for part in user_content if part["type"] == "text"
+        )
+        assert "EXTRACTED REPORT TEXT" in joined
+
+    def test_stream_same_endpoint_pdf_fallback_keeps_file_parts(self):
+        fallback = SharedEndpointPdfFakeLLM(stream_chunks=["fb"])
+        primary = SharedEndpointPdfFakeLLM(fail_at=0)
+        primary._fallback_llm = fallback
+
+        self._run_stream(primary, _parts_messages(), ATTACHMENTS)
+
+        user_content = fallback.last_messages_received[1]["content"]
+        assert isinstance(user_content, list)
+        assert {"type": "file", "file": {"file_id": "assistant-abc123"}} in user_content
+
+    def test_stream_file_part_without_extracted_content_becomes_note(self):
+        fallback = FakeLLM(stream_chunks=["fb"])
+        primary = FakeLLM(fail_at=0)
+        primary._fallback_llm = fallback
+
+        self._run_stream(primary, _parts_messages(), attachments=None)
+
+        user_content = fallback.last_messages_received[1]["content"]
+        assert isinstance(user_content, str)
+        assert "could not be included" in user_content
+        assert "assistant-abc123" not in user_content
+
+    def test_stream_string_messages_pass_through_unchanged(self):
+        fallback = FakeLLM(stream_chunks=["fb"])
+        primary = FakeLLM(fail_at=0)
+        primary._fallback_llm = fallback
+        messages = [{"role": "user", "content": "plain text"}]
+
+        self._run_stream(primary, messages, ATTACHMENTS)
+
+        assert fallback.last_messages_received == messages
+
+    def test_gen_fallback_gets_flattened_string_content(self):
+        fallback = FakeLLM(responses=["fb answer"])
+        primary = FakeLLM(fail_at=0)
+        primary._fallback_llm = fallback
+
+        result = primary.gen(
+            model="test-model",
+            messages=_parts_messages(),
+            _usage_attachments=ATTACHMENTS,
+        )
+
+        assert result == "fb answer"
+        user_content = fallback.last_messages_received[1]["content"]
+        assert isinstance(user_content, str)
+        assert "EXTRACTED REPORT TEXT" in user_content
+        assert "assistant-abc123" not in user_content