test(integration): drive scripted wires and provider wiring from data

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/integration/README.md

@@ -2,7 +2,7 @@
 
 These tests exercise a running gateway, PostgreSQL and Redis with an owned local upstream. CircleCI owns this suite. Tests are grouped by behavior, with no automatic test retries or fallback to paid provider calls
 
-The `cost` group runs the scripted-wire cost matrix through the shared integration upstream. The upstream serves the test-owned cost map over loopback through `LITELLM_MODEL_COST_MAP_URL`; cost goldens are checked into the integration suite and must not be copied into the E2E coverage registry
+The `cost` group runs the scripted-wire cost matrix through the shared integration upstream. A provider speaking an existing response shape is a `wires.json` row, a `cases.json` `providers` row and cost-map entries; a new response shape needs a renderer in `scripted_wires.py`
 
 Use `tests/integration/run.py management`, `accounting`, `database`, `providers`, `extensions`, `sdk` or `cost` to run a selected group. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate
 

tests/integration/_support/scripted_wires.py

@@ -34,100 +34,26 @@ import time
 import zlib
 from collections.abc import Mapping
 from dataclasses import dataclass
+from pathlib import Path
 from types import MappingProxyType
-from typing import Final, Literal, TypeAlias
+from typing import Final, Literal, TypeAlias, assert_never
 from urllib.parse import unquote, urlsplit
 
 from pydantic import BaseModel, ConfigDict, TypeAdapter, model_validator
 
-Wire: TypeAlias = Literal[
+Wire: TypeAlias = str
+Shape: TypeAlias = Literal[
     "openai_chat",
     "openai_responses",
     "anthropic_messages",
     "gemini_generate",
-    "together_chat",
-    "fireworks_chat",
-    "azure_chat",
     "bedrock_converse",
-    "vertex_generate",
 ]
-
-WIRE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType(
-    {
-        "openai_chat": "openai",
-        "openai_responses": "openai",
-        "anthropic_messages": "anthropic",
-        "gemini_generate": "gemini",
-        "together_chat": "together",
-        "fireworks_chat": "fireworks",
-        "azure_chat": "azure",
-        "bedrock_converse": "bedrock",
-        "vertex_generate": "vertex",
-    }
-)
-
 StreamUsage: TypeAlias = Literal["final_chunk", "absent"]
 ServiceTier: TypeAlias = Literal["flex", "priority"]
 TerminalKind: TypeAlias = Literal["completed", "incomplete", "unvalidated", "prompt_blocked"]
 
-# Which terminal variant each wire can represent.
-_TERMINAL_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType(
-    {
-        "openai_responses": frozenset({"incomplete", "unvalidated"}),
-        "gemini_generate": frozenset({"prompt_blocked"}),
-        "vertex_generate": frozenset({"prompt_blocked"}),
-    }
-)
-
-
 _BASE_USAGE_FIELDS: Final = frozenset({"fresh_input_tokens", "output_tokens"})
-_OPENAI_FAMILY_USAGE: Final = frozenset(
-    {
-        "cache_read_tokens",
-        "reasoning_tokens",
-        "audio_input_tokens",
-        "audio_output_tokens",
-        "web_search_calls",
-    }
-)
-_CACHE_WRITE_USAGE: Final = frozenset({"cache_write_5m_tokens", "cache_write_1h_tokens"})
-_GEMINI_USAGE: Final = frozenset(
-    {
-        "cache_read_tokens",
-        "reasoning_tokens",
-        "audio_input_tokens",
-        "audio_output_tokens",
-        "image_input_tokens",
-        "video_input_tokens",
-        "web_search_calls",
-        "google_maps_calls",
-    }
-)
-
-_USAGE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType(
-    {
-        wire: usage
-        for wire, usage in (
-            ("openai_chat", _OPENAI_FAMILY_USAGE),
-            ("azure_chat", _OPENAI_FAMILY_USAGE),
-            ("together_chat", _OPENAI_FAMILY_USAGE),
-            ("fireworks_chat", _OPENAI_FAMILY_USAGE),
-            (
-                "openai_responses",
-                frozenset(
-                    {"cache_read_tokens", "reasoning_tokens", "web_search_calls", "file_search_calls"}
-                ),
-            ),
-            (
-                "anthropic_messages",
-                frozenset({"cache_read_tokens", "web_search_calls"}) | _CACHE_WRITE_USAGE,
-            ),
-            ("bedrock_converse", frozenset({"cache_read_tokens"}) | _CACHE_WRITE_USAGE),
-            ("gemini_generate", _GEMINI_USAGE),
-            ("vertex_generate", _GEMINI_USAGE),
-        )
-    }
-)
 
 
 class ScriptedToolCall(BaseModel):
@@ -166,6 +92,32 @@ class ScriptedUsage(BaseModel):
     file_search_calls: int = 0
 
 
+class WireSpec(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    shape: Shape
+    mount: str
+    usage: frozenset[str]
+    terminals: frozenset[TerminalKind]
+
+
+def _load_wires() -> Mapping[str, WireSpec]:
+    adapter: Final = TypeAdapter(dict[str, WireSpec])
+    loaded: Final = adapter.validate_json((Path(__file__).resolve().with_name("wires.json")).read_bytes())
+    known_usage_fields: Final = frozenset(ScriptedUsage.model_fields) - _BASE_USAGE_FIELDS
+    unknown: Final = {
+        wire: sorted(spec.usage - known_usage_fields)
+        for wire, spec in loaded.items()
+        if spec.usage - known_usage_fields
+    }
+    if unknown:
+        raise ValueError(f"wires.json has unknown usage fields: {unknown}")
+    return MappingProxyType(loaded)
+
+
+WIRES: Final[Mapping[str, WireSpec]] = _load_wires()
+
+
 class ScriptedOutput(BaseModel):
     model_config = ConfigDict(frozen=True)
 
@@ -205,9 +157,14 @@ class Scenario(BaseModel):
 
     @model_validator(mode="after")
     def _check_terminal_supported(self) -> Scenario:
+        spec: Final = WIRES.get(self.wire)
+        if spec is None:
+            raise ValueError(
+                f"unknown wire {self.wire}; known wires: {', '.join(sorted(WIRES))}"
+            )
         if (
             self.output.terminal != "completed"
-            and self.output.terminal not in _TERMINAL_CAPS.get(self.wire, frozenset())
+            and self.output.terminal not in spec.terminals
         ):
             raise ValueError(
                 f"wire {self.wire} cannot emit terminal={self.output.terminal}"
@@ -216,7 +173,7 @@ class Scenario(BaseModel):
             field
             for field in self.usage.model_fields_set
             if getattr(self.usage, field)
-            and field not in (_USAGE_CAPS.get(self.wire, frozenset()) | _BASE_USAGE_FIELDS)
+            and field not in (spec.usage | _BASE_USAGE_FIELDS)
         )
         if unsupported:
             raise ValueError(
@@ -230,7 +187,7 @@ class Scenario(BaseModel):
 
     @property
     def mount(self) -> str:
-        return WIRE_MOUNTS[self.wire]
+        return WIRES[self.wire].mount
 
 
 class ScenarioRegistered(BaseModel):
@@ -1266,33 +1223,32 @@ def _render(
         return RenderedResponse(
             200, "application/json", _json_bytes(_responses_body(scenario, requested_model))
         )
-    if scenario.wire == "bedrock_converse":
-        if stream:
-            return RenderedResponse(
-                200, "application/vnd.amazon.eventstream", _bedrock_eventstream(scenario)
-            )
-        return RenderedResponse(200, "application/json", _json_bytes(_bedrock_body(scenario)))
-    if scenario.wire == "vertex_generate":
-        if stream:
-            return RenderedResponse(200, "text/event-stream", _gemini_sse(scenario, requested_model))
-        return RenderedResponse(200, "application/json", _json_bytes(_gemini_body(scenario, requested_model)))
-    if scenario.wire == "anthropic_messages":
-        if stream:
-            return RenderedResponse(200, "text/event-stream", _anthropic_sse(scenario, requested_model))
-        return RenderedResponse(200, "application/json", _json_bytes(_anthropic_body(scenario, requested_model)))
-    if scenario.wire == "gemini_generate":
-        if stream:
-            return RenderedResponse(200, "text/event-stream", _gemini_sse(scenario, requested_model))
-        return RenderedResponse(200, "application/json", _json_bytes(_gemini_body(scenario, requested_model)))
-    if scenario.wire == "openai_responses":
-        if stream:
-            return RenderedResponse(200, "text/event-stream", _responses_sse(scenario, requested_model))
-        return RenderedResponse(200, "application/json", _json_bytes(_responses_body(scenario, requested_model)))
-    # openai_chat, together_chat, fireworks_chat and azure_chat share the
-    # OpenAI chat shape.
-    if stream:
-        return RenderedResponse(200, "text/event-stream", _openai_chat_sse(scenario, requested_model))
-    return RenderedResponse(200, "application/json", _json_bytes(_openai_chat_body(scenario, requested_model)))
+    shape: Final = WIRES[scenario.wire].shape
+    match shape:
+        case "bedrock_converse":
+            if stream:
+                return RenderedResponse(
+                    200, "application/vnd.amazon.eventstream", _bedrock_eventstream(scenario)
+                )
+            return RenderedResponse(200, "application/json", _json_bytes(_bedrock_body(scenario)))
+        case "gemini_generate":
+            if stream:
+                return RenderedResponse(200, "text/event-stream", _gemini_sse(scenario, requested_model))
+            return RenderedResponse(200, "application/json", _json_bytes(_gemini_body(scenario, requested_model)))
+        case "anthropic_messages":
+            if stream:
+                return RenderedResponse(200, "text/event-stream", _anthropic_sse(scenario, requested_model))
+            return RenderedResponse(200, "application/json", _json_bytes(_anthropic_body(scenario, requested_model)))
+        case "openai_responses":
+            if stream:
+                return RenderedResponse(200, "text/event-stream", _responses_sse(scenario, requested_model))
+            return RenderedResponse(200, "application/json", _json_bytes(_responses_body(scenario, requested_model)))
+        case "openai_chat":
+            if stream:
+                return RenderedResponse(200, "text/event-stream", _openai_chat_sse(scenario, requested_model))
+            return RenderedResponse(200, "application/json", _json_bytes(_openai_chat_body(scenario, requested_model)))
+        case _:
+            assert_never(shape)
 
 
 # ---------- registry + request routing ----------

tests/integration/_support/upstream.py

@@ -19,12 +19,12 @@ from starlette.routing import Route
 
 from _fake_openai_endpoint_server import chat_completions, completions, embeddings, health, moderations
 from integration._support.scripted_wires import (
-    WIRE_MOUNTS,
     RenderedResponse,
     Scenario,
     ScenarioDeleted,
     ScenarioRegistered,
     ScenarioStore,
+    WIRES,
     Wire,
     render,
 )
@@ -218,7 +218,7 @@ class ScenarioHandle:
         return f"{self.control_url}/{self.scenario_id}/{self._mount()}"
 
     def _mount(self) -> str:
-        return WIRE_MOUNTS[self.wire]
+        return WIRES[self.wire].mount
 
 
 def register_scenario(scenario: Scenario) -> ScenarioHandle:

tests/integration/_support/wires.json

@@ -0,0 +1,119 @@
+{
+  "openai_chat": {
+    "shape": "openai_chat",
+    "mount": "openai",
+    "usage": [
+      "cache_read_tokens",
+      "reasoning_tokens",
+      "audio_input_tokens",
+      "audio_output_tokens",
+      "web_search_calls"
+    ],
+    "terminals": []
+  },
+  "openai_responses": {
+    "shape": "openai_responses",
+    "mount": "openai",
+    "usage": [
+      "cache_read_tokens",
+      "reasoning_tokens",
+      "web_search_calls",
+      "file_search_calls"
+    ],
+    "terminals": [
+      "incomplete",
+      "unvalidated"
+    ]
+  },
+  "anthropic_messages": {
+    "shape": "anthropic_messages",
+    "mount": "anthropic",
+    "usage": [
+      "cache_read_tokens",
+      "web_search_calls",
+      "cache_write_5m_tokens",
+      "cache_write_1h_tokens"
+    ],
+    "terminals": []
+  },
+  "gemini_generate": {
+    "shape": "gemini_generate",
+    "mount": "gemini",
+    "usage": [
+      "cache_read_tokens",
+      "reasoning_tokens",
+      "audio_input_tokens",
+      "audio_output_tokens",
+      "image_input_tokens",
+      "video_input_tokens",
+      "web_search_calls",
+      "google_maps_calls"
+    ],
+    "terminals": [
+      "prompt_blocked"
+    ]
+  },
+  "together_chat": {
+    "shape": "openai_chat",
+    "mount": "together",
+    "usage": [
+      "cache_read_tokens",
+      "reasoning_tokens",
+      "audio_input_tokens",
+      "audio_output_tokens",
+      "web_search_calls"
+    ],
+    "terminals": []
+  },
+  "fireworks_chat": {
+    "shape": "openai_chat",
+    "mount": "fireworks",
+    "usage": [
+      "cache_read_tokens",
+      "reasoning_tokens",
+      "audio_input_tokens",
+      "audio_output_tokens",
+      "web_search_calls"
+    ],
+    "terminals": []
+  },
+  "azure_chat": {
+    "shape": "openai_chat",
+    "mount": "azure",
+    "usage": [
+      "cache_read_tokens",
+      "reasoning_tokens",
+      "audio_input_tokens",
+      "audio_output_tokens",
+      "web_search_calls"
+    ],
+    "terminals": []
+  },
+  "bedrock_converse": {
+    "shape": "bedrock_converse",
+    "mount": "bedrock",
+    "usage": [
+      "cache_read_tokens",
+      "cache_write_5m_tokens",
+      "cache_write_1h_tokens"
+    ],
+    "terminals": []
+  },
+  "vertex_generate": {
+    "shape": "gemini_generate",
+    "mount": "vertex",
+    "usage": [
+      "cache_read_tokens",
+      "reasoning_tokens",
+      "audio_input_tokens",
+      "audio_output_tokens",
+      "image_input_tokens",
+      "video_input_tokens",
+      "web_search_calls",
+      "google_maps_calls"
+    ],
+    "terminals": [
+      "prompt_blocked"
+    ]
+  }
+}

tests/integration/cost_calculation/cases.json

@@ -1,4 +1,78 @@
 {
+  "providers": [
+    {
+      "litellm_provider": "openai",
+      "mode": "chat",
+      "wire": "openai_chat",
+      "model_prefix": "openai",
+      "litellm_params": {}
+    },
+    {
+      "litellm_provider": "openai",
+      "mode": "responses",
+      "wire": "openai_responses",
+      "model_prefix": "openai/responses",
+      "litellm_params": {}
+    },
+    {
+      "litellm_provider": "anthropic",
+      "mode": "chat",
+      "wire": "anthropic_messages",
+      "model_prefix": "anthropic",
+      "litellm_params": {}
+    },
+    {
+      "litellm_provider": "gemini",
+      "mode": "chat",
+      "wire": "gemini_generate",
+      "model_prefix": null,
+      "litellm_params": {}
+    },
+    {
+      "litellm_provider": "together_ai",
+      "mode": "chat",
+      "wire": "together_chat",
+      "model_prefix": null,
+      "litellm_params": {}
+    },
+    {
+      "litellm_provider": "fireworks_ai",
+      "mode": "chat",
+      "wire": "fireworks_chat",
+      "model_prefix": null,
+      "litellm_params": {}
+    },
+    {
+      "litellm_provider": "azure",
+      "mode": "chat",
+      "wire": "azure_chat",
+      "model_prefix": null,
+      "litellm_params": {
+        "api_version": "2025-04-01-preview"
+      }
+    },
+    {
+      "litellm_provider": "bedrock_converse",
+      "mode": "chat",
+      "wire": "bedrock_converse",
+      "model_prefix": "bedrock/converse",
+      "litellm_params": {
+        "aws_access_key_id": "AKIASCRIPTEDPROVIDER",
+        "aws_secret_access_key": "scripted-secret",
+        "aws_region_name": "us-east-1"
+      }
+    },
+    {
+      "litellm_provider": "vertex_ai-language-models",
+      "mode": "chat",
+      "wire": "vertex_generate",
+      "model_prefix": "vertex_ai",
+      "litellm_params": {
+        "vertex_project": "cc-scripted-project",
+        "vertex_location": "us-central1"
+      }
+    }
+  ],
   "deployments": [
     {
       "map_key": "azure/gpt-5.4-mini",

tests/integration/cost_calculation/cost_matrix.py

@@ -27,7 +27,14 @@ from types import MappingProxyType
 from typing import Final, Literal
 
 from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
-from integration._support.scripted_wires import Scenario, ScriptedOutput, ScriptedToolCall, ScriptedUsage, Wire
+from integration._support.scripted_wires import (
+    WIRES,
+    Scenario,
+    ScriptedOutput,
+    ScriptedToolCall,
+    ScriptedUsage,
+    Wire,
+)
 
 COST_MAP_PATH: Final = Path(__file__).resolve().parent / "cost_map.json"
 CASES_PATH: Final = Path(__file__).resolve().parent / "cases.json"
@@ -251,9 +258,20 @@ class Case(BaseModel):
         )
 
 
+class _ProviderWiringRow(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    litellm_provider: str
+    mode: str
+    wire: str
+    model_prefix: str | None
+    litellm_params: Mapping[str, str]
+
+
 class _CasesFile(BaseModel):
     model_config = ConfigDict(frozen=True)
 
+    providers: tuple[_ProviderWiringRow, ...] = ()
     deployments: tuple[DeploymentSpec, ...] = ()
     cases: tuple[Case, ...] = ()
 
@@ -267,7 +285,7 @@ _DEPLOYMENTS: Final[Mapping[str, DeploymentSpec]] = MappingProxyType(
 
 @dataclass(frozen=True, slots=True)
 class _ProviderWiring:
-    """How a (litellm_provider, mode) pair maps to a sidecar wire, the provider
+    """How a (litellm_provider, mode) pair maps to a provider wire, the provider
     prefix on the registered litellm model string, and extra litellm_params."""
 
     wire: Wire
@@ -275,42 +293,26 @@ class _ProviderWiring:
     litellm_params: Mapping[str, str]
 
 
-_AZURE_PARAMS: Final[Mapping[str, str]] = MappingProxyType({"api_version": "2025-04-01-preview"})
-_BEDROCK_PARAMS: Final[Mapping[str, str]] = MappingProxyType(
-    {
-        "aws_access_key_id": "AKIASCRIPTEDPROVIDER",
-        "aws_secret_access_key": "scripted-secret",
-        "aws_region_name": "us-east-1",
-    }
-)
-_VERTEX_PARAMS: Final[Mapping[str, str]] = MappingProxyType(
-    {
-        "vertex_project": "cc-scripted-project",
-        "vertex_location": "us-central1",
-    }
-)
+def _provider_wiring(rows: tuple[_ProviderWiringRow, ...]) -> Mapping[tuple[str, str], _ProviderWiring]:
+    unknown_wires: Final = sorted({row.wire for row in rows if row.wire not in WIRES})
+    if unknown_wires:
+        raise ValueError(
+            f"cases.json providers has unknown wires: {unknown_wires}; "
+            f"known wires are {sorted(WIRES)}"
+        )
+    return MappingProxyType(
+        {
+            (row.litellm_provider, row.mode): _ProviderWiring(
+                row.wire,
+                row.model_prefix,
+                MappingProxyType(dict(row.litellm_params)),
+            )
+            for row in rows
+        }
+    )
 
-_PROVIDER_WIRING: Final[Mapping[tuple[str, str], _ProviderWiring]] = MappingProxyType(
-    {
-        ("openai", "chat"): _ProviderWiring("openai_chat", "openai", MappingProxyType({})),
-        ("openai", "responses"): _ProviderWiring(
-            "openai_responses", "openai/responses", MappingProxyType({})
-        ),
-        ("anthropic", "chat"): _ProviderWiring(
-            "anthropic_messages", "anthropic", MappingProxyType({})
-        ),
-        ("gemini", "chat"): _ProviderWiring("gemini_generate", None, MappingProxyType({})),
-        ("together_ai", "chat"): _ProviderWiring("together_chat", None, MappingProxyType({})),
-        ("fireworks_ai", "chat"): _ProviderWiring("fireworks_chat", None, MappingProxyType({})),
-        ("azure", "chat"): _ProviderWiring("azure_chat", None, _AZURE_PARAMS),
-        ("bedrock_converse", "chat"): _ProviderWiring(
-            "bedrock_converse", "bedrock/converse", _BEDROCK_PARAMS
-        ),
-        ("vertex_ai-language-models", "chat"): _ProviderWiring(
-            "vertex_generate", "vertex_ai", _VERTEX_PARAMS
-        ),
-    }
-)
+
+_PROVIDER_WIRING: Final[Mapping[tuple[str, str], _ProviderWiring]] = _provider_wiring(CASES_FILE.providers)
 
 
 @dataclass(frozen=True, slots=True)
@@ -390,11 +392,7 @@ def _frontier() -> tuple[FrontierModel, ...]:
         pair = (entry.litellm_provider, entry.mode)
         wiring = _PROVIDER_WIRING.get(pair)
         if wiring is None:
-            raise ValueError(
-                f"cost_map entry {map_key} has no wiring for "
-                f"(litellm_provider={pair[0]}, mode={pair[1]}); add a "
-                f"_ProviderWiring row in cost_matrix.py"
-            )
+            continue
         siblings = groups[pair]
         override_key = (
             siblings[(siblings.index(map_key) + 1) % len(siblings)] if len(siblings) > 1 else None
@@ -567,6 +565,13 @@ def matrix_data_errors() -> tuple[str, ...]:
         for case in CASES
         if (case.family == "transport") != (not case.owns and not case.fallback_for)
     )
+    missing_provider_rows: Final = sorted(
+        f"cost_map entry {map_key} has no providers row for "
+        f"(litellm_provider={entry.litellm_provider}, mode={entry.mode}); "
+        f"add a providers row in cases.json"
+        for map_key, entry in COST_MAP.items()
+        if (entry.litellm_provider, entry.mode) not in _PROVIDER_WIRING
+    )
     input_rates: Final = tuple(entry.input_cost_per_token for entry in COST_MAP.values())
     findings: Final = (
         (
@@ -615,5 +620,10 @@ def matrix_data_errors() -> tuple[str, ...]:
             if family_violations
             else None
         ),
+        (
+            f"cost_map entries without providers rows: {missing_provider_rows}"
+            if missing_provider_rows
+            else None
+        ),
     )
     return tuple(finding for finding in findings if finding is not None)

tests/integration/cost_calculation/test_token_pricing.py

@@ -9,7 +9,7 @@ import pytest
 from pydantic import JsonValue
 
 from integration._support.client import JSON_OBJECT, Gateway
-from integration._support.scripted_wires import ScriptedUsage, Wire
+from integration._support.scripted_wires import WIRES, ScriptedUsage, Wire
 from integration.cost_calculation.conftest import (
     approx_equal,
     assert_total_is_sum_of_components,
@@ -50,12 +50,12 @@ _MATRIX: Final = tuple(
     for model in FRONTIER_MODELS
     for case in cases_for(model)
 )
-_CACHE_WIRES: Final = frozenset({"anthropic_messages", "bedrock_converse"})
-_WEB_SEARCH_OPTION_WIRES: Final = frozenset({"openai_chat", "azure_chat", "openai_responses"})
+_CACHE_SHAPES: Final = frozenset({"anthropic_messages", "bedrock_converse"})
+_WEB_SEARCH_OPTION_SHAPES: Final = frozenset({"openai_chat", "openai_responses"})
 
 
 def _cache_control(usage: ScriptedUsage, wire: Wire) -> dict[str, JsonValue] | None:
-    if wire not in _CACHE_WIRES:
+    if WIRES[wire].shape not in _CACHE_SHAPES:
         return None
     if not (usage.cache_read_tokens or usage.cache_write_5m_tokens or usage.cache_write_1h_tokens):
         return None
@@ -148,7 +148,7 @@ def _chat_body(model: FrontierModel, case: Case, model_name: str, marker: str) -
         **({"audio": {"voice": "alloy", "format": "pcm16"}} if case.audio_output else {}),
         **(
             {"web_search_options": {"search_context_size": case.web_search}}
-            if case.web_search is not None and model.wire in _WEB_SEARCH_OPTION_WIRES
+            if case.web_search is not None and WIRES[model.wire].shape in _WEB_SEARCH_OPTION_SHAPES
             else {}
         ),
         **({"tools": tools} if tools else {}),