diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ea4c011d809..abbecc0c197 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3862,9 +3862,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"}) -DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset( - {"supports_bedrock_native_structured_output"} -) +DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset({"supports_bedrock_native_structured_output"}) DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset({"token_count_framing", "token_count_model"}) diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index d2706fdeb38..90f33dfbcce 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -267,3 +267,25 @@ def team_admin_permissions(gateway: Gateway, fields: Sequence[str]) -> Iterator[ yield finally: _set_team_admin_permissions(gateway, [str(field) for field in original] if isinstance(original, list) else ()) + + +def held(read: Callable[[], T], satisfied: Callable[[T], bool], *, holding: float, seconds: float = 10) -> T: + """`eventually`, and then `satisfied` must hold on every read for `holding` more seconds; a miss restarts the hold.""" + deadline: Final = time.monotonic() + seconds + while True: + converged: Final = eventually(read, satisfied, seconds=max(deadline - time.monotonic(), 0.0)) + misses: Final = _misses_within(read, satisfied, holding) + if not misses: + return converged + assert time.monotonic() < deadline, f"State did not hold: {misses[0]!r}" + time.sleep(0.1) + + +def _misses_within(read: Callable[[], T], satisfied: Callable[[T], bool], holding: float) -> tuple[T, ...]: + until: Final = time.monotonic() + holding + while time.monotonic() < until: + time.sleep(0.1) + observed: Final = read() + if not satisfied(observed): + return (observed,) + return () diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index e645126032b..34e4a4ca63d 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -423,3 +423,21 @@ def main() -> None: if __name__ == "__main__": main() + + +def _aws_int_header(name: str, value: int) -> bytes: + name_bytes: Final = name.encode() + return struct.pack("!B", len(name_bytes)) + name_bytes + struct.pack("!B", 4) + struct.pack("!i", value) + + +def aws_event_stream_frame(headers: Mapping[str, str | int], payload: bytes) -> bytes: + """One AWS event-stream frame: a string header is wire type 7, an int header wire type 4 (int32).""" + headers_bytes: Final = b"".join( + _aws_str_header(name, value) if isinstance(value, str) else _aws_int_header(name, value) + for name, value in headers.items() + ) + total_length: Final = 12 + len(headers_bytes) + len(payload) + 4 + prelude: Final = struct.pack("!II", total_length, len(headers_bytes)) + prelude_crc: Final = struct.pack("!I", zlib.crc32(prelude) & 0xFFFFFFFF) + message: Final = prelude + prelude_crc + headers_bytes + payload + return message + struct.pack("!I", zlib.crc32(message) & 0xFFFFFFFF) diff --git a/tests/integration/mcp/test_mcp_transports.py b/tests/integration/mcp/test_mcp_transports.py index 30d74cf5fdb..7b47c260504 100644 --- a/tests/integration/mcp/test_mcp_transports.py +++ b/tests/integration/mcp/test_mcp_transports.py @@ -1,12 +1,11 @@ -from integration._support.process import owned_proxy -import yaml -from pathlib import Path import json import uuid +from pathlib import Path from typing import Final import httpx import pytest +import yaml from integration._support.client import Gateway from integration._support.mcp import ( ENTRY_POINTS, @@ -25,6 +24,7 @@ from integration._support.mcp import ( text_result, tool_calls, ) +from integration._support.process import owned_proxy ADD: Final = {"http": "add", "sse": "add", "stdio": "add", "openapi": "getpet"} ARGUMENTS: Final = {"add": {"a": 3, "b": 4}, "getpet": {"petId": "7"}} diff --git a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py index 9c472e4bfc3..cce7c2eaeee 100644 --- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -1,111 +1,101 @@ -import asyncio import base64 import functools import json import re -import signal -import threading -import time import uuid -from collections.abc import Callable, Iterator, Mapping -from concurrent.futures import ThreadPoolExecutor -from contextlib import ExitStack +from collections.abc import Mapping from dataclasses import dataclass -from hashlib import sha256 -from pathlib import Path -from queue import SimpleQueue from typing import Final -import anthropic import httpx -import openai -import psutil import pytest -import yaml -from integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value -from integration._support.database import read_rows -from integration._support.ports import allocate_port -from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.client import Gateway, Scenario, list_value, object_value, string_value from integration._support.upstream import aws_event_stream_frame from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter _JSON_TOOL: Final = "json_tool_call" + + _USER_TOOL: Final = "lookup" -_STRUCTURED_OUTPUTS_BETA: Final = "structured-outputs-2025-11-13" + + _TOOL_MODEL: Final = "claude-sonnet-5" -_EFFORT_GATED_MODEL: Final = "claude-haiku-4-5" + + _NATIVE_MODEL: Final = "claude-opus-5-5" + + _FORCED_TOOL_NATIVE_MODEL: Final = "claude-haiku-5-5" -_UNMAPPED_MODEL: Final = "claude-scripted-9" -_BEDROCK_INVOKE_MODEL: Final = "anthropic.claude-sonnet-5" + + _BEDROCK_JSON_TOOL_MODEL: Final = "us.anthropic.claude-opus-4-8" + + _THINKING_SIGNATURE: Final = "c2NyaXB0ZWQtc2lnbmF0dXJl" + + _PROJECT: Final = "scripted-project" + + _LOCATION: Final = "us-east5" + + _VERTEX_MODELS: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/" + + _VERTEX_STREAM_SUFFIX: Final = ":streamRawPredict?alt=sse" + + _ANTHROPIC_KEY: Final = "synthetic-anthropic-key" + + _AZURE_KEY: Final = "synthetic-azure-ai-key" + + _BEDROCK_TOKEN: Final = "synthetic-bedrock-bearer" -_PLATFORM_KEY: Final = "synthetic-platform-key" -_MANTLE_KEY: Final = "synthetic-bedrock-mantle-key" -_SCRIPTED_KEY: Final = "scripted" -_NIM_GROUP: Final = "scripted-nim" -_WATSONX_VERSION: Final = "2024-03-13" -_OWNED_MODEL: Final = "partner-claude" + + _DEPRECATED_OUTPUT_FORMAT: Final = "output_format: this field is deprecated, use output_config.format" + + _BAD_OUTPUT_CONFIG: Final = "output_config: Input should be a valid dictionary" + + _BAD_FORMAT: Final = "output_config.format: Input should be a json_schema object" + + _SCHEMA: Final[dict[str, JsonValue]] = { "type": "object", "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}}, "required": ["tier", "summary"], "additionalProperties": False, } + + _RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { "type": "json_schema", "json_schema": {"name": "risk", "schema": _SCHEMA}, } -_TEXT_FORMAT: Final[dict[str, JsonValue]] = { - "format": {"type": "json_schema", "name": "risk", "schema": _SCHEMA}, -} -_OUTPUT_FORMAT: Final[dict[str, JsonValue]] = {"type": "json_schema", "schema": _SCHEMA} -_NAME_SCHEMA: Final[dict[str, JsonValue]] = { - "type": "object", - "properties": {"name": {"type": "string"}, "tier": {"type": "integer"}}, - "required": ["name", "tier"], - "additionalProperties": False, -} -_NAME_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { - "type": "json_schema", - "json_schema": {"name": "named", "schema": _NAME_SCHEMA}, -} -_OPTIONAL_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": _SCHEMA["properties"]} + + _VALUES_SCHEMA: Final[dict[str, JsonValue]] = { "type": "object", "properties": {"values": {"type": "object"}}, "required": ["values"], } + + _VALUES_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { "type": "json_schema", "json_schema": {"name": "valued", "schema": _VALUES_SCHEMA}, } -_ADDITIONAL_OBJECT_SCHEMA: Final[dict[str, JsonValue]] = { - **_OPTIONAL_SCHEMA, - "additionalProperties": {"type": "object"}, -} -_PATTERN_SCHEMA: Final[dict[str, JsonValue]] = { - **_OPTIONAL_SCHEMA, - "patternProperties": {"^[a-z_]+$": {"type": "object"}}, -} -_ADDITIONAL_TRUE_SCHEMA: Final[dict[str, JsonValue]] = {**_OPTIONAL_SCHEMA, "additionalProperties": True} -_CLOSED_OPTIONAL_SCHEMA: Final[dict[str, JsonValue]] = {**_OPTIONAL_SCHEMA, "additionalProperties": False} -_CALLERS_FORMAT: Final[dict[str, JsonValue]] = {"type": "json_schema", "schema": {**_SCHEMA, "title": "callers"}} -_ANTHROPIC_JSON_TOOL: Final[dict[str, JsonValue]] = {"name": _JSON_TOOL, "input_schema": _SCHEMA} + + _FORCED_JSON_TOOL: Final[dict[str, JsonValue]] = {"type": "tool", "name": _JSON_TOOL} -_REASONING_FIELDS: Final = ("thinking", "output_config", "max_tokens") + + _USER_TOOL_SPEC: Final[dict[str, JsonValue]] = { "type": "function", "function": { @@ -113,99 +103,45 @@ _USER_TOOL_SPEC: Final[dict[str, JsonValue]] = { "parameters": {"type": "object", "properties": {"q": {"type": "string"}}, "required": ["q"]}, }, } + + _MARKER: Final = re.compile(r"\b([a-z_]+) ([0-9a-f]{32})\b") + + _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) -_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") -_OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120 -_BURST_SHAPES: Final = ("wrapped", "plain", "named", "values") -_STREAM_SHAPES: Final = ( - "plain", - "wrapped", - "named", - "values", - "large", - "two_candidates", - "inner_incomplete", - "self_shaped", - "scalar", - "empty", - "list_input", - "deep", - "other_tool", - "named_partial", - "parameter_label", - "function_label", - "label_unshared", - "label_two_objects", - "label_encoded", - "named_encoded", - "undeclared_encoded", - "encoded_deep", - "number_first", - "bool_first", - "null_first", -) -_HELD_STREAM_SHAPES: Final = frozenset( - { - "wrapped", - "named", - "values", - "large", - "two_candidates", - "inner_incomplete", - "deep", - "other_tool", - "named_partial", - "parameter_label", - "function_label", - "label_unshared", - "label_two_objects", - "label_encoded", - "named_encoded", - "undeclared_encoded", - "encoded_deep", - } -) + + _AS_WRITTEN_SHAPES: Final = ("label_non_ascii", "label_two_objects_non_ascii") + + _TEXT_ANSWER_SHAPES: Final = frozenset({"undeclared_wrapper"}) + + _LABEL: Final = "RiskAssessment" + + _NON_ASCII_LABEL: Final = "Évaluation des risques 风险" + + _ENCODED_DEPTH: Final = 200_000 -_FAMILY: Final = ( - "anthropic_unmapped", - "azure_ai_unmapped", - "bedrock_invoke", - "bedrock_mantle", - "claude_platform_unmapped", - "vertex_haiku", -) + + _FRAGMENT_SIZE: Final = 20 + + _TIMING_BYTES: Final = 512 * 1024 + + _TIMING_FRAGMENT: Final = 4096 + + _TIMING_SMALL_BYTES: Final = 128 * 1024 -_TIMING_GROWTH_LIMIT: Final = 6.0 -_HEALTH_LATENCY_SECONDS: Final = 1.0 + + _SSE: Final = "text/event-stream" + + _EVENT_STREAM: Final = "application/vnd.amazon.eventstream" -_RELAY_CHUNKS: Final = ( - b'data: {"choices":[{"delta":{"content":"relay one"}}]}\n\n', - b'data: {"choices":[{"delta":{"content":"relay two"}}]}\n\n', - b"data: [DONE]\n\n", -) -_PASS_THROUGH_ROUTES: Final = ( - "gemini", - "cohere", - "vllm", - "mistral", - "typesafe", - "openrouter", - "azure", - "openai", - "gigachat", - "watsonx", - "tinyfish", - "nvidia_nim", -) @dataclass(frozen=True, slots=True) @@ -333,30 +269,6 @@ def _timing_text(shape: str) -> str | None: return None -def _expected(shape: str, marker_hex: str) -> JsonValue: - match shape: - case ( - "plain" - | "wrapped" - | "named" - | "values" - | "named_input" - | "deep" - | "other_tool" - | "parameter_label" - | "function_label" - | "named_encoded" - ): - return _payload(marker_hex) - case "large": - return {**_payload(marker_hex), "summary": marker_hex + "x" * 5000} - case "named_partial" | "inner_incomplete": - return {"tier": 2} - case "name_property": - return {"name": marker_hex, "tier": 1} - return _tool_input(shape, marker_hex) - - def _fragments(text: str, size: int) -> tuple[str, ...]: return tuple(text[start : start + size] for start in range(0, len(text), size)) @@ -365,10 +277,6 @@ def _fragment_size(shape: str) -> int: return _TIMING_FRAGMENT if shape == "encoded_deep" or _timing_text(shape) is not None else _FRAGMENT_SIZE -def _fragment_count(shape: str, marker_hex: str) -> int: - return len(_fragments(_tool_text(shape, marker_hex), _fragment_size(shape))) - - def _message(message_id: str, content: list[JsonValue], stop_reason: str) -> dict[str, JsonValue]: return { "id": message_id, @@ -537,28 +445,6 @@ def _peer(request: Request) -> Reply: return _text_reply(message_id, f"plain text {marker_hex}", framing) -def _relay_peer(request: Request) -> Reply: - assert request.body, request.target - return Reply(content_type=_SSE, chunks=_RELAY_CHUNKS) - - -def _failing(status: int) -> Callable[[Request], Reply]: - def respond(request: Request) -> Reply: - _, marker_hex = _marker_of(_JSON_OBJECT.validate_json(request.body)) - return _error(status, f"scripted {status} {marker_hex}") - - return respond - - -def _holding(held: SimpleQueue[str], release: threading.Event) -> Callable[[Request], Reply]: - def respond(request: Request) -> Reply: - held.put(request.target) - assert release.wait(timeout=60), "Held request was never released" - return _peer(request) - - return respond - - def _delivered(wire: Wire) -> tuple[_Delivery, ...]: return tuple( _Delivery(request.target, request.headers.get("anthropic-beta"), _JSON_OBJECT.validate_json(request.body)) @@ -606,28 +492,6 @@ def _chat_body( return body if tools is None else {**body, "tools": tools} -def _effort_body(model: str, marker: str) -> dict[str, JsonValue]: - return {"model": model, "messages": [{"role": "user", "content": marker}], "reasoning_effort": "low"} - - -def _reasoning_fields(delivery: _Delivery) -> dict[str, JsonValue]: - return {field: delivery.body.get(field) for field in _REASONING_FIELDS} - - -def _messages_body(model: str, marker: str, *, stream: bool = False, **extra: JsonValue) -> dict[str, JsonValue]: - return { - "model": model, - "max_tokens": 64, - "messages": [{"role": "user", "content": marker}], - "stream": stream, - **extra, - } - - -def _responses_body(model: str, marker: str) -> dict[str, JsonValue]: - return {"model": model, "input": marker, "text": _TEXT_FORMAT} - - def _chat_message(response: httpx.Response) -> dict[str, JsonValue]: assert response.status_code == 200, response.text choice: Final = object_value(list_value(_JSON_OBJECT.validate_json(response.content)["choices"])[0]) @@ -639,113 +503,6 @@ def _chat_content(response: httpx.Response) -> JsonValue: return json.loads(string_value(_chat_message(response)["content"])) -def _response_id(response: httpx.Response) -> str: - return string_value(_JSON_OBJECT.validate_json(response.content)["id"]) - - -def _delta_content(frame: str) -> str: - choices: Final = list_value(_JSON_OBJECT.validate_json(frame).get("choices", [])) - if not choices: - return "" - content: Final = object_value(object_value(choices[0])["delta"]).get("content") - return "" if content is None else string_value(content) - - -def _content_frames(response: httpx.Response) -> tuple[str, ...]: - assert response.status_code == 200, response.text - assert response.text.rstrip().endswith("data: [DONE]"), response.text - contents: Final = tuple( - _delta_content(line.removeprefix("data: ")) for line in response.text.splitlines() if line.startswith("data: {") - ) - return tuple(content for content in contents if content) - - -def _timing_frames(response: httpx.Response) -> tuple[str, int]: - assert response.status_code == 200, response.text - frames: Final = tuple( - line.removeprefix("data: ") for line in response.text.splitlines() if line.startswith("data: {") - ) - error_frames: Final = sum(1 for frame in frames if "error" in _JSON_OBJECT.validate_json(frame)) - return "".join(_delta_content(frame) for frame in frames), error_frames - - -def _timed_stream( - client: httpx.Client, - key: str, - wire: Wire, - model: str, - shape: str, - response_format: dict[str, JsonValue], - record_property: Callable[[str, object], None], -) -> float: - marker, marker_hex = _marker(shape) - raw: Final = _timing_text(shape) - assert raw is not None - started: Final = time.perf_counter() - response: Final = _post( - client, key, "/v1/chat/completions", _chat_body(model, marker, stream=True, response_format=response_format) - ) - wall: Final = time.perf_counter() - started - text, error_frames = _timing_frames(response) - record_property(f"{shape}_seconds", round(wall, 3)) - record_property(f"{shape}_error_frames", error_frames) - assert text == raw, (shape, len(text), len(raw), text[:80]) - _one_delivery(wire, marker_hex) - return wall - - -def _stream_text(response: httpx.Response) -> str: - return "".join(_content_frames(response)) - - -def _stream_content(response: httpx.Response) -> JsonValue: - return json.loads(_stream_text(response)) - - -def _responses_stream_text(response: httpx.Response) -> JsonValue: - assert response.status_code == 200, response.text - frames: Final = tuple( - _JSON_OBJECT.validate_json(line.removeprefix("data: ")) - for line in response.text.splitlines() - if line.startswith("data: {") - ) - assert any(frame["type"] == "response.completed" for frame in frames), response.text - return json.loads( - "".join(string_value(frame["delta"]) for frame in frames if frame["type"] == "response.output_text.delta") - ) - - -def _responses_text(response: httpx.Response) -> JsonValue: - assert response.status_code == 200, response.text - items: Final = tuple( - object_value(item) for item in list_value(_JSON_OBJECT.validate_json(response.content)["output"]) - ) - (message,) = tuple(item for item in items if item["type"] == "message") - return json.loads(string_value(object_value(list_value(message["content"])[0])["text"])) - - -def _text_block_json(message: Mapping[str, JsonValue]) -> JsonValue: - block: Final = object_value(list_value(message["content"])[0]) - return json.loads(string_value(block["text"])) - - -def _messages_text(response: httpx.Response) -> JsonValue: - assert response.status_code == 200, response.text - return _text_block_json(_JSON_OBJECT.validate_json(response.content)) - - -def _messages_stream_text(response: httpx.Response) -> JsonValue: - assert response.status_code == 200, response.text - assert "event: message_stop" in response.text, response.text - frames: Final = tuple( - _JSON_OBJECT.validate_json(line.removeprefix("data: ")) - for line in response.text.splitlines() - if line.startswith("data: {") - ) - deltas: Final = tuple(object_value(frame["delta"]) for frame in frames if frame["type"] == "content_block_delta") - return json.loads("".join(string_value(delta["text"]) for delta in deltas if delta["type"] == "text_delta")) - - def _vertex_deployment(gateway: Gateway, scenario: Scenario, api_base: str, backend: str) -> str: return scenario.model( model=f"vertex_ai/{backend}", @@ -769,676 +526,6 @@ def _bedrock_invoke_deployment( ) -def _rejects_forced_tools_without_native_output(row: Mapping[str, JsonValue]) -> bool: - return row.get("supports_forced_tool_use") is False and row.get("supports_native_structured_output") is not True - - -def _claude_rejecting_forced_tools(gateway: Gateway, namespace: str, family: str) -> str: - rows: Final = gateway.get("/public/litellm_model_cost_map") - names: Final = sorted( - name.removeprefix(namespace) - for name, row in rows.items() - if name.startswith(namespace + family) and _rejects_forced_tools_without_native_output(object_value(row)) - ) - assert names, f"The cost map has no {namespace}{family} row rejecting forced tools without native output" - return names[0] - - -def _family_deployment(gateway: Gateway, scenario: Scenario, api_base: str, member: str) -> str: - match member: - case "anthropic_unmapped": - return scenario.model(model=f"anthropic/{_UNMAPPED_MODEL}", api_base=api_base, api_key=_ANTHROPIC_KEY) - case "azure_ai_unmapped": - return scenario.model(model=f"azure_ai/{_UNMAPPED_MODEL}", api_base=api_base, api_key=_AZURE_KEY) - case "bedrock_invoke": - return _bedrock_invoke_deployment(scenario, api_base, _BEDROCK_INVOKE_MODEL) - case "bedrock_mantle": - return scenario.model( - model=f"bedrock_mantle/{_BEDROCK_INVOKE_MODEL}", - api_base=api_base, - api_key=_MANTLE_KEY, - aws_region_name="us-east-1", - ) - case "claude_platform_unmapped": - return scenario.model( - model=f"bedrock/claude_platform/{_UNMAPPED_MODEL}", - api_base=api_base, - api_key=_PLATFORM_KEY, - aws_region_name="us-east-1", - workspace_id="scripted-workspace", - ) - case "vertex_haiku": - return _vertex_deployment(gateway, scenario, api_base, _EFFORT_GATED_MODEL) - raise AssertionError(member) - - -def _owned_config(path: Path, gateway: Gateway, api_base: str) -> Path: - config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) - path.write_text( - yaml.safe_dump( - { - **config, - "model_list": [ - { - "model_name": _OWNED_MODEL, - "litellm_params": { - "model": f"vertex_ai/{_TOOL_MODEL}", - "api_base": api_base, - "vertex_project": _PROJECT, - "vertex_location": _LOCATION, - "vertex_credentials": service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), - }, - }, - { - "model_name": _NIM_GROUP, - "litellm_params": { - "model": f"nvidia_nim/{_NIM_GROUP}", - "api_base": api_base, - "api_key": _SCRIPTED_KEY, - }, - }, - ], - } - ) - ) - return path - - -def _pass_through_environment(api_base: str) -> dict[str, str]: - bases: Final = ( - "GEMINI_API_BASE", - "COHERE_API_BASE", - "VLLM_API_BASE", - "MISTRAL_API_BASE", - "TYPESAFE_API_BASE", - "OPENROUTER_API_BASE", - "AZURE_API_BASE", - "OPENAI_API_BASE", - "GIGACHAT_API_BASE", - "WATSONX_API_BASE", - "TINYFISH_AGENT_API_BASE", - ) - keys: Final = ( - "GEMINI_API_KEY", - "COHERE_API_KEY", - "VLLM_API_KEY", - "MISTRAL_API_KEY", - "TYPESAFE_API_KEY", - "OPENROUTER_API_KEY", - "AZURE_API_KEY", - "OPENAI_API_KEY", - "GIGACHAT_ACCESS_TOKEN", - "WATSONX_TOKEN", - "TINYFISH_API_KEY", - ) - return {**{name: api_base for name in bases}, **{name: _SCRIPTED_KEY for name in keys}} - - -@dataclass(frozen=True, slots=True) -class _PassThroughCall: - path: str - body: dict[str, JsonValue] - upstream_target: str - headers: dict[str, str] - - -def _pass_through_call(route: str, marker: str, key: str) -> _PassThroughCall: - chat: Final[dict[str, JsonValue]] = { - "model": "scripted-chat", - "messages": [{"role": "user", "content": marker}], - "stream": True, - } - bearer: Final = {"Authorization": f"Bearer {key}"} - match route: - case "gemini": - return _PassThroughCall( - "/gemini/v1beta/models/scripted-model:streamGenerateContent?alt=sse", - {"contents": [{"parts": [{"text": marker}]}]}, - "/v1beta/models/scripted-model:streamGenerateContent?alt=sse&key=scripted", - {"x-goog-api-key": key}, - ) - case "cohere": - return _PassThroughCall("/cohere/v2/chat", chat, "/v2/chat", bearer) - case "vllm": - return _PassThroughCall("/vllm/v1/chat/completions", chat, "/v1/chat/completions", bearer) - case "mistral": - return _PassThroughCall("/mistral/v1/chat/completions", chat, "/v1/chat/completions", bearer) - case "typesafe": - return _PassThroughCall("/typesafe/v1/chat/completions", chat, "/v1/chat/completions", bearer) - case "openrouter": - return _PassThroughCall("/openrouter/v1/chat/completions", chat, "/v1/chat/completions", bearer) - case "azure": - return _PassThroughCall( - "/azure/openai/deployments/scripted-deployment/chat/completions?api-version=2024-10-21", - chat, - "/openai/deployments/scripted-deployment/chat/completions?api-version=2024-10-21", - bearer, - ) - case "openai": - return _PassThroughCall("/openai/v1/chat/completions", chat, "/v1/chat/completions", bearer) - case "gigachat": - return _PassThroughCall("/gigachat/chat/completions", chat, "/chat/completions", bearer) - case "watsonx": - return _PassThroughCall( - "/watsonx/ml/v1/text/chat_stream", - {**chat, "model_id": "scripted-chat", "project_id": "scripted-project"}, - f"/ml/v1/text/chat_stream?version={_WATSONX_VERSION}", - bearer, - ) - case "tinyfish": - return _PassThroughCall( - "/tinyfish/v1/automation/run-sse", - {"url": "https://example.com", "goal": marker}, - "/v1/automation/run-sse", - bearer, - ) - case "nvidia_nim": - return _PassThroughCall( - f"/nvidia_nim/{_NIM_GROUP}/v1/chat/completions", chat, "/v1/chat/completions", bearer - ) - raise AssertionError(route) - - -def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]: - return tuple( - stack.enter_context(httpx.Client(base_url=base_url, timeout=60, trust_env=False)) for _ in range(count) - ) - - -def _post(client: httpx.Client, key: str, path: str, body: Mapping[str, JsonValue]) -> httpx.Response: - return client.post(path, json=dict(body), headers={"Authorization": f"Bearer {key}"}) - - -def _post_or_dropped(client: httpx.Client, key: str, path: str, body: Mapping[str, JsonValue]) -> httpx.Response | None: - try: - return _post(client, key, path, body) - except httpx.TransportError: - return None - - -def _local_port(client: httpx.Client) -> int: - with client.stream("GET", "/health/liveliness") as response: - port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1]) - response.read() - assert response.status_code == 200, response.text - return port - - -def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]: - return frozenset( - connection.raddr.port - for connection in psutil.Process(pid).net_connections(kind="tcp") - if connection.raddr and connection.laddr.port == proxy_port - ) - - -def _seen_once(delivered: tuple[_Delivery, ...], marker_hexes: tuple[str, ...]) -> None: - bodies: Final = tuple(json.dumps(sent.body) for sent in delivered) - seen: Final = tuple(sum(marker_hex in body for body in bodies) for marker_hex in marker_hexes) - assert seen == (1,) * len(marker_hexes), seen - - -def _pass_through_spend_rows(key: str) -> list[dict[str, JsonValue]]: - return eventually( - lambda: read_rows( - 'SELECT request_id, model, call_type FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND call_type = %s', - (sha256(key.encode()).hexdigest(), "pass_through_endpoint"), - ), - lambda rows: len(rows) == 1, - seconds=90, - ) - - -def _spend_rows(request_id: str) -> list[dict[str, JsonValue]]: - return eventually( - lambda: read_rows('SELECT request_id, model FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)), - lambda rows: len(rows) == 1, - seconds=70, - ) - - -def _openai_client(gateway: Gateway) -> openai.OpenAI: - return openai.OpenAI( - base_url=f"{gateway.client.base_url}/v1", - api_key=gateway.key, - max_retries=0, - http_client=httpx.Client(trust_env=False, timeout=60), - ) - - -@pytest.fixture(scope="module") -def pass_through_proxy( - session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory -) -> Iterator[tuple[OwnedProxy, Wire]]: - directory: Final = tmp_path_factory.mktemp("structured-output-pass-through") - with wire_server(_relay_peer) as wire: - config: Final = _owned_config(directory / "pass-through.yaml", session_gateway, wire.url) - with owned_proxy_process( - session_gateway, directory, _pass_through_environment(wire.url), config=config - ) as owned: - yield owned, wire - - -def _anthropic_client(gateway: Gateway) -> anthropic.Anthropic: - return anthropic.Anthropic( - base_url=str(gateway.client.base_url), - api_key=gateway.key, - max_retries=0, - http_client=httpx.Client(trust_env=False, timeout=60), - ) - - -@pytest.mark.parametrize( - "shape", - [ - "plain", - "wrapped", - "named", - "values", - "large", - "two_candidates", - "inner_incomplete", - "self_shaped", - "scalar", - "empty", - "list_input", - "deep", - "parameter_label", - "function_label", - "label_unshared", - "label_two_objects", - "label_encoded", - "label_then_declared", - "named_encoded", - "undeclared_encoded", - "encoded_deep", - "number_first", - "bool_first", - "null_first", - "label_non_ascii", - "label_two_objects_non_ascii", - ], -) -def test_vertex_chat_json_schema_content_is_the_schema_object(gateway: Gateway, shape: str) -> None: - marker, marker_hex = _marker(shape) - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) - assert _chat_content(response) == _expected(shape, marker_hex), response.text - assert "tool_calls" not in _chat_message(response) or _chat_message(response)["tool_calls"] is None - sent: Final = _one_delivery(wire, marker_hex) - assert sent.target == f"{_VERTEX_MODELS}{_TOOL_MODEL}:rawPredict", sent.target - assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body - assert "output_config" not in sent.body and "output_format" not in sent.body, sent.body - - -def test_vertex_chat_through_the_openai_sdk_unwraps_the_envelope(gateway: Gateway) -> None: - marker, marker_hex = _marker("wrapped") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - with _openai_client(gateway) as client: - reply: Final = client.chat.completions.create( - model=model, - messages=[{"role": "user", "content": marker}], - response_format={"type": "json_schema", "json_schema": {"name": "risk", "schema": _SCHEMA}}, - ) - assert json.loads(reply.choices[0].message.content or "") == _payload(marker_hex), reply - _one_delivery(wire, marker_hex) - - -def test_vertex_chat_through_the_async_openai_sdk_unwraps_the_envelope(gateway: Gateway) -> None: - marker, marker_hex = _marker("wrapped") - - async def generate(model: str) -> str | None: - async with openai.AsyncOpenAI( - base_url=f"{gateway.client.base_url}/v1", - api_key=gateway.key, - max_retries=0, - http_client=httpx.AsyncClient(trust_env=False, timeout=60), - ) as client: - reply: Final = await client.chat.completions.create( - model=model, - messages=[{"role": "user", "content": marker}], - response_format={"type": "json_schema", "json_schema": {"name": "risk", "schema": _SCHEMA}}, - ) - return reply.choices[0].message.content - - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - content: Final = asyncio.run(generate(model)) - assert json.loads(content or "") == _payload(marker_hex), content - _one_delivery(wire, marker_hex) - - -def test_vertex_responses_json_schema_output_text_is_the_schema_object(gateway: Gateway) -> None: - marker, marker_hex = _marker("wrapped") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request("POST", "/v1/responses", _responses_body(model, marker)) - assert _responses_text(response) == _payload(marker_hex), response.text - sent: Final = _one_delivery(wire, marker_hex) - assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body - - -@pytest.mark.parametrize("shape", _STREAM_SHAPES) -def test_vertex_chat_stream_json_schema_content_is_the_schema_object(gateway: Gateway, shape: str) -> None: - marker, marker_hex = _marker(shape) - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=True)) - frames: Final = _content_frames(response) - assert json.loads("".join(frames)) == _expected(shape, marker_hex), response.text - if shape in _HELD_STREAM_SHAPES: - assert len(frames) == 1, frames - else: - assert len(frames) >= _fragment_count(shape, marker_hex), frames - sent: Final = _one_delivery(wire, marker_hex) - assert sent.target == f"{_VERTEX_MODELS}{_TOOL_MODEL}{_VERTEX_STREAM_SUFFIX}", sent.target - assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body - - -def test_vertex_chat_stream_label_before_a_declared_key_keeps_streaming(gateway: Gateway) -> None: - marker, marker_hex = _marker("label_then_declared") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=True)) - frames: Final = _content_frames(response) - assert json.loads("".join(frames)) == _tool_input("label_then_declared", marker_hex), response.text - assert len(frames) > 1, frames - _one_delivery(wire, marker_hex) - - -@pytest.mark.parametrize("shape", _AS_WRITTEN_SHAPES) -def test_vertex_chat_stream_held_answer_without_an_envelope_is_sent_as_the_model_wrote_it( - gateway: Gateway, shape: str -) -> None: - marker, marker_hex = _marker(shape) - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=True)) - assert _content_frames(response) == (_tool_text(shape, marker_hex),), response.text - _one_delivery(wire, marker_hex) - - -@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) -@pytest.mark.parametrize("schema", [_OPTIONAL_SCHEMA, {"type": "object"}], ids=["no_required", "no_properties"]) -@pytest.mark.parametrize( - ("shape", "unwrapped"), [("wrapped", True), ("parameter_label", False)], ids=["wrapped", "label"] -) -def test_vertex_chat_json_schema_without_required_keys_unwraps_only_the_tool_named_envelope( - gateway: Gateway, schema: dict[str, JsonValue], shape: str, unwrapped: bool, stream: bool -) -> None: - marker, marker_hex = _marker(shape) - response_format: Final[dict[str, JsonValue]] = { - "type": "json_schema", - "json_schema": {"name": "risk", "schema": schema}, - } - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request( - "POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream, response_format=response_format) - ) - content: Final = _stream_content(response) if stream else _chat_content(response) - assert content == (_payload(marker_hex) if unwrapped else _tool_input(shape, marker_hex)), response.text - if stream: - assert (len(_content_frames(response)) == 1) is unwrapped, response.text - _one_delivery(wire, marker_hex) - - -@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) -def test_vertex_chat_json_schema_that_declares_the_wrapper_key_keeps_the_answer_whole( - gateway: Gateway, stream: bool -) -> None: - marker, marker_hex = _marker("values") - response_format: Final[dict[str, JsonValue]] = { - "type": "json_schema", - "json_schema": {"name": "valued", "schema": _VALUES_SCHEMA}, - } - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request( - "POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream, response_format=response_format) - ) - content: Final = _stream_content(response) if stream else _chat_content(response) - assert content == _tool_input("values", marker_hex), response.text - if stream: - assert len(_content_frames(response)) >= _fragment_count("values", marker_hex), response.text - _one_delivery(wire, marker_hex) - - -@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) -@pytest.mark.parametrize( - ("schema", "wrapper_unwrapped"), - [ - (_ADDITIONAL_OBJECT_SCHEMA, False), - (_PATTERN_SCHEMA, False), - (_ADDITIONAL_TRUE_SCHEMA, False), - (_CLOSED_OPTIONAL_SCHEMA, True), - ], - ids=["additional_object", "pattern_properties", "additional_true", "additional_false"], -) -@pytest.mark.parametrize("shape", ["undeclared_wrapper", "wrapped"]) -def test_vertex_chat_json_schema_that_allows_undeclared_keys_keeps_an_undeclared_wrapper_whole( - gateway: Gateway, shape: str, schema: dict[str, JsonValue], wrapper_unwrapped: bool, stream: bool -) -> None: - marker, marker_hex = _marker(shape) - unwrapped: Final = shape == "wrapped" or wrapper_unwrapped - response_format: Final[dict[str, JsonValue]] = { - "type": "json_schema", - "json_schema": {"name": "dynamic", "schema": schema}, - } - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request( - "POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream, response_format=response_format) - ) - content: Final = _stream_content(response) if stream else _chat_content(response) - assert content == (_payload(marker_hex) if unwrapped else _tool_input(shape, marker_hex)), response.text - if stream and unwrapped: - assert len(_content_frames(response)) == 1, response.text - elif stream: - assert len(_content_frames(response)) >= _fragment_count(shape, marker_hex), response.text - _one_delivery(wire, marker_hex) - - -def test_vertex_responses_stream_json_schema_output_text_is_the_schema_object(gateway: Gateway) -> None: - marker, marker_hex = _marker("wrapped") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request("POST", "/v1/responses", {**_responses_body(model, marker), "stream": True}) - assert _responses_stream_text(response) == _payload(marker_hex), response.text - sent: Final = _one_delivery(wire, marker_hex) - assert sent.target == f"{_VERTEX_MODELS}{_TOOL_MODEL}{_VERTEX_STREAM_SUFFIX}", sent.target - - -@pytest.mark.parametrize("shape", ["other_tool", "named_input", "named_two_objects", "named_string", "named_partial"]) -def test_vertex_chat_json_schema_unwraps_a_tool_named_envelope_only_around_one_object( - gateway: Gateway, shape: str -) -> None: - marker, marker_hex = _marker(shape) - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) - assert _chat_content(response) == _expected(shape, marker_hex), response.text - _one_delivery(wire, marker_hex) - - -@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) -@pytest.mark.parametrize("shape", ["name_property", "name_value_plain"]) -def test_vertex_chat_json_schema_with_a_name_property_unwraps_only_the_tool_named_wrapper( - gateway: Gateway, shape: str, stream: bool -) -> None: - marker, marker_hex = _marker(shape) - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request( - "POST", - "/v1/chat/completions", - _chat_body(model, marker, stream=stream, response_format=_NAME_RESPONSE_FORMAT), - ) - content: Final = _stream_content(response) if stream else _chat_content(response) - assert content == _expected(shape, marker_hex), response.text - if stream and shape == "name_property": - assert len(_content_frames(response)) == 1, response.text - sent: Final = _one_delivery(wire, marker_hex) - assert _json_tool_schema(sent)["properties"] == _NAME_SCHEMA["properties"], sent.body - - -@pytest.mark.parametrize("shape", ["wrapped", "inner_incomplete"]) -def test_vertex_chat_json_schema_validation_runs_on_the_unwrapped_object(gateway: Gateway, shape: str) -> None: - marker, marker_hex = _marker(shape) - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request( - "POST", "/v1/chat/completions", {**_chat_body(model, marker), "enable_json_schema_validation": True} - ) - if shape == "wrapped": - assert _chat_content(response) == _payload(marker_hex), response.text - else: - assert response.status_code >= 400, response.text - assert "JSONSchemaValidationError" in response.text, response.text - _one_delivery(wire, marker_hex) - assert gateway.request("GET", "/health/liveliness").status_code == 200 - - -def test_vertex_chat_json_schema_next_to_a_user_tool_unwraps_the_envelope_and_keeps_the_tool_call( - gateway: Gateway, -) -> None: - marker, marker_hex = _marker("wrapped") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request( - "POST", "/v1/chat/completions", _chat_body(model, marker, tools=[_USER_TOOL_SPEC]) - ) - assert response.status_code == 200, response.text - message: Final = object_value( - object_value(list_value(_JSON_OBJECT.validate_json(response.content)["choices"])[0])["message"] - ) - assert json.loads(string_value(message["content"])) == _payload(marker_hex), response.text - (tool_call,) = tuple(object_value(call) for call in list_value(message["tool_calls"])) - function: Final = object_value(tool_call["function"]) - assert function["name"] == _USER_TOOL, response.text - assert json.loads(string_value(function["arguments"])) == {"q": marker_hex}, response.text - sent: Final = _one_delivery(wire, marker_hex) - assert _tool_names(sent.body) == frozenset({_USER_TOOL, _JSON_TOOL}), sent.body - - -@pytest.mark.parametrize("member", _FAMILY) -def test_tool_path_family_unwraps_the_envelope(gateway: Gateway, member: str) -> None: - marker, marker_hex = _marker("wrapped") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _family_deployment(gateway, scenario, wire.url, member) - response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) - assert _chat_content(response) == _payload(marker_hex), response.text - sent: Final = _one_delivery(wire, marker_hex) - assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body - - -@pytest.mark.parametrize("member", _FAMILY) -def test_tool_path_family_stream_unwraps_the_envelope(gateway: Gateway, member: str) -> None: - marker, marker_hex = _marker("wrapped") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _family_deployment(gateway, scenario, wire.url, member) - response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=True)) - frames: Final = _content_frames(response) - assert frames, response.text - assert json.loads("".join(frames)) == _payload(marker_hex), response.text - assert len(frames) == 1, frames - sent: Final = _one_delivery(wire, marker_hex) - assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body - - -@pytest.mark.parametrize("accepts_forced_tools", [True, False], ids=["accepts_forced_tools", "rejects_forced_tools"]) -def test_vertex_chat_json_schema_forces_the_json_tool_only_when_the_model_accepts_forced_tools( - gateway: Gateway, accepts_forced_tools: bool -) -> None: - marker, marker_hex = _marker("plain") - backend: Final = ( - _TOOL_MODEL if accepts_forced_tools else _claude_rejecting_forced_tools(gateway, "vertex_ai/", "claude") - ) - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, backend) - response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) - assert _chat_content(response) == _payload(marker_hex), response.text - sent: Final = _one_delivery(wire, marker_hex) - assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body - assert sent.body.get("tool_choice") == (_FORCED_JSON_TOOL if accepts_forced_tools else None), sent.body - assert "output_config" not in sent.body and "output_format" not in sent.body, sent.body - - -def test_vertex_chat_json_schema_on_a_model_rejecting_forced_tools_maps_reasoning_effort_like_plain_chat( - gateway: Gateway, -) -> None: - backend: Final = _claude_rejecting_forced_tools(gateway, "vertex_ai/", "claude") - structured_marker, structured_hex = _marker("plain") - plain_marker, plain_hex = _marker("plain") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, backend) - structured: Final = gateway.request( - "POST", - "/v1/chat/completions", - {**_effort_body(model, structured_marker), "response_format": _RESPONSE_FORMAT}, - ) - assert _chat_content(structured) == _payload(structured_hex), structured.text - structured_sent: Final = _one_delivery(wire, structured_hex) - plain: Final = gateway.request("POST", "/v1/chat/completions", _effort_body(model, plain_marker)) - assert plain.status_code == 200, plain.text - plain_sent: Final = _one_delivery(wire, plain_hex) - assert "thinking" in plain_sent.body, plain_sent.body - assert _reasoning_fields(structured_sent) == _reasoning_fields(plain_sent), (structured_sent.body, plain_sent.body) - - -def test_bedrock_invoke_json_schema_sends_the_json_tool_unforced_when_the_model_rejects_forced_tools( - gateway: Gateway, -) -> None: - marker, marker_hex = _marker("plain") - backend: Final = _claude_rejecting_forced_tools(gateway, "", "anthropic.claude") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _bedrock_invoke_deployment(scenario, wire.url, backend) - response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) - assert _chat_content(response) == _payload(marker_hex), response.text - sent: Final = _one_delivery(wire, marker_hex) - assert sent.target == f"/model/{backend}/invoke", sent.target - assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body - assert "tool_choice" not in sent.body, sent.body - - -@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) -def test_vertex_chat_json_schema_native_folds_output_format_into_output_config(gateway: Gateway, stream: bool) -> None: - marker, marker_hex = _marker("native") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) - response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream)) - content: Final = _stream_content(response) if stream else _chat_content(response) - assert content == _payload(marker_hex), response.text - sent: Final = _one_delivery(wire, marker_hex) - assert _folded_format(sent) == _OUTPUT_FORMAT, sent.body - assert "tools" not in sent.body, sent.body - - -def test_vertex_chat_native_through_the_openai_sdk_folds_output_format(gateway: Gateway) -> None: - marker, marker_hex = _marker("native") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) - with _openai_client(gateway) as client: - reply: Final = client.chat.completions.create( - model=model, - messages=[{"role": "user", "content": marker}], - response_format={"type": "json_schema", "json_schema": {"name": "risk", "schema": _SCHEMA}}, - ) - assert json.loads(reply.choices[0].message.content or "") == _payload(marker_hex), reply - assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT - - -def test_vertex_responses_native_folds_output_format(gateway: Gateway) -> None: - marker, marker_hex = _marker("native") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) - response: Final = gateway.request("POST", "/v1/responses", _responses_body(model, marker)) - assert _responses_text(response) == _payload(marker_hex), response.text - assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT - - def _native_deployment(gateway: Gateway, scenario: Scenario, api_base: str, provider: str, backend: str) -> str: match provider: case "vertex_ai": @@ -1587,580 +674,3 @@ def test_vertex_chat_native_schema_keeps_min_items_one_and_closes_a_nullable_obj sent: Final = _one_delivery(wire, marker_hex) assert _folded_format(sent) == {"type": "json_schema", "schema": expected}, sent.body assert "tools" not in sent.body, sent.body - - -@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) -def test_vertex_messages_folds_output_format_into_output_config(gateway: Gateway, stream: bool) -> None: - marker, marker_hex = _marker("native") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) - response: Final = gateway.request( - "POST", - "/v1/messages", - _messages_body(model, marker, stream=stream, output_format=_OUTPUT_FORMAT), - headers={"anthropic-beta": _STRUCTURED_OUTPUTS_BETA}, - ) - text: Final = _messages_stream_text(response) if stream else _messages_text(response) - assert text == _payload(marker_hex), response.text - assert f"msg_{marker_hex}" in response.text, response.text - sent: Final = _one_delivery(wire, marker_hex) - assert _folded_format(sent) == _OUTPUT_FORMAT, sent.body - assert sent.target.endswith(_VERTEX_STREAM_SUFFIX if stream else ":rawPredict"), sent.target - - -def test_vertex_messages_through_the_anthropic_sdk_folds_output_format(gateway: Gateway) -> None: - marker, marker_hex = _marker("native") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) - with _anthropic_client(gateway) as client: - reply: Final = client.beta.messages.create( - model=model, - max_tokens=64, - messages=[{"role": "user", "content": marker}], - output_format={"type": "json_schema", "schema": _SCHEMA}, - betas=[_STRUCTURED_OUTPUTS_BETA], - ) - assert reply.id == f"msg_{marker_hex}", reply - assert _text_block_json(_JSON_OBJECT.validate_json(reply.model_dump_json())) == _payload(marker_hex), reply - assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT - - -def test_vertex_messages_through_the_async_anthropic_sdk_folds_output_format(gateway: Gateway) -> None: - marker, marker_hex = _marker("native") - - async def generate(model: str) -> tuple[str, str]: - async with anthropic.AsyncAnthropic( - base_url=str(gateway.client.base_url), - api_key=gateway.key, - max_retries=0, - http_client=httpx.AsyncClient(trust_env=False, timeout=60), - ) as client: - reply: Final = await client.beta.messages.create( - model=model, - max_tokens=64, - messages=[{"role": "user", "content": marker}], - output_format={"type": "json_schema", "schema": _SCHEMA}, - betas=[_STRUCTURED_OUTPUTS_BETA], - ) - return reply.id, reply.model_dump_json() - - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) - reply_id, dumped = asyncio.run(generate(model)) - assert reply_id == f"msg_{marker_hex}", dumped - assert _text_block_json(_JSON_OBJECT.validate_json(dumped)) == _payload(marker_hex), dumped - assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT - - -@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) -@pytest.mark.parametrize( - ("callers_config", "expected_output_config"), - [ - ({"format": _CALLERS_FORMAT}, {"format": _OUTPUT_FORMAT}), - ({"effort": "high", "format": _CALLERS_FORMAT}, {"effort": "high", "format": _OUTPUT_FORMAT}), - ], - ids=["format", "effort_and_format"], -) -def test_vertex_messages_output_format_wins_over_the_callers_output_config_format( - gateway: Gateway, callers_config: dict[str, JsonValue], expected_output_config: dict[str, JsonValue], stream: bool -) -> None: - marker, marker_hex = _marker("native") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) - response: Final = gateway.request( - "POST", - "/v1/messages", - _messages_body(model, marker, stream=stream, output_format=_OUTPUT_FORMAT, output_config=callers_config), - ) - text: Final = _messages_stream_text(response) if stream else _messages_text(response) - assert text == _payload(marker_hex), response.text - sent: Final = _one_delivery(wire, marker_hex) - assert "output_format" not in sent.body, sent.body - assert sent.body["output_config"] == expected_output_config, sent.body - - -@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) -@pytest.mark.parametrize( - "schema", - [_ADDITIONAL_OBJECT_SCHEMA, _PATTERN_SCHEMA, _ADDITIONAL_TRUE_SCHEMA], - ids=["additional_object", "pattern_properties", "additional_true"], -) -def test_vertex_messages_json_schema_that_allows_undeclared_keys_keeps_the_answer_whole( - gateway: Gateway, schema: dict[str, JsonValue], stream: bool -) -> None: - marker, marker_hex = _marker("undeclared_wrapper") - output_format: Final[dict[str, JsonValue]] = {"type": "json_schema", "schema": schema} - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) - response: Final = gateway.request( - "POST", - "/v1/messages", - _messages_body(model, marker, stream=stream, output_format=output_format), - headers={"anthropic-beta": _STRUCTURED_OUTPUTS_BETA}, - ) - text: Final = _messages_stream_text(response) if stream else _messages_text(response) - assert text == _tool_input("undeclared_wrapper", marker_hex), response.text - assert _folded_format(_one_delivery(wire, marker_hex)) == output_format - - -def test_vertex_messages_replaces_a_non_dict_output_config_with_the_folded_format(gateway: Gateway) -> None: - marker, marker_hex = _marker("native") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) - response: Final = gateway.request( - "POST", - "/v1/messages", - _messages_body(model, marker, output_format=_OUTPUT_FORMAT, output_config="bogus"), - ) - assert _messages_text(response) == _payload(marker_hex), response.text - sent: Final = _one_delivery(wire, marker_hex) - assert sent.body["output_config"] == {"format": _OUTPUT_FORMAT}, sent.body - - -def test_vertex_messages_drops_effort_but_keeps_the_folded_format_on_haiku(gateway: Gateway) -> None: - marker, marker_hex = _marker("native") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _EFFORT_GATED_MODEL) - response: Final = gateway.request( - "POST", - "/v1/messages", - _messages_body(model, marker, output_format=_OUTPUT_FORMAT, output_config={"effort": "high"}), - ) - assert _messages_text(response) == _payload(marker_hex), response.text - sent: Final = _one_delivery(wire, marker_hex) - assert sent.body["output_config"] == {"format": _OUTPUT_FORMAT}, sent.body - - -def test_vertex_messages_forwards_an_explicit_output_config_format_unchanged(gateway: Gateway) -> None: - marker, marker_hex = _marker("native") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) - response: Final = gateway.request( - "POST", "/v1/messages", _messages_body(model, marker, output_config={"format": _OUTPUT_FORMAT}) - ) - assert _messages_text(response) == _payload(marker_hex), response.text - sent: Final = _one_delivery(wire, marker_hex) - assert sent.body["output_config"] == {"format": _OUTPUT_FORMAT}, sent.body - - -@pytest.mark.parametrize( - ("extra", "expected_output_config"), - [ - ({"output_format": _OUTPUT_FORMAT, "output_config": None}, {"format": _OUTPUT_FORMAT}), - ({"output_format": _OUTPUT_FORMAT, "output_config": {"format": None}}, {"format": _OUTPUT_FORMAT}), - ({"output_format": _OUTPUT_FORMAT, "output_config": {}}, {"format": _OUTPUT_FORMAT}), - ( - {"output_format": _OUTPUT_FORMAT, "output_config": {"effort": "high"}}, - {"effort": "high", "format": _OUTPUT_FORMAT}, - ), - ({"output_format": None, "output_config": {"format": _OUTPUT_FORMAT}}, {"format": _OUTPUT_FORMAT}), - ({"output_format": None}, None), - ({"output_config": None}, None), - ], - ids=[ - "null_config", - "null_format_key", - "empty_config", - "effort_config", - "null_output_format_beside_callers", - "null_output_format_alone", - "null_config_alone", - ], -) -def test_vertex_messages_output_format_beside_an_explicit_output_config( - gateway: Gateway, extra: dict[str, JsonValue], expected_output_config: JsonValue -) -> None: - marker, marker_hex = _marker("native") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) - response: Final = gateway.request("POST", "/v1/messages", _messages_body(model, marker, **extra)) - if expected_output_config is None: - assert response.status_code == 200, response.text - assert f"plain text {marker_hex}" in response.text, response.text - else: - assert _messages_text(response) == _payload(marker_hex), response.text - sent: Final = _one_delivery(wire, marker_hex) - assert "output_format" not in sent.body, sent.body - if expected_output_config is None: - assert "output_config" not in sent.body, sent.body - else: - assert sent.body["output_config"] == expected_output_config, sent.body - - -def test_anthropic_messages_forwards_top_level_output_format_verbatim(gateway: Gateway) -> None: - marker, marker_hex = _marker("native") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = scenario.model(model=f"anthropic/{_NATIVE_MODEL}", api_base=wire.url, api_key=_ANTHROPIC_KEY) - response: Final = gateway.request( - "POST", - "/v1/messages", - _messages_body(model, marker, output_format=_OUTPUT_FORMAT), - headers={"anthropic-beta": _STRUCTURED_OUTPUTS_BETA}, - ) - assert _messages_text(response) == _payload(marker_hex), response.text - sent: Final = _one_delivery(wire, marker_hex) - assert sent.target == "/v1/messages", sent.target - assert sent.body["output_format"] == _OUTPUT_FORMAT, sent.body - assert "output_config" not in sent.body, sent.body - - -def test_vertex_chat_without_response_format_sends_no_tool_and_no_output_config(gateway: Gateway) -> None: - marker, marker_hex = _marker("plain_text") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request( - "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]} - ) - assert response.status_code == 200, response.text - assert f"plain text {marker_hex}" in response.text, response.text - sent: Final = _one_delivery(wire, marker_hex) - assert not {"tools", "output_config", "output_format"} & sent.body.keys(), sent.body - - -@pytest.mark.parametrize( - "output_format", - [5, [], {}, "x" * 5000], - ids=["int", "list", "empty", "5kb"], -) -def test_vertex_messages_malformed_output_format_reaches_the_caller_as_the_peer_400( - gateway: Gateway, output_format: JsonValue -) -> None: - marker, marker_hex = _marker("native") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) - response: Final = gateway.request( - "POST", "/v1/messages", _messages_body(model, marker, output_format=output_format) - ) - assert response.status_code == 400, response.text - assert _BAD_FORMAT in response.text, response.text - sent: Final = _one_delivery(wire, marker_hex) - assert "output_format" not in sent.body, sent.body - assert object_value(sent.body["output_config"])["format"] == output_format, sent.body - assert gateway.request("GET", "/health/liveliness").status_code == 200 - - -def test_vertex_messages_duplicate_output_format_keys_fold_the_last_value(gateway: Gateway) -> None: - marker, marker_hex = _marker("native") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) - raw: Final = ( - json.dumps(_messages_body(model, marker, output_format=5))[:-1] - + f', "output_format": {json.dumps(_OUTPUT_FORMAT)}}}' - ) - response: Final = gateway.client.post( - "/v1/messages", - content=raw, - headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, - ) - assert _messages_text(response) == _payload(marker_hex), response.text - assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT - - -@pytest.mark.parametrize( - "schema", - [ - {"type": "object", "properties": ["tier", "summary"]}, - {"type": "object", "properties": "tier"}, - {"type": "object", "properties": None}, - ], - ids=["properties_list", "properties_string", "properties_null"], -) -@pytest.mark.parametrize("shape", ["wrapped", "other_tool"]) -@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) -def test_vertex_chat_malformed_response_format_schema_keeps_the_tool_input( - gateway: Gateway, schema: JsonValue, shape: str, stream: bool -) -> None: - marker, marker_hex = _marker(shape) - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request( - "POST", - "/v1/chat/completions", - _chat_body( - model, - marker, - stream=stream, - response_format={"type": "json_schema", "json_schema": {"name": "risk", "schema": schema}}, - ), - ) - content: Final = _stream_content(response) if stream else _chat_content(response) - assert content == _tool_input(shape, marker_hex), response.text - _one_delivery(wire, marker_hex) - assert gateway.request("GET", "/health/liveliness").status_code == 200 - - -@pytest.mark.parametrize("status", [400, 403, 404]) -def test_vertex_chat_peer_error_reaches_the_caller_and_the_deployment_keeps_serving( - gateway: Gateway, status: int -) -> None: - marker, marker_hex = _marker("wrapped") - after, after_hex = _marker("wrapped") - port: Final = allocate_port() - with gateway.scenario() as scenario: - with wire_server(_failing(status), port=port) as failing: - model: Final = _vertex_deployment(gateway, scenario, failing.url, _TOOL_MODEL) - response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) - assert response.status_code == status, response.text - assert f"scripted {status} {marker_hex}" in response.text, response.text - _one_delivery(failing, marker_hex) - with wire_server(_peer, port=port) as wire: - recovered: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, after)) - assert _chat_content(recovered) == _payload(after_hex), recovered.text - _one_delivery(wire, after_hex) - - -def test_vertex_chat_unauthenticated_request_never_reaches_the_peer(gateway: Gateway) -> None: - marker, _ = _marker("wrapped") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - response: Final = gateway.request( - "POST", "/v1/chat/completions", _chat_body(model, marker), key="sk-not-a-key-this-proxy-issued" - ) - assert response.status_code == 401, response.text - assert wire.drain() == () - - -def test_vertex_chat_repeated_requests_each_reach_the_peer_and_log_spend_once(gateway: Gateway) -> None: - markers: Final = tuple(_marker("wrapped") for _ in range(2)) - with wire_server(_peer) as wire, gateway.scenario() as scenario: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - responses: Final = tuple( - gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) for marker, _ in markers - ) - for response, (_, marker_hex) in zip(responses, markers, strict=True): - assert _chat_content(response) == _payload(marker_hex), response.text - delivered: Final = _delivered(wire) - assert len(delivered) == 2, delivered - _seen_once(delivered, tuple(marker_hex for _, marker_hex in markers)) - request_ids: Final = tuple(_response_id(response) for response in responses) - assert len(set(request_ids)) == 2, request_ids - for request_id in request_ids: - assert _spend_rows(request_id)[0]["request_id"] == request_id - - -def test_concurrent_burst_across_both_workers_unwraps_every_envelope(gateway: Gateway) -> None: - with ExitStack() as stack: - wire: Final = stack.enter_context(wire_server(_peer)) - scenario: Final = stack.enter_context(gateway.scenario()) - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - shapes: Final = tuple(_BURST_SHAPES[index % len(_BURST_SHAPES)] for index in range(12)) - markers: Final = tuple(_marker(shape) for shape in shapes) - streams: Final = tuple(index % 3 == 2 for index in range(len(markers))) - bodies: Final = tuple( - _chat_body(model, marker, stream=stream) for (marker, _), stream in zip(markers, streams, strict=True) - ) - clients: Final = _clients(stack, str(gateway.client.base_url), len(bodies)) - pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(bodies))) - responses: Final = tuple( - pool.map( - lambda pair: _post(pair[0], gateway.key, "/v1/chat/completions", pair[1]), - zip(clients, bodies, strict=True), - ) - ) - for response, shape, (_, marker_hex), stream in zip(responses, shapes, markers, streams, strict=True): - observed: Final = _stream_content(response) if stream else _chat_content(response) - expected: Final = _expected(shape, marker_hex) - assert observed == expected, (shape, stream, response.text) - _seen_once(_delivered(wire), tuple(marker_hex for _, marker_hex in markers)) - - -def test_upstream_outage_mid_run_recovers_with_envelopes_still_unwrapped(gateway: Gateway) -> None: - with ExitStack() as stack: - clients: Final = _clients(stack, str(gateway.client.base_url), 8) - pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) - scenario: Final = stack.enter_context(gateway.scenario()) - - streams: Final = tuple(index % 2 == 1 for index in range(len(clients))) - - def wave() -> tuple[tuple[httpx.Response | None, str, bool], ...]: - markers: Final = tuple(_marker("wrapped") for _ in clients) - responses: Final = tuple( - pool.map( - lambda item: _post_or_dropped( - item[0], gateway.key, "/v1/chat/completions", _chat_body(model, item[1][0], stream=item[2]) - ), - zip(clients, markers, streams, strict=True), - ) - ) - return tuple(zip(responses, tuple(marker_hex for _, marker_hex in markers), streams, strict=True)) - - def unwrapped(wire: Wire, results: tuple[tuple[httpx.Response | None, str, bool], ...]) -> None: - for response, marker_hex, stream in results: - assert response is not None, results - observed: Final = _stream_content(response) if stream else _chat_content(response) - assert observed == _payload(marker_hex), response.text - _seen_once(_delivered(wire), tuple(marker_hex for _, marker_hex, _ in results)) - - port: Final = allocate_port() - with wire_server(_peer, port=port) as wire: - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - unwrapped(wire, wave()) - outage: Final = wave() - assert all(response is not None and response.status_code >= 500 for response, _, _ in outage), outage - assert gateway.request("GET", "/health/liveliness").status_code == 200 - with wire_server(_peer, port=port) as revived: - unwrapped(revived, wave()) - - -@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) -def test_worker_sigkill_mid_burst_leaves_the_replacement_unwrapping_envelopes(gateway: Gateway, tmp_path: Path) -> None: - held: Final[SimpleQueue[str]] = SimpleQueue() - release: Final = threading.Event() - with ExitStack() as stack: - wire: Final = stack.enter_context(wire_server(_holding(held, release))) - stack.callback(release.set) - config: Final = _owned_config(tmp_path / "structured-output-worker-kill.yaml", gateway, wire.url) - owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2)) - proxy_url: Final = owned.gateway.client.base_url - workers: Final = eventually( - lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), - lambda pids: len(pids) == 2, - seconds=30, - ) - clients: Final = _clients(stack, str(proxy_url), 12) - pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) - ports: Final = tuple(_local_port(client) for client in clients) - markers: Final = tuple(_marker("wrapped") for _ in clients) - futures: Final = tuple( - pool.submit(_post_or_dropped, client, gateway.key, "/v1/chat/completions", _chat_body(_OWNED_MODEL, marker)) - for client, (marker, _) in zip(clients, markers, strict=True) - ) - eventually(held.qsize, lambda size: size == len(clients), seconds=30) - shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in workers} - assert sum(map(len, shares.values())) == len(clients), shares - victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) - psutil.Process(victim).send_signal(signal.SIGKILL) - release.set() - results: Final = tuple(future.result(timeout=60) for future in futures) - for port, result, (_, marker_hex) in zip(ports, results, markers, strict=True): - if port in shares[victim]: - assert result is None, (port, result, shares) - else: - assert result is not None and _chat_content(result) == _payload(marker_hex), (port, result, shares) - eventually(lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda started: started >= 3, 120) - second_wave: Final = tuple(_marker("wrapped") for _ in range(6)) - for marker, marker_hex in second_wave: - after: Final = _post( - owned.gateway.client, gateway.key, "/v1/chat/completions", _chat_body(_OWNED_MODEL, marker) - ) - assert _chat_content(after) == _payload(marker_hex), after.text - _seen_once(_delivered(wire), tuple(marker_hex for _, marker_hex in markers + second_wave)) - assert owned.process.poll() is None - - -@pytest.mark.parametrize( - ("shape", "response_format"), - [ - ("unclosed_key", _RESPONSE_FORMAT), - ("unclosed_name", _NAME_RESPONSE_FORMAT), - ("unclosed_label", _RESPONSE_FORMAT), - ("unclosed_label_key", _RESPONSE_FORMAT), - ], - ids=["unclosed_key", "unclosed_name", "unclosed_label", "unclosed_label_key"], -) -def test_vertex_chat_stream_never_closing_tool_input_stays_linear_and_keeps_health_responsive( - gateway: Gateway, shape: str, response_format: dict[str, JsonValue], record_property: Callable[[str, object], None] -) -> None: - latencies: Final[SimpleQueue[float]] = SimpleQueue() - done: Final = threading.Event() - with ExitStack() as stack: - wire: Final = stack.enter_context(wire_server(_peer)) - scenario: Final = stack.enter_context(gateway.scenario()) - model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) - streaming, probing = _clients(stack, str(gateway.client.base_url), 2) - _local_port(probing) - - def probe() -> None: - while not done.wait(0.05): - started: Final = time.perf_counter() - assert probing.get("/health/liveliness").status_code == 200 - latencies.put(time.perf_counter() - started) - - prober: Final = threading.Thread(target=probe) - prober.start() - stack.callback(prober.join, 30) - stack.callback(done.set) - small, large = ( - _timed_stream(streaming, gateway.key, wire, model, size_shape, response_format, record_property) - for size_shape in (f"{shape}_small", shape) - ) - done.set() - prober.join(timeout=30) - probes: Final = tuple(latencies.get_nowait() for _ in range(latencies.qsize())) - slowest: Final = max(probes, default=0.0) - record_property("slowest_health_seconds", round(slowest, 3)) - record_property("health_probes", len(probes)) - assert large <= _TIMING_GROWTH_LIMIT * small, (shape, small, large) - assert slowest <= _HEALTH_LATENCY_SECONDS, (shape, slowest, len(probes)) - - -@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) -def test_anthropic_pass_through_stream_relays_the_tool_stream_verbatim_and_logs_spend( - gateway: Gateway, tmp_path: Path -) -> None: - marker, marker_hex = _marker("wrapped") - overrides: Final = {"ANTHROPIC_API_BASE": "", "ANTHROPIC_API_KEY": _ANTHROPIC_KEY} - with wire_server(_peer) as wire, gateway.scenario() as scenario: - key: Final = scenario.key() - with owned_proxy_process(gateway, tmp_path, {**overrides, "ANTHROPIC_API_BASE": wire.url}) as owned: - response: Final = _post( - owned.gateway.client, - key, - "/anthropic/v1/messages", - _messages_body(_UNMAPPED_MODEL, marker, stream=True, tools=[_ANTHROPIC_JSON_TOOL]), - ) - assert response.status_code == 200, response.text - assert response.headers["content-type"].startswith(_SSE), response.headers - expected: Final = _tool_reply(f"msg_{marker_hex}", "wrapped", marker_hex, False, "sse").chunks - assert expected is not None - assert response.content == b"".join(expected), response.text - sent: Final = _one_delivery(wire, marker_hex) - assert sent.target == "/v1/messages", sent.target - assert _json_tool_schema(sent) == _SCHEMA, sent.body - (row,) = _pass_through_spend_rows(key) - assert row["model"] == _UNMAPPED_MODEL, row - - -def test_bedrock_pass_through_invoke_stream_relays_the_event_stream_verbatim(gateway: Gateway) -> None: - marker, marker_hex = _marker("wrapped") - with wire_server(_peer) as wire, gateway.scenario() as scenario: - deployment: Final = scenario.model( - model=f"bedrock/{_BEDROCK_INVOKE_MODEL}", - api_base=wire.url, - aws_access_key_id="AKIASCRIPTEDPROVIDER", - aws_secret_access_key="scripted-secret", - aws_region_name="us-east-1", - ) - response: Final = gateway.request( - "POST", - f"/bedrock/model/{deployment}/invoke-with-response-stream", - { - "anthropic_version": "bedrock-2023-05-31", - "max_tokens": 64, - "messages": [{"role": "user", "content": marker}], - "tools": [_ANTHROPIC_JSON_TOOL], - }, - ) - assert response.status_code == 200, response.text - assert response.headers["content-type"] == _EVENT_STREAM, response.headers - expected: Final = _tool_reply(f"msg_{marker_hex}", "wrapped", marker_hex, False, "aws").chunks - assert expected is not None - assert response.content == b"".join(expected), response.content[:200] - sent: Final = _one_delivery(wire, marker_hex) - assert sent.target == f"/model/{_BEDROCK_INVOKE_MODEL}/invoke-with-response-stream", sent.target - - -@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) -@pytest.mark.parametrize("route", _PASS_THROUGH_ROUTES) -def test_pass_through_stream_relays_the_scripted_frames_verbatim( - pass_through_proxy: tuple[OwnedProxy, Wire], session_gateway: Gateway, route: str -) -> None: - owned, wire = pass_through_proxy - marker_hex: Final = uuid.uuid4().hex - call: Final = _pass_through_call(route, f"relay {marker_hex}", session_gateway.key) - response: Final = owned.gateway.client.post(call.path, json=call.body, headers=call.headers) - assert response.status_code == 200, response.text - assert response.headers["content-type"].startswith(_SSE), response.headers - assert response.content == b"".join(_RELAY_CHUNKS), response.text - (seen,) = tuple(request for request in wire.drain() if marker_hex in request.body.decode()) - assert seen.target == call.upstream_target, seen.target diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py index cfd91ff96f3..bc0b81bed16 100644 --- a/tests/local_testing/test_get_model_info.py +++ b/tests/local_testing/test_get_model_info.py @@ -1,12 +1,10 @@ -from typing import Final # What is this? ## Unit testing for the 'get_model_info()' function import os import re from collections.abc import Collection, Mapping - - -from typing import List, Dict, Any, Final, Literal +from typing import Any, Dict, Final, List, Literal +from unittest.mock import MagicMock, patch import pytest @@ -15,7 +13,6 @@ from litellm import get_model_info from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.types.utils import ModelInfoBase from litellm.utils import _invalidate_model_cost_lowercase_map -from unittest.mock import MagicMock, patch def test_get_model_info_simple_model_name(): @@ -251,8 +248,7 @@ def test_get_model_info_custom_provider(): def test_get_model_info_custom_model_router(): - from litellm import Router - from litellm import get_model_info + from litellm import Router, get_model_info litellm._turn_on_debug() diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index feb5d1a2c43..b46ae42a8c2 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -44,17 +44,25 @@ import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at im import litellm.router as litellm_router_module # noqa: E402 # same import-time dependency import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency from litellm._logging import ALL_LOGGERS # noqa: E402 # same import-time dependency -from litellm.anthropic_beta_headers_manager import reload_beta_headers_config # noqa: E402 # same import-time dependency +from litellm.anthropic_beta_headers_manager import ( + reload_beta_headers_config, # noqa: E402 # same import-time dependency +) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER # noqa: E402 # same import-time dependency -from litellm.litellm_core_utils.prompt_templates import factory as prompt_factory_module # noqa: E402 # same import-time dependency +from litellm.litellm_core_utils.prompt_templates import ( + factory as prompt_factory_module, # noqa: E402 # same import-time dependency +) from litellm.litellm_core_utils.prompt_templates import ( # noqa: E402 # same import-time dependency image_handling as image_handling_module, ) -from litellm.llms.gemini.chat import transformation as gemini_chat_transformation_module # noqa: E402 # same import-time dependency from litellm.llms.custom_httpx.async_client_cleanup import ( # noqa: E402 # same import-time dependency close_litellm_async_clients, ) -from litellm.proxy.db import tool_registry_writer as tool_registry_writer_module # noqa: E402 # same import-time dependency +from litellm.llms.gemini.chat import ( + transformation as gemini_chat_transformation_module, # noqa: E402 # same import-time dependency +) +from litellm.proxy.db import ( + tool_registry_writer as tool_registry_writer_module, # noqa: E402 # same import-time dependency +) LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1", "localhost"] AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index 6b644338020..2f51828cfb5 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -1653,7 +1653,9 @@ async def test_http_response_handler_preserves_success_and_http_errors(status_co @pytest.mark.asyncio async def test_http_status_check_allows_auth_refresh_before_rejecting() -> None: - from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ClientCredentialsBearerAuth + from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( + ClientCredentialsBearerAuth, + ) seen = [] @@ -1908,6 +1910,7 @@ async def test_sse_read_failure_is_preserved() -> None: @pytest.mark.parametrize("mode", ["ok", "closed", "silent"]) async def test_transport_completion_and_normal_messages(transport: MCPTransport, mode: str, protocol_version: str) -> None: from mcp import ClientSession + from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message logging_callback: Final = AsyncMock() diff --git a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py index 7caef1e8467..5a75b8d384b 100644 --- a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -1,4 +1,5 @@ import litellm + """Tests for `BedrockConverseLLM.completion`. AWS credential resolution is stubbed so nothing reaches STS. diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index d33556dd555..2300670bfc9 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -2255,7 +2255,6 @@ async def test_assistant_message_cache_control(): """Test that assistant messages with cache_control generate cachePoint blocks.""" from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, - _bedrock_converse_messages_pt, ) # Test assistant message with string content and cache_control @@ -2308,7 +2307,6 @@ async def test_assistant_message_list_content_cache_control(): """Test assistant messages with list content and cache_control.""" from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, - _bedrock_converse_messages_pt, ) messages = [ @@ -2352,7 +2350,6 @@ async def test_tool_message_cache_control(): """Test that tool messages with cache_control generate cachePoint blocks.""" from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, - _bedrock_converse_messages_pt, ) messages = [ @@ -2416,7 +2413,6 @@ async def test_tool_message_string_content_cache_control(): """Test tool messages with string content and message-level cache_control.""" from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, - _bedrock_converse_messages_pt, ) messages = [ @@ -2472,7 +2468,6 @@ async def test_tool_message_search_results_maps_to_bedrock_search_result_block() """OpenAI tool message search_results should map to Bedrock searchResult blocks.""" from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, - _bedrock_converse_messages_pt, ) messages = [ @@ -2533,7 +2528,6 @@ async def test_tool_message_empty_search_results_falls_back_to_content(): """Empty search_results must not skip normal tool content processing.""" from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, - _bedrock_converse_messages_pt, ) messages = [ @@ -2702,7 +2696,6 @@ async def test_assistant_tool_calls_cache_control(): """Test that assistant tool_calls with cache_control generate cachePoint blocks.""" from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, - _bedrock_converse_messages_pt, ) messages = [ @@ -2754,7 +2747,6 @@ async def test_multiple_tool_calls_with_mixed_cache_control(): """Test multiple tool calls where only some have cache_control.""" from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, - _bedrock_converse_messages_pt, ) messages = [ @@ -2815,7 +2807,6 @@ async def test_no_cache_control_no_cache_point(): """Test that messages without cache_control don't generate cachePoint blocks.""" from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, - _bedrock_converse_messages_pt, ) messages = [ @@ -2860,9 +2851,6 @@ async def test_no_cache_control_no_cache_point(): def test_guarded_text_wraps_in_guardrail_converse_content(): """Test that guarded_text content type gets wrapped in guardContent blocks.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - _bedrock_converse_messages_pt, - ) messages = [ { @@ -2963,9 +2951,6 @@ def test_guarded_text_with_system_messages(): def test_guarded_text_with_mixed_content_types(): """Test guarded_text with mixed content types including images.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - _bedrock_converse_messages_pt, - ) messages = [ { @@ -3052,9 +3037,6 @@ async def test_async_guarded_text(): def test_guarded_text_with_tool_calls(): """Test guarded_text with tool calls in the conversation.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - _bedrock_converse_messages_pt, - ) messages = [ { @@ -3704,9 +3686,6 @@ def test_empty_assistant_message_handling(): """ # Import the litellm module that factory.py uses to ensure we patch the correct reference import litellm.litellm_core_utils.prompt_templates.factory as factory_module - from litellm.litellm_core_utils.prompt_templates.factory import ( - _bedrock_converse_messages_pt, - ) # Test case 1: Empty string content - test with modify_params=True to prevent merging messages = [ @@ -3796,9 +3775,6 @@ def test_empty_assistant_message_handling(): def test_bedrock_converse_trailing_prefix_assistant_skips_user_continue(): """Assistant prefill (prefix: true) must not inject a dummy user 'Please continue.' turn.""" import litellm.litellm_core_utils.prompt_templates.factory as factory_module - from litellm.litellm_core_utils.prompt_templates.factory import ( - _bedrock_converse_messages_pt, - ) messages = [ {"role": "user", "content": "Hello, how are you?"}, @@ -3825,9 +3801,6 @@ def test_bedrock_converse_trailing_prefix_assistant_skips_user_continue(): def test_bedrock_converse_leading_prefix_assistant_skips_user_continue(): """Leading assistant with prefix: true should not prepend dummy user.""" import litellm.litellm_core_utils.prompt_templates.factory as factory_module - from litellm.litellm_core_utils.prompt_templates.factory import ( - _bedrock_converse_messages_pt, - ) messages = [ {"role": "assistant", "content": "Partial", "prefix": True}, @@ -6338,9 +6311,6 @@ def test_bedrock_tool_message_openai_file_pdf_becomes_document(): inside a tool message content list should translate to a Bedrock toolResult.content[].document block. """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - _bedrock_converse_messages_pt, - ) pdf_b64 = "JVBERi0xLjQKJeLjz9MK" # tiny "%PDF-1.4\n" header messages = [ @@ -6381,9 +6351,6 @@ def test_bedrock_tool_message_image_url_pdf_data_uri_becomes_document(): BedrockImageProcessor correctly returns a {"document": ...} block, but the tool-result wrapper only appended the "image" case, silently dropping documents. """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - _bedrock_converse_messages_pt, - ) pdf_b64 = "JVBERi0xLjQKJeLjz9MK" messages = [ @@ -6426,7 +6393,6 @@ def test_bedrock_tool_message_file_id_http_url_becomes_document(): from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockImageProcessor, - _bedrock_converse_messages_pt, ) pdf_url = "https://example.com/whitepaper.pdf" @@ -6480,9 +6446,6 @@ def test_bedrock_tool_message_file_without_data_or_id_raises(): see an empty tool result and obscures the caller bug. """ import litellm - from litellm.litellm_core_utils.prompt_templates.factory import ( - _bedrock_converse_messages_pt, - ) messages = [ {"role": "user", "content": "Summarize."}, @@ -6509,9 +6472,6 @@ def test_bedrock_tool_message_image_url_png_still_becomes_image(): to a Bedrock image block (not document). Locks in existing behavior after the document-passthrough fix. """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - _bedrock_converse_messages_pt, - ) png_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGNgYGBgAAAABQABXvMqOgAAAABJRU5ErkJggg==" messages = [ @@ -6708,7 +6668,6 @@ async def test_grounding_source_and_query_rendered_as_text(): otherwise vanish from the prompt.""" from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, - _bedrock_converse_messages_pt, ) messages = [ @@ -7207,7 +7166,6 @@ async def test_message_level_cache_control_honors_ttl_for_supported_model( hardcoded `{"type": "default"}`), silently downgrading 1h to 5m.""" from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, - _bedrock_converse_messages_pt, ) messages = _agentic_messages_with_ttl(ttl_target) @@ -7233,9 +7191,6 @@ async def test_message_level_cache_control_honors_ttl_for_supported_model( def test_message_level_cache_control_drops_ttl_for_unsupported_model(ttl_target): """Models outside the extended-caching allow-list must keep emitting the plain `{"type": "default"}` cachePoint (Bedrock rejects `ttl` for them).""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - _bedrock_converse_messages_pt, - ) result = _bedrock_converse_messages_pt( messages=_agentic_messages_with_ttl(ttl_target), diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 2fe636d8f68..1f22ea10a7e 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -46,7 +46,7 @@ def test_bedrock_response_stream_shape_lazy_loads_once(): """ get_bedrock_response_stream_shape() loads from botocore at most once per process. """ - from unittest.mock import MagicMock, patch + from unittest.mock import MagicMock import litellm.llms.bedrock.common_utils as mod @@ -74,7 +74,6 @@ def test_bedrock_response_stream_shape_load_failure_returns_none(): should return None rather than propagating the exception, so the module still imports cleanly. """ - from unittest.mock import patch import litellm.llms.bedrock.common_utils as mod @@ -136,7 +135,7 @@ def test_bedrock_parse_message_from_event_raises_on_none_shape(): _parse_message_from_event must raise BedrockError before touching the botocore parser — not an opaque AttributeError from inside botocore. """ - from unittest.mock import MagicMock, patch + from unittest.mock import MagicMock import litellm.llms.bedrock.common_utils as mod from litellm.llms.bedrock.common_utils import ( @@ -260,7 +259,6 @@ def test_context_window_suffix_stripped_for_cost_lookup(): Models configured like `bedrock/us.anthropic.claude-opus-4-6-v1[1m]` should resolve to the base model name so pricing can be found. """ - from litellm.llms.bedrock.common_utils import get_bedrock_base_model assert get_bedrock_base_model("us.anthropic.claude-opus-4-6-v1[1m]") == "anthropic.claude-opus-4-6-v1" assert get_bedrock_base_model("us.anthropic.claude-sonnet-4-6[1m]") == "anthropic.claude-sonnet-4-6" @@ -280,7 +278,6 @@ def test_context_window_suffix_stripped_for_cost_lookup(): def test_legacy_mantle_route_prefix_stripped_for_cost_lookup(): """The mantle/ route token is a routing prefix like openai/, so a bedrock/mantle/ deployment must resolve the bare Bedrock model for cost lookup while still routing to Mantle.""" - from litellm.llms.bedrock.common_utils import get_bedrock_base_model, strip_bedrock_routing_prefix assert strip_bedrock_routing_prefix("mantle/anthropic.claude-sonnet-5") == "anthropic.claude-sonnet-5" assert get_bedrock_base_model("bedrock/mantle/anthropic.claude-sonnet-5") == "anthropic.claude-sonnet-5" @@ -417,7 +414,6 @@ def test_capability_lookups_fall_back_to_base_model_when_regional_entry_lacks_fi must not shadow a base entry that has it (`get(model) or get(base)` used to short-circuit on the truthy regional dict and drop the capability). """ - import litellm from litellm.llms.bedrock.common_utils import ( bedrock_converse_supports_parallel_tool_use_config, is_claude_4_5_on_bedrock, @@ -543,7 +539,6 @@ def test_apply_structured_output_prefers_legacy_output_format(local_model_cost_m def test_sign_aws_request_assumes_role_with_external_id(monkeypatch): """A trust policy requiring sts:ExternalId must be satisfied when signing batch API requests.""" import datetime - from unittest.mock import patch import boto3 from botocore.exceptions import ClientError @@ -596,7 +591,6 @@ def test_sign_aws_request_assumes_role_with_external_id(monkeypatch): def test_sign_aws_request_assumes_role_with_session_tags(monkeypatch): """Batch and file signing must carry the deployment's session tags into the AssumeRole call too.""" import datetime - from unittest.mock import patch import boto3 from botocore.exceptions import ClientError @@ -855,7 +849,6 @@ def test_bedrock_mantle_get_error_class_keeps_provider_headers(config): def _bedrock_configs_with_get_error_class(): - import importlib import inspect import pathlib @@ -940,7 +933,6 @@ def test_unmapped_openai_family_model_routes_to_converse(): """ from typing import Final - import litellm unmapped: Final = "bedrock/global.openai.gpt-99-unmapped" assert unmapped.removeprefix("bedrock/") not in litellm.bedrock_converse_models diff --git a/tests/unit/test_constants.py b/tests/unit/test_constants.py index 21b75ceed76..e926f258295 100644 --- a/tests/unit/test_constants.py +++ b/tests/unit/test_constants.py @@ -1,7 +1,8 @@ -import sys import ast +import importlib import inspect import json +import sys from unittest import mock import httpx @@ -9,9 +10,6 @@ import pytest import respx from fastapi.testclient import TestClient - -import importlib - import litellm from litellm import constants diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 494a129f313..88f15de0078 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -2,14 +2,11 @@ import asyncio import base64 import contextlib import contextvars -import copy -import importlib import io import json import logging import os import queue -import re import threading from collections.abc import Callable, Iterator, Mapping from concurrent.futures import Future, ThreadPoolExecutor @@ -40,7 +37,6 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.thread_pool_executor import executor as logging_executor -from litellm.llms.base_llm.audio_transcription.transformation import BaseAudioTranscriptionConfig from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.proxy.utils import is_valid_api_key from litellm.types.caching import CachingSupportedCallTypes @@ -76,19 +72,12 @@ from litellm.utils import ( _is_streaming_request, _run_success_deployment_hook_on_converted_chat_stream, _snapshot_exception_for_hook, - add_custom_logger_callback_to_specific_event, async_post_call_failure_deployment_hook, async_post_call_success_deployment_hook, calculate_max_parallel_requests, - check_valid_key, client, - filter_out_litellm_params, - get_applied_guardrails, - get_llm_provider, get_non_default_completion_params, - get_optional_params_embeddings, get_optional_params_image_gen, - get_potential_model_names, get_prompt_cache_min_tokens, is_cached_message, is_prompt_caching_valid_prompt, @@ -3166,7 +3155,9 @@ class TestExtraBodyCannotOverrideModel: } ] - untouched = litellm.get_optional_params(model="my-vllm-model", custom_llm_provider="hosted_vllm", tools=tools()) + untouched = litellm.get_optional_params( + model="my-vllm-model", custom_llm_provider="hosted_vllm", tools=tools() + ) assert untouched["tools"][0]["function"]["custom_marker"] == "LEAK", untouched result = litellm.get_optional_params( @@ -4135,9 +4126,6 @@ def test_is_prompt_caching_valid_prompt_explicit_min_token_count_overrides_model ) - - - def test_custom_logger_guards_ignore_subclass_instances(monkeypatch: pytest.MonkeyPatch) -> None: """Regression LIT-4392: the success/failure existence guards used isinstance, so a user subclass of a built-in logger already promoted into the callback lists made the guard @@ -4583,13 +4571,7 @@ _SUCCESS_RESPONSES_BY_CALL_TYPE: Final = ( pytest.param(EmbeddingResponse(model="text-embedding-3-small"), CallTypes.aembedding, id="embedding"), pytest.param( ResponsesAPIResponse( - id="resp_abc", - created_at=1, - output=[], - parallel_tool_calls=False, - tool_choice="auto", - tools=[], - model="gpt-5.6", + id="resp_abc", created_at=1, output=[], parallel_tool_calls=False, tool_choice="auto", tools=[], model="gpt-5.6" ), CallTypes.aresponses, id="responses", @@ -4617,9 +4599,7 @@ async def test_success_deployment_hook_raising_keeps_response_and_runs_later_hoo assert result is response assert second_hook.seen_responses == (response,) - failure_logs: Final = tuple( - r for r in caplog.records if "async_post_call_success_deployment_hook error" in r.message - ) + failure_logs: Final = tuple(r for r in caplog.records if "async_post_call_success_deployment_hook error" in r.message) assert len(failure_logs) == 1 assert "_ChatShapedSuccessDeploymentHook" in failure_logs[0].message assert str(call_type) in failure_logs[0].message @@ -4743,9 +4723,6 @@ async def test_wrapper_async_logs_converted_responses_stream_with_standard_loggi assert success_kwargs["stream"] is True - - - @pytest.mark.asyncio async def test_wrapper_async_replays_cached_converted_chat_stream_as_stream( monkeypatch: pytest.MonkeyPatch, @@ -6546,9 +6523,6 @@ def test_function_setup_never_logs_the_ocr_data_uri_payload() -> None: assert payload not in str(logged) - - - @pytest.mark.asyncio async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: pytest.MonkeyPatch) -> None: """Chat over the Responses bridge exits two @client wrappers with one logging object. Issue @@ -6611,6 +6585,3 @@ async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: assert len(counting_logger.logged_results) == 1, counting_logger.logged_results assert counting_logger.logged_results[0] is inner_result - - -