Per-provider attachment validation: registry overrides + inline file-MIME gating

tee_gateway/llm_backend.py

@@ -493,17 +493,29 @@ class AttachmentValidationError(ValueError):
 
 
 def get_model_capabilities(model: str) -> Dict[str, Any]:
-    """Return the LangChain capability profile for a model (``image_inputs``,
-    ``pdf_inputs``, ...), or ``{}`` when the model has no profile data.
+    """Return the capability profile for a model (``image_inputs``,
+    ``pdf_inputs``, ...), or ``{}`` when the model has no capability data.
 
     Reads the public ``.profile`` attribute of the instantiated chat model, which
-    each ``langchain-<provider>`` package populates from maintained model data.
+    each ``langchain-<provider>`` package populates from maintained model data,
+    then applies any explicit overrides from the model registry. The overrides
+    are the only capability source for models the profile registry doesn't cover
+    (e.g. ByteDance ModelArk models routed through the generic OpenAI client).
     """
     try:
         chat = get_chat_model_cached(model, 0.0, 16)
-        return getattr(chat, "profile", None) or {}
+        caps = dict(getattr(chat, "profile", None) or {})
     except Exception:
-        return {}
+        caps = {}
+    try:
+        cfg = get_model_config(model)
+    except ValueError:
+        return caps
+    if cfg.image_inputs is not None:
+        caps["image_inputs"] = cfg.image_inputs
+    if cfg.pdf_inputs is not None:
+        caps["pdf_inputs"] = cfg.pdf_inputs
+    return caps
 
 
 def _iter_content_parts(messages: list) -> Generator[Dict[str, Any], None, None]:
@@ -519,13 +531,53 @@ def _iter_content_parts(messages: list) -> Generator[Dict[str, Any], None, None]
                     yield part
 
 
+# MIME types each provider accepts as inline ``file`` content parts on its
+# chat endpoint (per provider docs, July 2026). ``None`` = no gateway-side
+# restriction (OpenAI's file-inputs accept a broad set — PDFs plus Office,
+# CSV, text and code formats — so the provider validates specifics). An empty
+# set = the provider has no inline file input on the chat path at all (its
+# document features live on a separate endpoint or URL/file-id mechanism).
+# Text-based files don't normally reach this gate — clients fold them into
+# the message as plain text — but SDK callers may send them as file parts.
+_PROVIDER_FILE_INPUT_MIMES: Dict[str, Optional[frozenset]] = {
+    # Anthropic base64 document blocks accept only PDFs (text/plain documents
+    # use a different, non-base64 source shape the gateway doesn't emit).
+    "anthropic": frozenset({"application/pdf"}),
+    # Gemini inline_data: PDFs get document vision; text-family types are
+    # accepted and extracted as plain text.
+    "google": frozenset(
+        {
+            "application/pdf",
+            "text/plain",
+            "text/csv",
+            "text/markdown",
+            "text/x-markdown",
+            "text/html",
+            "text/xml",
+            "text/rtf",
+        }
+    ),
+    "openai": None,
+    "x-ai": frozenset(),
+    "bytedance": frozenset(),
+    "nous": frozenset(),
+    "zai": frozenset(),
+}
+
+
 def validate_attachments(messages: list, model: str) -> None:
     """Reject attachments the target model can't handle.
 
-    Fails *open*: a modality is only rejected when the model's profile explicitly
-    marks it unsupported, so models without profile data are never wrongly blocked
-    (the provider would still reject a truly unsupported combination). Raises
-    ``AttachmentValidationError``.
+    Two layers, both raising ``AttachmentValidationError``:
+
+    - Modality gating fails *open*: images/documents are only rejected when the
+      model's capability profile (or registry override) explicitly marks the
+      modality unsupported, so models without capability data are never wrongly
+      blocked (the provider would still reject a truly unsupported combination).
+    - File-part MIME gating per provider (``_PROVIDER_FILE_INPUT_MIMES``):
+      inline file uploads are limited to what the provider's chat endpoint
+      documents, so e.g. a docx to Anthropic or a PDF to xAI gets a clean 400
+      here instead of the provider's raw error.
     """
     # Fast path: if the request contains no multimodal parts, skip model lookup.
     for part in _iter_content_parts(messages):
@@ -538,8 +590,13 @@ def validate_attachments(messages: list, model: str) -> None:
     image_supported = caps.get("image_inputs")
     pdf_supported = caps.get("pdf_inputs")
 
-    if image_supported is not False and pdf_supported is not False:
-        return  # model accepts every modality; nothing to gate
+    allowed_file_mimes: Optional[frozenset] = None
+    try:
+        provider = get_model_config(model).provider
+    except ValueError:
+        provider = None
+    if provider is not None:
+        allowed_file_mimes = _PROVIDER_FILE_INPUT_MIMES.get(provider)
 
     for part in _iter_content_parts(messages):
         block = _convert_content_part(part)
@@ -549,10 +606,18 @@ def validate_attachments(messages: list, model: str) -> None:
             raise AttachmentValidationError(
                 f"Model {model!r} does not support image attachments."
             )
-        if block["type"] == "file" and pdf_supported is False:
-            raise AttachmentValidationError(
-                f"Model {model!r} does not support document attachments."
-            )
+        if block["type"] == "file":
+            if pdf_supported is False:
+                raise AttachmentValidationError(
+                    f"Model {model!r} does not support document attachments."
+                )
+            # Missing mime falls back to PDF, mirroring what the LangChain
+            # provider packages assume when serializing the block.
+            mime = block.get("mime_type") or "application/pdf"
+            if allowed_file_mimes is not None and mime not in allowed_file_mimes:
+                raise AttachmentValidationError(
+                    f"Model {model!r} does not support {mime} file attachments."
+                )
 
 
 def convert_messages(messages: list) -> List[Any]:

tee_gateway/model_registry.py

@@ -72,6 +72,15 @@ class ModelConfig:
     # means "use the provider default" (see WEB_SEARCH_PRICE_USD_BY_PROVIDER);
     # set an explicit value here to override a single model's web-search price.
     web_search_price_usd: Optional[Decimal] = None
+    # ── Attachment capability overrides (chat models) ──
+    # Consulted by llm_backend.get_model_capabilities ahead of the LangChain
+    # model profile. ``None`` defers to the profile, which fails open when the
+    # model has no profile data. Set an explicit ``False`` for models the
+    # profile registry doesn't cover (e.g. OpenAI-compatible providers like
+    # ByteDance ModelArk) so unsupported attachments are rejected up front with
+    # a clean error instead of surfacing the provider's raw 400.
+    image_inputs: Optional[bool] = None
+    pdf_inputs: Optional[bool] = None
 
 
 # Default per-search USD price charged when a model uses native web search.
@@ -454,17 +463,23 @@ class SupportedModel(Enum):
         input_price_usd=Decimal("0.0000001"),
         output_price_usd=Decimal("0.0000004"),
     )
+    # DeepSeek V4 is text-only: ModelArk rejects multimodal content parts
+    # ("Model do not support image input"), so gate attachments at the gateway.
     DEEPSEEK_V4_FLASH = ModelConfig(
         provider="bytedance",
         api_name="deepseek-v4-flash-260425",
         input_price_usd=Decimal("0.00000014"),
         output_price_usd=Decimal("0.00000028"),
+        image_inputs=False,
+        pdf_inputs=False,
     )
     DEEPSEEK_V4_PRO = ModelConfig(
         provider="bytedance",
         api_name="deepseek-v4-pro-260425",
         input_price_usd=Decimal("0.00000174"),
         output_price_usd=Decimal("0.00000348"),
+        image_inputs=False,
+        pdf_inputs=False,
     )
     # Seedream 4.0 image generation via ModelArk's OpenAI-compatible
     # /images/generations endpoint. Billed at a flat $0.03 per generated image;
@@ -518,26 +533,36 @@ class SupportedModel(Enum):
     # model-id matching. The bare lowercase "hermes-4-70b" is NOT an accepted
     # alias and is rejected with HTTP 400; the accepted form is the capitalized
     # "Hermes-4-70B" (the canonical id "nousresearch/hermes-4-70b" also works).
+    # Hermes 4 is text-only (Llama-3.1 text base): the Nous endpoint has no
+    # image or file input, so gate attachments at the gateway.
     HERMES_4_405B = ModelConfig(
         provider="nous",
         api_name="Hermes-4-405B",
         input_price_usd=Decimal("0.00000009"),
         output_price_usd=Decimal("0.00000037"),
+        image_inputs=False,
+        pdf_inputs=False,
     )
     HERMES_4_70B = ModelConfig(
         provider="nous",
         api_name="Hermes-4-70B",
         input_price_usd=Decimal("0.00000013"),
         output_price_usd=Decimal("0.0000004"),
+        image_inputs=False,
+        pdf_inputs=False,
     )
 
     # ── Z.ai (Model API, OpenAI-compatible) ─────────────────────────────
     # Z.ai publishes GLM-5.2 prices per 1M tokens: $1.40 input, $4.40 output.
+    # GLM-5.2 is text-only (Z.ai restricts image/video/file content parts to
+    # its V-series vision models), so gate attachments at the gateway.
     GLM_5_2 = ModelConfig(
         provider="zai",
         api_name="glm-5.2",
         input_price_usd=Decimal("0.0000014"),
         output_price_usd=Decimal("0.0000044"),
+        image_inputs=False,
+        pdf_inputs=False,
     )
     # GLM-Image uses Z.ai's image endpoint and is billed per generated image.
     # Z.ai returns hosted URLs only (fetched and inlined by the gateway) and

tee_gateway/test/test_tee_core.py

@@ -26,6 +26,7 @@
     canonical_user_content,
     convert_messages,
     extract_usage,
+    get_model_capabilities,
     validate_attachments,
 )
 from tee_gateway.model_registry import get_model_config, get_rate_card
@@ -828,6 +829,86 @@ def test_pdf_blocked_when_model_lacks_support(self):
             with self.assertRaises(AttachmentValidationError):
                 validate_attachments(self._pdf_msg("JVBERi0="), "grok-4")
 
+    DOCX = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+
+    @staticmethod
+    def _file_msg(mime, filename):
+        return [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "file",
+                        "file": {
+                            "filename": filename,
+                            "file_data": f"data:{mime};base64,QUJD",
+                        },
+                    }
+                ],
+            }
+        ]
+
+    def test_provider_file_mime_allowlist(self):
+        # OpenAI file-inputs accept Office formats — no gateway restriction.
+        validate_attachments(self._file_msg(self.DOCX, "a.docx"), "gpt-5")
+        # Anthropic documents are PDF-only; Office must be converted first.
+        with self.assertRaises(AttachmentValidationError):
+            validate_attachments(self._file_msg(self.DOCX, "a.docx"), "claude-sonnet-5")
+        validate_attachments(self._pdf_msg("JVBERi0="), "claude-sonnet-5")
+        # Gemini accepts PDFs and text-family types inline, but not Office.
+        validate_attachments(self._pdf_msg("JVBERi0="), "gemini-3.5-flash")
+        validate_attachments(self._file_msg("text/csv", "a.csv"), "gemini-3.5-flash")
+        with self.assertRaises(AttachmentValidationError):
+            validate_attachments(
+                self._file_msg(self.DOCX, "a.docx"), "gemini-3.5-flash"
+            )
+
+    def test_providers_without_inline_file_input_reject_files(self):
+        # xAI's document path is a separate Responses API workflow, and
+        # ModelArk's is Responses-API-only too — neither takes an inline file
+        # part on chat completions.
+        for model in ("grok-4.5", "seed-1.8"):
+            with self.assertRaises(AttachmentValidationError):
+                validate_attachments(self._pdf_msg("JVBERi0="), model)
+
+    def test_registry_override_blocks_text_only_models(self):
+        # Hermes (Llama-3.1 text base) and GLM-5.2 are text-only per their
+        # official docs; the registry override rejects both modalities.
+        for model in ("hermes-4-405b", "glm-5.2"):
+            for msgs in (self._pdf_msg("JVBERi0="), self._image_msg("aGVsbG8=")):
+                with self.assertRaises(AttachmentValidationError):
+                    validate_attachments(msgs, model)
+
+    def test_images_fail_open_for_seed_models(self):
+        # ModelArk Seed models take images on the chat API; they have no
+        # profile data, so the image gate stays open and the request passes.
+        validate_attachments(self._image_msg("aGVsbG8="), "seed-1.8")
+
+    def test_registry_override_blocks_deepseek_attachments(self):
+        # DeepSeek V4 rides the generic OpenAI-compatible client and has no
+        # LangChain profile data, so the model-registry override is what
+        # rejects attachments here — without it this fell open and ModelArk
+        # returned a raw "Model do not support image input" 400.
+        for model in ("deepseek-v4-pro", "deepseek-v4-flash"):
+            for msgs in (self._pdf_msg("JVBERi0="), self._image_msg("aGVsbG8=")):
+                with self.assertRaises(AttachmentValidationError):
+                    validate_attachments(msgs, model)
+
+    def test_registry_override_wins_over_profile(self):
+        # An explicit registry False must take precedence even if a LangChain
+        # profile exists and claims support.
+        profiled = mock.Mock()
+        profiled.profile = {"image_inputs": True, "pdf_inputs": True}
+        with mock.patch(
+            "tee_gateway.llm_backend.get_chat_model_cached", return_value=profiled
+        ):
+            caps = get_model_capabilities("deepseek-v4-pro")
+        self.assertIs(caps["image_inputs"], False)
+        self.assertIs(caps["pdf_inputs"], False)
+
+    def test_text_only_request_passes_for_deepseek(self):
+        validate_attachments([{"role": "user", "content": "hi"}], "deepseek-v4-pro")
+
 
 # ---------------------------------------------------------------------------
 # llm_backend.canonical_user_content (request-hashing canonicalization)