diff --git a/tee_gateway/llm_backend.py b/tee_gateway/llm_backend.py index da7031c..85885d9 100644 --- a/tee_gateway/llm_backend.py +++ b/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-`` package populates from maintained model data. + each ``langchain-`` 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]: diff --git a/tee_gateway/model_registry.py b/tee_gateway/model_registry.py index ae6c51f..fd51683 100644 --- a/tee_gateway/model_registry.py +++ b/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 diff --git a/tee_gateway/test/test_tee_core.py b/tee_gateway/test/test_tee_core.py index 55f87aa..4d4c930 100644 --- a/tee_gateway/test/test_tee_core.py +++ b/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)