diff --git a/application/llm/base.py b/application/llm/base.py index df6ed773..395a2d71 100644 --- a/application/llm/base.py +++ b/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: diff --git a/tests/llm/test_fallback.py b/tests/llm/test_fallback.py index e06b1edd..dd97e715 100644 --- a/tests/llm/test_fallback.py +++ b/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