diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 2f1a4147544..d67b659fe37 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1354,12 +1354,32 @@ def drop_non_python_regex_patterns(schema: Mapping[str, object]) -> Mapping[str, at more schema levels than a JSON parser admits, so a cyclic schema built in code cannot spin it. """ + return _schema_without_rejected_regex(schema, _is_not_python_regex) + + +def drop_lookaround_regex_patterns(schema: Mapping[str, object]) -> Mapping[str, object]: + """Drop every regex in a schema position that uses a lookaround assertion. + + Some Bedrock Converse families compile tool schema regexes with an engine that + has no lookahead or lookbehind and refuse the whole request over one. The ``(?=``, + ``(?!``, ``(?<=`` and ``(? Mapping[str, object]: rebuilt: dict[int, Mapping[str, object]] = {} # mutable-ok: per-call memo of rewritten nodes, deepest level first for level in reversed(tuple(islice(_schema_levels(schema), _MAX_SCHEMA_NESTING))): rebuilt.update( (id(node), rewritten) for node in level - if (rewritten := _node_without_non_python_regex(node, rebuilt)) is not node + if (rewritten := _node_without_rejected_regex(node, rebuilt, rejected)) is not node ) return rebuilt.get(id(schema), schema) @@ -1381,23 +1401,56 @@ def _subschemas(node: Mapping[str, object]) -> Iterator[Mapping[str, object]]: yield value -def _node_without_non_python_regex( - node: Mapping[str, object], rebuilt: Mapping[int, Mapping[str, object]] +def _node_without_rejected_regex( + node: Mapping[str, object], + rebuilt: Mapping[int, Mapping[str, object]], + rejected: Callable[[str], bool], ) -> Mapping[str, object]: kept: Final = { - key: _keyword_value_rebuilt(key, value, rebuilt) + key: _keyword_value_rebuilt(key, value, rebuilt, rejected) for key, value in node.items() - if key != "pattern" or not isinstance(value, str) or _is_python_regex(value) + if key != "pattern" or not isinstance(value, str) or not rejected(value) } - return node if len(kept) == len(node) and all(kept[key] is node[key] for key in kept) else kept + if len(kept) == len(node) and all(kept[key] is node[key] for key in kept): + return node + dropped_pattern_properties: Final = _dropped_pattern_properties(node, kept, rebuilt) + if not dropped_pattern_properties or kept.get("additionalProperties") is not False: + return kept + return {**kept, "additionalProperties": _any_of(dropped_pattern_properties)} -def _keyword_value_rebuilt(key: str, value: object, rebuilt: Mapping[int, Mapping[str, object]]) -> object: +def _dropped_pattern_properties( + node: Mapping[str, object], + kept: Mapping[str, object], + rebuilt: Mapping[int, Mapping[str, object]], +) -> tuple[object, ...]: + before: Final = _schema_at(node, "patternProperties") + after: Final = _schema_at(kept, "patternProperties") + if before is None or after is None: + return () + return tuple(rebuilt.get(id(sub), sub) for name, sub in before.items() if name not in after) + + +def _schema_at(container: Mapping[str, object], key: str) -> Mapping[str, object] | None: + value: Final = container.get(key) + return value if isinstance(value, dict) else None + + +def _any_of(schemas: tuple[object, ...]) -> object: + return schemas[0] if len(schemas) == 1 else {"anyOf": list(schemas)} + + +def _keyword_value_rebuilt( + key: str, + value: object, + rebuilt: Mapping[int, Mapping[str, object]], + rejected: Callable[[str], bool], +) -> object: if key in _SUBSCHEMA_MAP_KEYWORDS and isinstance(value, dict): kept: Final = { name: rebuilt.get(id(sub), sub) for name, sub in value.items() - if key != "patternProperties" or not isinstance(name, str) or _is_python_regex(name) + if key != "patternProperties" or not isinstance(name, str) or not rejected(name) } return value if len(kept) == len(value) and all(kept[name] is value[name] for name in kept) else kept if key in _SUBSCHEMA_LIST_KEYWORDS and isinstance(value, list): @@ -1408,12 +1461,19 @@ def _keyword_value_rebuilt(key: str, value: object, rebuilt: Mapping[int, Mappin return value -def _is_python_regex(pattern: str) -> bool: +def _is_not_python_regex(pattern: str) -> bool: try: re.compile(pattern) except (re.error, RecursionError): - return False - return True + return True + return False + + +_REGEX_LOOKAROUND_RE: Final = re.compile(r"\(\? bool: + return _REGEX_LOOKAROUND_RE.search(pattern) is not None def flatten_combinators_and_drop_non_python_regex_patterns(schema: Mapping[str, object]) -> Mapping[str, object]: @@ -1424,16 +1484,23 @@ def tool_with_sanitized_parameters( tool: Mapping[str, object], sanitize: Callable[[Mapping[str, object]], Mapping[str, object]], ) -> Mapping[str, object]: - function: Final = tool.get("function") - if not isinstance(function, dict): + """Run the tool's JSON schema through ``sanitize``: ``function.parameters`` on an + OpenAI tool, ``input_schema`` on an Anthropic one. The same object comes back when + nothing changed.""" + function: Final = _schema_at(tool, "function") + if function is not None: + parameters: Final = _schema_at(function, "parameters") + if parameters is None: + return tool + sanitized_parameters: Final = sanitize(parameters) + if sanitized_parameters is parameters: + return tool + return {**tool, "function": {**function, "parameters": sanitized_parameters}} + input_schema: Final = _schema_at(tool, "input_schema") + if input_schema is None: return tool - parameters: Final = function.get("parameters") - if not isinstance(parameters, dict): - return tool - sanitized: Final = sanitize(parameters) - if sanitized is parameters: - return tool - return {**tool, "function": {**function, "parameters": sanitized}} + sanitized_schema: Final = sanitize(input_schema) + return tool if sanitized_schema is input_schema else {**tool, "input_schema": sanitized_schema} def _get_image_mime_type_from_url(url: str) -> str | None: diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 48a8b1b44bb..9f094521bfe 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -12,6 +12,7 @@ from itertools import chain from typing import TYPE_CHECKING, Final, Literal, cast, overload import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger @@ -28,6 +29,8 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.prompt_templates.common_utils import ( _parse_content_for_reasoning, + drop_lookaround_regex_patterns, + tool_with_sanitized_parameters, ) from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, @@ -49,6 +52,7 @@ from litellm.llms.anthropic.chat.transformation import ( ) from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.bedrock.common_utils import bedrock_model_supports_regex_lookaround from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, bedrock_request_metadata_is_owned, @@ -128,6 +132,17 @@ UNSUPPORTED_BEDROCK_CONVERSE_BETA_PATTERNS: Final = [ ] +_TOOLS_AS_SENT: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +def _tools_the_model_accepts( + tools: Sequence[Mapping[str, object]], model: str, litellm_params: Mapping[str, object] | None +) -> list[Mapping[str, object]]: + if bedrock_model_supports_regex_lookaround(model, litellm_params): + return list(tools) + return [tool_with_sanitized_parameters(tool, drop_lookaround_regex_patterns) for tool in tools] + + class AmazonConverseConfig(BaseConfig): """ Reference - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html @@ -1689,6 +1704,7 @@ class AmazonConverseConfig(BaseConfig): model: str, headers: dict | None, additional_request_params: dict, + litellm_params: Mapping[str, object] | None = None, ) -> tuple[list[ToolBlock], list]: """Process tools and collect anthropic_beta values.""" bedrock_tools: list[ToolBlock] = [] @@ -1729,7 +1745,9 @@ class AmazonConverseConfig(BaseConfig): computer_use_tools, regular_tools = self._separate_computer_use_tools(filtered_tools, model) # Process regular function tools using existing logic - bedrock_tools = _bedrock_tools_pt(regular_tools, model=model) + bedrock_tools = _bedrock_tools_pt( + _tools_the_model_accepts(regular_tools, model, litellm_params), model=model + ) # Add computer use tools and anthropic_beta if needed (only when computer use tools are present) if computer_use_tools: @@ -1793,7 +1811,10 @@ class AmazonConverseConfig(BaseConfig): additional_request_params["tools"] = transformed_computer_tools else: # No computer use tools, process all tools as regular tools - bedrock_tools = _bedrock_tools_pt(filtered_tools, model=model) + bedrock_tools = _bedrock_tools_pt( + _tools_the_model_accepts(_TOOLS_AS_SENT.validate_python(filtered_tools), model, litellm_params), + model=model, + ) # Append pre-formatted tools (systemTool etc.) after transformation bedrock_tools.extend(pre_formatted_tools) @@ -1905,7 +1926,7 @@ class AmazonConverseConfig(BaseConfig): # Process tools and collect beta values bedrock_tools, anthropic_beta_list = self._process_tools_and_beta( - original_tools, model, headers, additional_request_params + original_tools, model, headers, additional_request_params, litellm_params ) # Append cachePoint to tools if cache_control_injection_points has tool_config diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 323890624a0..61c8e3e9d97 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -11,7 +11,7 @@ import os import re from collections.abc import Iterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias from typing_extensions import ReadOnly, TypedDict @@ -972,6 +972,7 @@ def is_claude_4_5_on_bedrock(model: str) -> bool: _BEDROCK_MODEL_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$") +_DEPLOYMENT_MODEL_INFO: Final = TypeAdapter(dict[str, object]) def bedrock_converse_supports_strict_tools(model: str) -> bool: @@ -989,12 +990,38 @@ def bedrock_converse_supports_strict_tools(model: str) -> bool: base: Final = get_bedrock_base_model(model) if not base.startswith("anthropic"): return False - flag: Final = _get_bedrock_converse_strict_tools_flag(base) + flag: Final = _bedrock_converse_model_flag(base, "bedrock_converse_supports_strict_tools") return flag if flag is not None else True -def _get_bedrock_converse_strict_tools_flag(base_model: str) -> bool | None: - candidates: Final = dict.fromkeys((base_model, _BEDROCK_MODEL_VERSION_SUFFIX_RE.sub("", base_model))) +def bedrock_model_supports_regex_lookaround(model: str, litellm_params: Mapping[str, object] | None = None) -> bool: + """ + Whether ``model`` accepts lookahead and lookbehind assertions in tool schema regexes. + + The deployment's ``model_info.supports_regex_lookaround`` wins, then the + ``model_prices_and_context_window.json`` entry of its ``base_model``, then the + entry of ``model`` itself. A model nobody flagged keeps its schema as sent. + """ + params: Final = litellm_params or {} + model_info: Final = _DEPLOYMENT_MODEL_INFO.validate_python(params.get("model_info") or {}) + deployment_flag: Final = model_info.get("supports_regex_lookaround") + if isinstance(deployment_flag, bool): + return deployment_flag + base_model: Final = params.get("base_model") + candidates: Final = (*((base_model,) if isinstance(base_model, str) else ()), model) + flags: Final = (_bedrock_converse_model_flag(candidate, "supports_regex_lookaround") for candidate in candidates) + return next((flag for flag in flags if flag is not None), True) + + +_BedrockConverseModelFlag: TypeAlias = Literal[ + "bedrock_converse_supports_strict_tools", + "supports_regex_lookaround", +] + + +def _bedrock_converse_model_flag(model: str, key: _BedrockConverseModelFlag) -> bool | None: + base: Final = get_bedrock_base_model(model) + candidates: Final = dict.fromkeys((model, base, _BEDROCK_MODEL_VERSION_SUFFIX_RE.sub("", base))) for candidate in candidates: with contextlib.suppress(Exception): model_info = get_cached_model_info()( @@ -1002,15 +1029,13 @@ def _get_bedrock_converse_strict_tools_flag(base_model: str) -> bool | None: custom_llm_provider="bedrock", ) - flag = model_info.get("bedrock_converse_supports_strict_tools") + flag = model_info.get(key) if isinstance(flag, bool): return flag model_cost_key = model_info.get("key") if isinstance(model_cost_key, str): - local_flag = ( - _get_local_model_cost_map().get(model_cost_key, {}).get("bedrock_converse_supports_strict_tools") - ) + local_flag = _get_local_model_cost_map().get(model_cost_key, {}).get(key) if isinstance(local_flag, bool): return local_flag return None diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 21bdee73464..a089fa984b4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -47460,6 +47460,7 @@ "supports_tool_choice": true }, "us-gov.xai.grok-4.6": { + "supports_regex_lookaround": false, "input_cost_per_token": 2.64e-06, "output_cost_per_token": 7.92e-06, "cache_read_input_token_cost": 6.6e-07, @@ -59073,6 +59074,7 @@ "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-5-5.html" }, "us.xai.grok-4.6": { + "supports_regex_lookaround": false, "input_cost_per_token": 2.2e-06, "output_cost_per_token": 6.6e-06, "cache_read_input_token_cost": 5.5e-07, @@ -59089,6 +59091,7 @@ "supports_vision": true }, "global.xai.grok-4.6": { + "supports_regex_lookaround": false, "input_cost_per_token": 2e-06, "output_cost_per_token": 6e-06, "cache_read_input_token_cost": 5e-07, @@ -76781,6 +76784,7 @@ "supports_web_search": true }, "moonshotai.kimi-k3": { + "supports_regex_lookaround": false, "cache_creation_input_token_cost": 4.125e-06, "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, @@ -76801,6 +76805,7 @@ "supports_vision": true }, "global.moonshotai.kimi-k3": { + "supports_regex_lookaround": false, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, @@ -76821,6 +76826,7 @@ "supports_vision": true }, "us.moonshotai.kimi-k3": { + "supports_regex_lookaround": false, "cache_creation_input_token_cost": 4.125e-06, "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, @@ -79377,6 +79383,7 @@ "supports_vision": false }, "global.xai.grok-4.7": { + "supports_regex_lookaround": false, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "bedrock_converse", @@ -79393,6 +79400,7 @@ "supports_vision": true }, "us.xai.grok-4.7": { + "supports_regex_lookaround": false, "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 2.2e-06, "litellm_provider": "bedrock_converse", @@ -79409,6 +79417,7 @@ "supports_vision": true }, "xai.grok-4.7": { + "supports_regex_lookaround": false, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "bedrock_converse", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8c10b9e3497..6919fd6fd27 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -211,6 +211,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): vertex_ai_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"] | None] bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None bedrock_converse_supports_strict_tools: bool | None + supports_regex_lookaround: ReadOnly[bool | None] class SearchContextCostPerQuery(TypedDict, total=False): @@ -3844,10 +3845,13 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"}) +DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset({"supports_regex_lookaround"}) + SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__) - frozenset(CustomPricingLiteLLMParams.model_fields) - DEPLOYMENT_SCOPED_PRICING_FIELDS + - DEPLOYMENT_SCOPED_CAPABILITY_FIELDS ) diff --git a/litellm/utils.py b/litellm/utils.py index 05d5986885c..bea8cfffd3c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6321,6 +6321,7 @@ def _get_model_info_helper( default_reasoning_effort=_model_info.get("default_reasoning_effort", None), bedrock_output_config_effort_ceiling=_model_info.get("bedrock_output_config_effort_ceiling", None), bedrock_converse_supports_strict_tools=_model_info.get("bedrock_converse_supports_strict_tools", None), + supports_regex_lookaround=_model_info.get("supports_regex_lookaround", None), supports_computer_use=_model_info.get("supports_computer_use", None), search_context_cost_per_query=_model_info.get("search_context_cost_per_query", None), web_search_billing_unit=_model_info.get("web_search_billing_unit", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 21bdee73464..a089fa984b4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -47460,6 +47460,7 @@ "supports_tool_choice": true }, "us-gov.xai.grok-4.6": { + "supports_regex_lookaround": false, "input_cost_per_token": 2.64e-06, "output_cost_per_token": 7.92e-06, "cache_read_input_token_cost": 6.6e-07, @@ -59073,6 +59074,7 @@ "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-5-5.html" }, "us.xai.grok-4.6": { + "supports_regex_lookaround": false, "input_cost_per_token": 2.2e-06, "output_cost_per_token": 6.6e-06, "cache_read_input_token_cost": 5.5e-07, @@ -59089,6 +59091,7 @@ "supports_vision": true }, "global.xai.grok-4.6": { + "supports_regex_lookaround": false, "input_cost_per_token": 2e-06, "output_cost_per_token": 6e-06, "cache_read_input_token_cost": 5e-07, @@ -76781,6 +76784,7 @@ "supports_web_search": true }, "moonshotai.kimi-k3": { + "supports_regex_lookaround": false, "cache_creation_input_token_cost": 4.125e-06, "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, @@ -76801,6 +76805,7 @@ "supports_vision": true }, "global.moonshotai.kimi-k3": { + "supports_regex_lookaround": false, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, @@ -76821,6 +76826,7 @@ "supports_vision": true }, "us.moonshotai.kimi-k3": { + "supports_regex_lookaround": false, "cache_creation_input_token_cost": 4.125e-06, "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, @@ -79377,6 +79383,7 @@ "supports_vision": false }, "global.xai.grok-4.7": { + "supports_regex_lookaround": false, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "bedrock_converse", @@ -79393,6 +79400,7 @@ "supports_vision": true }, "us.xai.grok-4.7": { + "supports_regex_lookaround": false, "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 2.2e-06, "litellm_provider": "bedrock_converse", @@ -79409,6 +79417,7 @@ "supports_vision": true }, "xai.grok-4.7": { + "supports_regex_lookaround": false, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "bedrock_converse", diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index cdf023e71ef..bec34b32966 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -1062,6 +1062,9 @@ "supports_reasoning": { "type": "boolean" }, + "supports_regex_lookaround": { + "type": "boolean" + }, "supports_response_schema": { "type": "boolean" }, diff --git a/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py b/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py new file mode 100644 index 00000000000..b92360a2193 --- /dev/null +++ b/tests/integration/providers/test_bedrock_converse_lookaround_regex_chaos.py @@ -0,0 +1,379 @@ +import asyncio +import json +import re +import signal +import threading +import uuid +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import unquote, urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_KIMI: Final = "global.moonshotai.kimi-k3" +_NOVA: Final = "us.amazon.nova-lite-v1:0" +_AWS: Final[dict[str, JsonValue]] = { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", +} +_LOOKAHEAD: Final = r"^(?!\.\.?(?:\/|$))[A-Za-z0-9_\-.~:@+]{1,200}$" +_PLAIN: Final = r"^[a-z][a-z0-9_]*$" +_TOOL: Final = "ArtifactData" +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_LIST: Final = TypeAdapter(list[JsonValue]) +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_USAGE: Final[dict[str, JsonValue]] = {"inputTokens": 21, "outputTokens": 7, "totalTokens": 28} +_WIRE_AS_SENT: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "collection": {"type": "string", "pattern": _LOOKAHEAD}, + "doc_id": {"type": "string", "pattern": _PLAIN}, + }, + "required": ["collection"], +} +_WIRE_LOOKAROUND_FREE: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"collection": {"type": "string"}, "doc_id": {"type": "string", "pattern": _PLAIN}}, + "required": ["collection"], +} +_SCHEMA_AS_SENT: Final[dict[str, JsonValue]] = {**_WIRE_AS_SENT, "additionalProperties": False} + +Endpoint = Literal["chat", "messages", "responses"] +_ENDPOINTS: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses") + + +@dataclass(frozen=True, slots=True) +class _Fleet: + kimi_bare: str + kimi_flagged_true: str + nova_off: str + nova_bare: str + + def names(self) -> tuple[str, ...]: + return (self.kimi_bare, self.kimi_flagged_true, self.nova_off, self.nova_bare) + + def expected_schema(self, model: str) -> dict[str, JsonValue]: + return _WIRE_LOOKAROUND_FREE if model in (self.kimi_bare, self.nova_off) else _WIRE_AS_SENT + + +@dataclass(frozen=True, slots=True) +class _Call: + model: str + endpoint: Endpoint + stream: bool + marker: str + + +@dataclass(frozen=True, slots=True) +class _Served: + call: _Call + status: int + text: str + + +def _answer(marker: str) -> str: + return f"answer marker-{marker}" + + +def _frame(event_type: str, payload: Mapping[str, JsonValue]) -> bytes: + return _aws_event_frame(event_type, payload, "sc", "u") + + +def _stream_frames(marker: str) -> tuple[bytes, ...]: + return ( + _frame("messageStart", {"role": "assistant"}), + _frame("contentBlockDelta", {"delta": {"text": "answer "}, "contentBlockIndex": 0}), + _frame("contentBlockDelta", {"delta": {"text": f"marker-{marker}"}, "contentBlockIndex": 0}), + _frame("contentBlockStop", {"contentBlockIndex": 0}), + _frame("messageStop", {"stopReason": "end_turn"}), + _frame("metadata", {"usage": _USAGE}), + ) + + +def _text_reply(marker: str, stream: bool, abort_after: int | None = None) -> Reply: + if stream: + return Reply(content_type=_EVENT_STREAM, chunks=_stream_frames(marker), abort_after=abort_after) + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": _answer(marker)}]}}, + "stopReason": "end_turn", + "usage": _USAGE, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + +def _marker_of(request: Request) -> str: + found: Final = _MARKER.search(request.body.decode()) + assert found is not None, request.body + return found.group(1) + + +def _is_stream(request: Request) -> bool: + return unquote(request.target).endswith("/converse-stream") + + +def _echo(request: Request) -> Reply: + return _text_reply(_marker_of(request), _is_stream(request)) + + +def _path(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "/v1/chat/completions" + case "messages": + return "/v1/messages" + case "responses": + return "/v1/responses" + + +def _body(call: _Call) -> dict[str, JsonValue]: + question: Final = f"Question marker-{call.marker}" + common: Final[dict[str, JsonValue]] = { + "model": call.model, + "stream": call.stream, + "num_retries": 0, + "cache": {"no-cache": True}, + } + tool: Final[dict[str, JsonValue]] = {"description": f"{_TOOL} tool"} + match call.endpoint: + case "chat": + return { + **common, + "messages": [{"role": "user", "content": question}], + "max_tokens": 64, + "tools": [{"type": "function", "function": {"name": _TOOL, **tool, "parameters": _SCHEMA_AS_SENT}}], + } + case "messages": + return { + **common, + "messages": [{"role": "user", "content": question}], + "max_tokens": 64, + "tools": [{"name": _TOOL, **tool, "input_schema": _SCHEMA_AS_SENT}], + } + case "responses": + return { + **common, + "input": question, + "max_output_tokens": 64, + "tools": [{"type": "function", "name": _TOOL, **tool, "parameters": _SCHEMA_AS_SENT}], + } + + +def _received_schema(request: Request) -> dict[str, JsonValue]: + body: Final = _JSON.validate_json(request.body) + (tool,) = _LIST.validate_python(object_value(body["toolConfig"])["tools"]) + spec: Final = object_value(object_value(tool)["toolSpec"]) + assert spec["name"] == _TOOL, spec + return object_value(object_value(spec["inputSchema"])["json"]) + + +def _assert_schemas_by_marker(received: tuple[Request, ...], calls: tuple[_Call, ...], fleet: _Fleet) -> None: + by_marker: Final = MappingProxyType({call.marker: call for call in calls}) + assert sorted(_marker_of(request) for request in received) == sorted(by_marker), len(received) + for request in received: + call: Final = by_marker[_marker_of(request)] + assert _is_stream(request) == call.stream, (call, request.target) + assert _received_schema(request) == fleet.expected_schema(call.model), (call, request.body) + + +def _spend_statuses(model: str, expected: int) -> list[JsonValue]: + rows: Final = eventually( + lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)), + lambda found: len(found) >= expected, + seconds=60, + ) + assert len({row["request_id"] for row in rows}) == len(rows), rows + return [row["status"] for row in rows] + + +async def _send(client: httpx.AsyncClient, key: str, call: _Call) -> _Served: + async with client.stream( + "POST", + _path(call.endpoint), + json=_body(call), + headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"}, + ) as response: + raw: Final = await response.aread() + return _Served(call=call, status=response.status_code, text=raw.decode()) + + +async def _burst( + base_url: str, key: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, key, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Served)) + + +def _mixed_calls(fleet: _Fleet, count: int) -> tuple[_Call, ...]: + names: Final = fleet.names() + return tuple( + _Call( + model=names[index % len(names)], + endpoint=_ENDPOINTS[(index // len(names)) % len(_ENDPOINTS)], + stream=(index // (len(names) * len(_ENDPOINTS))) % 2 == 0, + marker=uuid.uuid4().hex, + ) + for index in range(count) + ) + + +def _assert_answered_with_its_own_marker(served: _Served) -> None: + assert served.status == 200, served.text + assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text + + +def _fleet_config(wire: Wire, tmp_path: Path) -> tuple[Path, _Fleet]: + run_id: Final = uuid.uuid4().hex[:8] + fleet: Final = _Fleet( + kimi_bare=f"kimi-bare-{run_id}", + kimi_flagged_true=f"kimi-flagged-true-{run_id}", + nova_off=f"nova-off-{run_id}", + nova_bare=f"nova-bare-{run_id}", + ) + params: Final[dict[str, JsonValue]] = {"api_base": wire.url, **_AWS} + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["model_list"] = [ + {"model_name": fleet.kimi_bare, "litellm_params": {"model": f"bedrock/{_KIMI}", **params}}, + { + "model_name": fleet.kimi_flagged_true, + "litellm_params": {"model": f"bedrock/{_KIMI}", **params}, + "model_info": {"supports_regex_lookaround": True}, + }, + { + "model_name": fleet.nova_off, + "litellm_params": {"model": f"bedrock/converse/{_NOVA}", **params}, + "model_info": {"supports_regex_lookaround": False}, + }, + {"model_name": fleet.nova_bare, "litellm_params": {"model": f"bedrock/converse/{_NOVA}", **params}}, + ] + path: Final = tmp_path / "bedrock-lookaround-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path, fleet + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +@pytest.mark.timeout(600) +async def test_a_mixed_burst_across_two_workers_cleans_only_the_flagged_deployments( + gateway: Gateway, tmp_path: Path +) -> None: + with wire_server(_echo) as wire: + path, fleet = _fleet_config(wire, tmp_path) + calls: Final = _mixed_calls(fleet, 36) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + served: Final = await _burst(str(candidate.client.base_url), candidate.key, calls) + assert len(served) == 36 + for item in served: + _assert_answered_with_its_own_marker(item) + _assert_schemas_by_marker(wire.drain(), calls, fleet) + for name in fleet.names(): + assert _spend_statuses(name, 9) == ["success"] * 9 + + +@pytest.mark.timeout(600) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_cleaning_schemas(gateway: Gateway, tmp_path: Path) -> None: + release: Final = threading.Event() + held_markers: Final[SimpleQueue[str]] = SimpleQueue() + + def held(request: Request) -> Reply: + held_markers.put(_marker_of(request)) + assert release.wait(timeout=60), "The burst was never released" + return _echo(request) + + with wire_server(held) as wire: + path, fleet = _fleet_config(wire, tmp_path) + calls: Final = tuple( + _Call(model=fleet.kimi_bare, endpoint="chat", stream=False, marker=uuid.uuid4().hex) for _ in range(20) + ) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + burst: Final = asyncio.create_task( + _burst(str(candidate.client.base_url), candidate.key, calls, tolerate_transport_errors=True) + ) + await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + victim: Final = psutil.Process(victim_pid) + victim.suspend() + victim.send_signal(signal.SIGKILL) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) + follow_up: Final = _Call(model=fleet.kimi_bare, endpoint="chat", stream=False, marker=uuid.uuid4().hex) + (answered,) = await _burst(str(candidate.client.base_url), candidate.key, (follow_up,)) + _assert_answered_with_its_own_marker(answered) + _assert_schemas_by_marker(wire.drain(), (*calls, follow_up), fleet) + + +@pytest.mark.timeout(600) +async def test_peer_stream_aborts_reach_callers_while_the_rest_of_the_burst_is_cleaned( + gateway: Gateway, tmp_path: Path +) -> None: + markers: Final = tuple(uuid.uuid4().hex for _ in range(12)) + aborted: Final = frozenset(marker for index, marker in enumerate(markers) if index % 3 == 0) + + def respond(request: Request) -> Reply: + marker: Final = _marker_of(request) + return _text_reply(marker, stream=True, abort_after=0 if marker in aborted else None) + + with wire_server(respond) as wire: + path, fleet = _fleet_config(wire, tmp_path) + calls: Final = tuple( + _Call(model=fleet.kimi_bare, endpoint=_ENDPOINTS[index % 3], stream=True, marker=marker) + for index, marker in enumerate(markers) + ) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + served: Final = await _burst(str(candidate.client.base_url), candidate.key, calls) + assert len(served) == 12 + for item in served: + if item.call.marker in aborted: + assert "marker-" not in item.text, item.text + assert item.status >= 500 or "error" in item.text.lower(), (item.status, item.text) + else: + _assert_answered_with_its_own_marker(item) + recovery: Final = _Call(model=fleet.kimi_bare, endpoint="chat", stream=True, marker=uuid.uuid4().hex) + (recovered,) = await _burst(str(candidate.client.base_url), candidate.key, (recovery,)) + _assert_answered_with_its_own_marker(recovered) + _assert_schemas_by_marker(wire.drain(), (*calls, recovery), fleet) diff --git a/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py b/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py new file mode 100644 index 00000000000..97b64c663bb --- /dev/null +++ b/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py @@ -0,0 +1,833 @@ +import json +import threading +import time +from collections.abc import Mapping, Sequence +from typing import Final, Literal +from urllib.parse import unquote + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_KIMI: Final = "global.moonshotai.kimi-k3" +_GROK: Final = "us.xai.grok-4.7" +_NOVA: Final = "us.amazon.nova-lite-v1:0" +_CLAUDE: Final = "global.anthropic.claude-opus-4-8" +_PROFILE_ARN: Final = "arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/lookaround0" +_AWS: Final[dict[str, JsonValue]] = { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", +} +_NO_CACHE: Final[dict[str, JsonValue]] = {"cache": {"no-cache": True}, "num_retries": 0} +_LOOKAHEAD: Final = r"^(?!\.\.?(?:\/|$))[A-Za-z0-9_\-.~:@+]{1,200}$" +_NEGATIVE_LOOKBEHIND: Final = r"^(? dict[str, JsonValue]: + return { + "type": schema["type"], + "properties": schema.get("properties", {}), + "required": schema.get("required", []), + } + + +_WIRE_AS_SENT: Final = _converse_root(_SCHEMA_AS_SENT) +_WIRE_LOOKAROUND_FREE: Final = _converse_root(_SCHEMA_LOOKAROUND_FREE) +_WIRE_PLAIN: Final = _converse_root(_PLAIN_SCHEMA) + +Endpoint = Literal["chat", "messages", "responses"] +_ENDPOINTS: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses") + + +def _frame(event_type: str, payload: Mapping[str, JsonValue]) -> bytes: + return _aws_event_frame(event_type, payload, "sc", "u") + + +_TOOL_USE_RESPONSE: Final = json.dumps( + { + "output": { + "message": { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "tooluse_lookaround_1", "name": _TOOL, "input": _TOOL_INPUT}}], + } + }, + "stopReason": "tool_use", + "usage": _USAGE, + "metrics": {"latencyMs": 1}, + } +).encode() +_STREAM_FRAMES: Final = b"".join( + ( + _frame("messageStart", {"role": "assistant"}), + _frame("contentBlockDelta", {"delta": {"text": _ANSWER}, "contentBlockIndex": 0}), + _frame("contentBlockStop", {"contentBlockIndex": 0}), + _frame("messageStop", {"stopReason": "end_turn"}), + _frame("metadata", {"usage": _USAGE}), + ) +) + + +def _bedrock_peer(request: Request) -> Reply: + if unquote(request.target).endswith("/converse-stream"): + return Reply(body=_STREAM_FRAMES, content_type=_EVENT_STREAM) + return Reply(body=_TOOL_USE_RESPONSE) + + +def _rejecting_peer(request: Request) -> Reply: + return Reply(status=400, body=json.dumps({"message": _BEDROCK_REJECTION}).encode()) + + +def _openai_tool(name: str, schema: Mapping[str, JsonValue], **extra: JsonValue) -> dict[str, JsonValue]: + return { + "type": "function", + "function": {"name": name, "description": f"{name} tool", "parameters": dict(schema), **extra}, + } + + +def _anthropic_tool(name: str, schema: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {"name": name, "description": f"{name} tool", "input_schema": dict(schema)} + + +def _responses_tool(name: str, schema: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {"type": "function", "name": name, "description": f"{name} tool", "parameters": dict(schema)} + + +def _tool_for(endpoint: Endpoint, name: str, schema: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + match endpoint: + case "chat": + return _openai_tool(name, schema) + case "messages": + return _anthropic_tool(name, schema) + case "responses": + return _responses_tool(name, schema) + + +def _path(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "/v1/chat/completions" + case "messages": + return "/v1/messages" + case "responses": + return "/v1/responses" + + +def _body( + endpoint: Endpoint, + model: str, + tools: Sequence[Mapping[str, JsonValue]], + *, + stream: bool = False, + **extra: JsonValue, +) -> dict[str, JsonValue]: + tool_list: Final[list[JsonValue]] = [dict(tool) for tool in tools] + match endpoint: + case "chat": + return { + "model": model, + "messages": [{"role": "user", "content": _PROMPT}], + "max_tokens": 64, + "stream": stream, + "tools": tool_list, + **_NO_CACHE, + **extra, + } + case "messages": + return { + "model": model, + "messages": [{"role": "user", "content": _PROMPT}], + "max_tokens": 64, + "stream": stream, + "tools": tool_list, + **_NO_CACHE, + **extra, + } + case "responses": + return { + "model": model, + "input": _PROMPT, + "max_output_tokens": 64, + "stream": stream, + "tools": tool_list, + **_NO_CACHE, + **extra, + } + + +def _deployment( + scenario: Scenario, + wire: Wire, + model: str, + *, + model_info: Mapping[str, JsonValue] | None = None, + **params: JsonValue, +) -> str: + return scenario.model(model=model, api_base=wire.url, **_AWS, **params, model_info=model_info) + + +def _received_specs(wire: Wire) -> tuple[dict[str, JsonValue], ...]: + received: Final = wire.drain() + assert len(received) == 1, [request.target for request in received] + body: Final = _JSON.validate_json(received[0].body) + tools: Final = _LIST.validate_python(object_value(body["toolConfig"])["tools"]) + return tuple(object_value(object_value(tool)["toolSpec"]) for tool in tools) + + +def _schema_of(spec: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return object_value(object_value(spec["inputSchema"])["json"]) + + +def _only_schema(wire: Wire) -> dict[str, JsonValue]: + (spec,) = _received_specs(wire) + assert spec["name"] == _TOOL, spec + return _schema_of(spec) + + +def _assert_tool_call_relayed(endpoint: Endpoint, response: httpx.Response) -> None: + assert response.status_code == 200, response.text + body: Final = _JSON.validate_json(response.content) + match endpoint: + case "chat": + message: Final = object_value(object_value(_LIST.validate_python(body["choices"])[0])["message"]) + (call,) = _LIST.validate_python(message["tool_calls"]) + function: Final = object_value(object_value(call)["function"]) + assert function["name"] == _TOOL and json.loads(string_value(function["arguments"])) == _TOOL_INPUT, ( + response.text + ) + case "messages": + blocks: Final = tuple(object_value(block) for block in _LIST.validate_python(body["content"])) + (tool_use,) = tuple(block for block in blocks if block.get("type") == "tool_use") + assert tool_use["name"] == _TOOL and tool_use["input"] == _TOOL_INPUT, response.text + case "responses": + items: Final = tuple(object_value(item) for item in _LIST.validate_python(body["output"])) + (call_item,) = tuple(item for item in items if item.get("type") == "function_call") + assert call_item["name"] == _TOOL and json.loads(string_value(call_item["arguments"])) == _TOOL_INPUT, ( + response.text + ) + + +def _stream_text(gateway: Gateway, endpoint: Endpoint, body: Mapping[str, JsonValue]) -> str: + headers: Final = {"Authorization": f"Bearer {gateway.key}"} + with gateway.client.stream("POST", _path(endpoint), json=body, headers=headers) as response: + lines: Final = tuple(line for line in response.iter_lines() if line) + assert response.status_code == 200, "\n".join(lines) + return "\n".join(lines) + + +def _openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) + + +def _async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) + + +def _schema_sent_through( + gateway: Gateway, wire: Wire, endpoint: Endpoint, model: str, tool: Mapping[str, JsonValue], **extra: JsonValue +) -> dict[str, JsonValue]: + response: Final = gateway.request("POST", _path(endpoint), _body(endpoint, model, (tool,), **extra)) + _assert_tool_call_relayed(endpoint, response) + return _only_schema(wire) + + +@pytest.mark.parametrize("endpoint", _ENDPOINTS) +def test_flagged_model_receives_a_lookaround_free_schema_and_the_tool_call_comes_back( + gateway: Gateway, endpoint: Endpoint +) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + tool: Final = _tool_for(endpoint, _TOOL, _SCHEMA_AS_SENT) + assert _schema_sent_through(gateway, wire, endpoint, model, tool) == _WIRE_LOOKAROUND_FREE + + +@pytest.mark.parametrize("endpoint", _ENDPOINTS) +def test_flagged_model_streams_after_the_schema_lost_its_lookarounds(gateway: Gateway, endpoint: Endpoint) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + tool: Final = _tool_for(endpoint, _TOOL, _SCHEMA_AS_SENT) + streamed: Final = _stream_text(gateway, endpoint, _body(endpoint, model, (tool,), stream=True)) + assert _ANSWER in streamed, streamed + received: Final = wire.drain() + assert len(received) == 1 and unquote(received[0].target).endswith("/converse-stream"), received + (tool_block,) = _LIST.validate_python( + object_value(_JSON.validate_json(received[0].body)["toolConfig"])["tools"] + ) + assert _schema_of(object_value(object_value(tool_block)["toolSpec"])) == _WIRE_LOOKAROUND_FREE + + +def test_openai_sdk_sync_chat_sends_a_lookaround_free_schema(gateway: Gateway) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + client: Final = _openai_client(gateway) + completion: Final = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": _PROMPT}], + tools=[_openai_tool(_TOOL, _SCHEMA_AS_SENT)], + max_tokens=64, + extra_body=_NO_CACHE, + ) + (call,) = completion.choices[0].message.tool_calls or () + assert call.function.name == _TOOL and json.loads(call.function.arguments) == _TOOL_INPUT + assert _only_schema(wire) == _WIRE_LOOKAROUND_FREE + chunks: Final = tuple( + client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": _PROMPT}], + tools=[_openai_tool(_TOOL, _SCHEMA_AS_SENT)], + max_tokens=64, + stream=True, + extra_body=_NO_CACHE, + ) + ) + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == _ANSWER + assert _only_schema(wire) == _WIRE_LOOKAROUND_FREE + + +async def test_openai_sdk_async_chat_sends_a_lookaround_free_schema(gateway: Gateway) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + client: Final = _async_openai_client(gateway) + completion: Final = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": _PROMPT}], + tools=[_openai_tool(_TOOL, _SCHEMA_AS_SENT)], + max_tokens=64, + extra_body=_NO_CACHE, + ) + (call,) = completion.choices[0].message.tool_calls or () + assert call.function.name == _TOOL and json.loads(call.function.arguments) == _TOOL_INPUT + assert _only_schema(wire) == _WIRE_LOOKAROUND_FREE + stream: Final = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": _PROMPT}], + tools=[_openai_tool(_TOOL, _SCHEMA_AS_SENT)], + max_tokens=64, + stream=True, + extra_body=_NO_CACHE, + ) + text: Final = "".join([chunk.choices[0].delta.content or "" async for chunk in stream if chunk.choices]) + assert text == _ANSWER + assert _only_schema(wire) == _WIRE_LOOKAROUND_FREE + + +def test_anthropic_sdk_sync_messages_send_a_lookaround_free_schema(gateway: Gateway) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + message: Final = client.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": _PROMPT}], + tools=[_anthropic_tool(_TOOL, _SCHEMA_AS_SENT)], + extra_body=_NO_CACHE, + ) + (tool_use,) = tuple(block for block in message.content if block.type == "tool_use") + assert tool_use.name == _TOOL and tool_use.input == _TOOL_INPUT + assert _only_schema(wire) == _WIRE_LOOKAROUND_FREE + with client.messages.stream( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": _PROMPT}], + tools=[_anthropic_tool(_TOOL, _SCHEMA_AS_SENT)], + extra_body=_NO_CACHE, + ) as stream: + text: Final = "".join(stream.text_stream) + assert text == _ANSWER + assert _only_schema(wire) == _WIRE_LOOKAROUND_FREE + + +async def test_anthropic_sdk_async_messages_send_a_lookaround_free_schema(gateway: Gateway) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + client: Final = anthropic.AsyncAnthropic( + base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0 + ) + message: Final = await client.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": _PROMPT}], + tools=[_anthropic_tool(_TOOL, _SCHEMA_AS_SENT)], + extra_body=_NO_CACHE, + ) + (tool_use,) = tuple(block for block in message.content if block.type == "tool_use") + assert tool_use.name == _TOOL and tool_use.input == _TOOL_INPUT + assert _only_schema(wire) == _WIRE_LOOKAROUND_FREE + async with client.messages.stream( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": _PROMPT}], + tools=[_anthropic_tool(_TOOL, _SCHEMA_AS_SENT)], + extra_body=_NO_CACHE, + ) as stream: + text: Final = "".join([piece async for piece in stream.text_stream]) + assert text == _ANSWER + assert _only_schema(wire) == _WIRE_LOOKAROUND_FREE + + +def test_openai_sdk_sync_responses_send_a_lookaround_free_schema(gateway: Gateway) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + client: Final = _openai_client(gateway) + response: Final = client.responses.create( + model=model, + input=_PROMPT, + tools=[_responses_tool(_TOOL, _SCHEMA_AS_SENT)], + max_output_tokens=64, + extra_body=_NO_CACHE, + ) + (call,) = tuple(item for item in response.output if item.type == "function_call") + assert call.name == _TOOL and json.loads(call.arguments) == _TOOL_INPUT + assert _only_schema(wire) == _WIRE_LOOKAROUND_FREE + events: Final = tuple( + client.responses.create( + model=model, + input=_PROMPT, + tools=[_responses_tool(_TOOL, _SCHEMA_AS_SENT)], + max_output_tokens=64, + stream=True, + extra_body=_NO_CACHE, + ) + ) + deltas: Final = "".join(event.delta for event in events if event.type == "response.output_text.delta") + assert deltas == _ANSWER + assert _only_schema(wire) == _WIRE_LOOKAROUND_FREE + + +async def test_openai_sdk_async_responses_send_a_lookaround_free_schema(gateway: Gateway) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + client: Final = _async_openai_client(gateway) + response: Final = await client.responses.create( + model=model, + input=_PROMPT, + tools=[_responses_tool(_TOOL, _SCHEMA_AS_SENT)], + max_output_tokens=64, + extra_body=_NO_CACHE, + ) + (call,) = tuple(item for item in response.output if item.type == "function_call") + assert call.name == _TOOL and json.loads(call.arguments) == _TOOL_INPUT + assert _only_schema(wire) == _WIRE_LOOKAROUND_FREE + stream: Final = await client.responses.create( + model=model, + input=_PROMPT, + tools=[_responses_tool(_TOOL, _SCHEMA_AS_SENT)], + max_output_tokens=64, + stream=True, + extra_body=_NO_CACHE, + ) + deltas: Final = "".join([event.delta async for event in stream if event.type == "response.output_text.delta"]) + assert deltas == _ANSWER + assert _only_schema(wire) == _WIRE_LOOKAROUND_FREE + + +def test_grok_on_the_explicit_converse_route_is_flagged_too(gateway: Gateway) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/converse/{_GROK}") + tool: Final = _openai_tool(_TOOL, _SCHEMA_AS_SENT) + assert _schema_sent_through(gateway, wire, "chat", model, tool) == _WIRE_LOOKAROUND_FREE + + +def test_a_tool_without_lookarounds_beside_a_cleaned_one_is_forwarded_untouched(gateway: Gateway) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + tools: Final = (_openai_tool(_TOOL, _SCHEMA_AS_SENT), _openai_tool(_PLAIN_TOOL, _PLAIN_SCHEMA)) + response: Final = gateway.request("POST", _path("chat"), _body("chat", model, tools)) + _assert_tool_call_relayed("chat", response) + cleaned, plain = _received_specs(wire) + assert (cleaned["name"], _schema_of(cleaned)) == (_TOOL, _WIRE_LOOKAROUND_FREE) + assert plain == { + "name": _PLAIN_TOOL, + "description": f"{_PLAIN_TOOL} tool", + "inputSchema": {"json": _WIRE_PLAIN}, + }, plain + + +@pytest.mark.parametrize("model_id", (_NOVA, _CLAUDE)) +def test_models_without_the_flag_keep_their_schema_as_sent(gateway: Gateway, model_id: str) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/converse/{model_id}") + tool: Final = _openai_tool(_TOOL, _SCHEMA_AS_SENT) + assert _schema_sent_through(gateway, wire, "chat", model, tool) == _WIRE_AS_SENT + + +@pytest.mark.parametrize( + ("model_id", "model_info", "params", "expected"), + ( + (_KIMI, {"supports_regex_lookaround": True}, {}, _WIRE_AS_SENT), + (_NOVA, {"supports_regex_lookaround": False}, {}, _WIRE_LOOKAROUND_FREE), + (_PROFILE_ARN, None, {"base_model": f"bedrock/{_KIMI}"}, _WIRE_LOOKAROUND_FREE), + (_PROFILE_ARN, None, {}, _WIRE_AS_SENT), + (_KIMI, {"supports_regex_lookaround": None}, {}, _WIRE_LOOKAROUND_FREE), + (_NOVA, {"supports_regex_lookaround": "false"}, {}, _WIRE_AS_SENT), + (_PROFILE_ARN, {"supports_regex_lookaround": True}, {"base_model": f"bedrock/{_KIMI}"}, _WIRE_AS_SENT), + (_KIMI, None, {"base_model": ""}, _WIRE_LOOKAROUND_FREE), + ), + ids=( + "deployment-true-wins-over-map", + "deployment-false-flags-an-unflagged-model", + "base-model-flags-a-profile-arn", + "bare-profile-arn-keeps-the-schema", + "null-falls-back-to-the-map", + "string-false-is-not-a-flag", + "deployment-true-wins-over-base-model", + "empty-base-model-falls-back-to-the-model", + ), +) +def test_deployment_settings_decide_before_the_cost_map( + gateway: Gateway, + model_id: str, + model_info: Mapping[str, JsonValue] | None, + params: Mapping[str, JsonValue], + expected: Mapping[str, JsonValue], +) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{model_id}", model_info=model_info, **params) + tool: Final = _openai_tool(_TOOL, _SCHEMA_AS_SENT) + assert _schema_sent_through(gateway, wire, "chat", model, tool) == expected + + +@pytest.mark.parametrize( + ("model_id", "flag", "expected_for_the_bare_sibling"), + ((_KIMI, True, _WIRE_LOOKAROUND_FREE), (_NOVA, False, _WIRE_AS_SENT)), + ids=("kimi-sibling-keeps-the-map-false", "nova-sibling-keeps-the-map-absence"), +) +@pytest.mark.parametrize("flagged_first", (True, False), ids=("flagged-registered-first", "bare-registered-first")) +def test_a_deployment_flag_never_reaches_its_sibling_on_the_same_model( + gateway: Gateway, + model_id: str, + flag: bool, + expected_for_the_bare_sibling: Mapping[str, JsonValue], + flagged_first: bool, +) -> None: + flag_info: Final[dict[str, JsonValue]] = {"supports_regex_lookaround": flag} + first_info, second_info = (flag_info, None) if flagged_first else (None, flag_info) + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + first: Final = _deployment(scenario, wire, f"bedrock/{model_id}", model_info=first_info) + second: Final = _deployment(scenario, wire, f"bedrock/{model_id}", model_info=second_info) + bare: Final = second if flagged_first else first + tool: Final = _openai_tool(_TOOL, _SCHEMA_AS_SENT) + assert _schema_sent_through(gateway, wire, "chat", bare, tool) == expected_for_the_bare_sibling + + +@pytest.mark.parametrize( + ("model_id", "body_base_model", "expected"), + ((_NOVA, f"bedrock/{_KIMI}", _WIRE_LOOKAROUND_FREE), (_KIMI, f"bedrock/{_NOVA}", _WIRE_LOOKAROUND_FREE)), + ids=("client-base-model-can-loosen-an-unflagged-deployment", "client-base-model-cannot-restore-a-flagged-one"), +) +def test_a_base_model_in_the_request_body_only_ever_loosens( + gateway: Gateway, model_id: str, body_base_model: str, expected: Mapping[str, JsonValue] +) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{model_id}") + tool: Final = _openai_tool(_TOOL, _SCHEMA_AS_SENT) + assert _schema_sent_through(gateway, wire, "chat", model, tool, base_model=body_base_model) == expected + + +@pytest.mark.parametrize( + ("subschema", "expected"), + ( + ( + { + "type": "object", + "patternProperties": {_POSITIVE_LOOKAHEAD_KEY: {"type": "string"}, r"^y_(?!z)": {"type": "integer"}}, + "additionalProperties": False, + }, + { + "type": "object", + "patternProperties": {}, + "additionalProperties": {"anyOf": [{"type": "string"}, {"type": "integer"}]}, + }, + ), + ( + {"type": "object", "patternProperties": {_POSITIVE_LOOKAHEAD_KEY: {"type": "string"}}}, + {"type": "object", "patternProperties": {}}, + ), + ( + {"type": "object", "properties": {"name": {"type": "string", "pattern": r"\(?=x"}}}, + {"type": "object", "properties": {"name": {"type": "string"}}}, + ), + ( + { + "type": "object", + "properties": {"name": {"type": "string"}}, + "dependencies": {"name": {"properties": {"alias": {"type": "string", "pattern": _LOOKAHEAD}}}}, + }, + { + "type": "object", + "properties": {"name": {"type": "string"}}, + "dependencies": {"name": {"properties": {"alias": {"type": "string", "pattern": _LOOKAHEAD}}}}, + }, + ), + ), + ids=( + "two-dropped-pattern-properties-become-an-anyof", + "an-open-object-just-loses-the-key", + "an-escaped-literal-spelling-an-opener-is-dropped-too", + "draft-07-dependencies-are-not-walked", + ), +) +def test_schema_shapes_at_the_edges_of_the_walk( + gateway: Gateway, subschema: Mapping[str, JsonValue], expected: Mapping[str, JsonValue] +) -> None: + schema: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"labels": dict(subschema)}} + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + assert _schema_sent_through(gateway, wire, "chat", model, _openai_tool(_TOOL, schema)) == { + "type": "object", + "properties": {"labels": dict(expected)}, + "required": [], + } + + +def test_strict_is_still_withheld_from_a_flagged_non_anthropic_model(gateway: Gateway) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + tool: Final = _openai_tool(_TOOL, _SCHEMA_AS_SENT, strict=True) + response: Final = gateway.request("POST", _path("chat"), _body("chat", model, (tool,))) + _assert_tool_call_relayed("chat", response) + (spec,) = _received_specs(wire) + assert spec == {"name": _TOOL, "description": f"{_TOOL} tool", "inputSchema": {"json": _WIRE_LOOKAROUND_FREE}} + + +def test_a_json_schema_response_format_rides_the_same_tool_path(gateway: Gateway) -> None: + schema: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"collection": {"type": "string", "pattern": _LOOKAHEAD}}, + "required": ["collection"], + } + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + response: Final = gateway.request( + "POST", + _path("chat"), + { + "model": model, + "messages": [{"role": "user", "content": _PROMPT}], + "max_tokens": 64, + "response_format": {"type": "json_schema", "json_schema": {"name": "document", "schema": schema}}, + **_NO_CACHE, + }, + ) + assert response.status_code == 200, response.text + (spec,) = _received_specs(wire) + assert spec["name"] == "json_tool_call", spec + assert _schema_of(spec) == { + "type": "object", + "properties": {"collection": {"type": "string"}}, + "required": ["collection"], + }, spec + + +@pytest.mark.parametrize( + ("pattern", "expected_property"), + ( + (5, {"type": "string", "pattern": 5}), + ([_LOOKAHEAD], {"type": "string", "pattern": [_LOOKAHEAD]}), + ("", {"type": "string", "pattern": ""}), + ("a" * 5120, {"type": "string", "pattern": "a" * 5120}), + ("a" * 5120 + "(?=b)", {"type": "string"}), + ), + ids=("int", "list", "empty", "5kb-plain", "5kb-ending-in-a-lookahead"), +) +def test_odd_pattern_values_are_forwarded_unless_they_are_a_lookaround_string( + gateway: Gateway, pattern: JsonValue, expected_property: Mapping[str, JsonValue] +) -> None: + schema: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "collection": {"type": "string", "pattern": pattern}, + "doc_id": {"type": "string", "pattern": pattern}, + }, + } + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + assert _schema_sent_through(gateway, wire, "chat", model, _openai_tool(_TOOL, schema)) == { + "type": "object", + "properties": {"collection": dict(expected_property), "doc_id": dict(expected_property)}, + "required": [], + } + + +@pytest.mark.parametrize( + "parameters", + (None, {"type": "object", "properties": [{"name": "collection", "pattern": _LOOKAHEAD}]}), + ids=("null-parameters", "properties-as-a-list"), +) +def test_malformed_tool_parameters_never_take_the_proxy_down(gateway: Gateway, parameters: JsonValue) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + tool: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": _TOOL, "description": f"{_TOOL} tool", "parameters": parameters}, + } + response: Final = gateway.request("POST", _path("chat"), _body("chat", model, (tool,))) + assert response.status_code in (200, 400), response.text + if response.status_code == 400: + assert "error" in _JSON.validate_json(response.content), response.text + wire.drain() + control: Final = gateway.request( + "POST", _path("chat"), _body("chat", model, (_openai_tool(_TOOL, _SCHEMA_AS_SENT),)) + ) + _assert_tool_call_relayed("chat", control) + assert _only_schema(wire) == _WIRE_LOOKAROUND_FREE + + +def test_an_unauthenticated_request_never_reaches_the_peer(gateway: Gateway) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + response: Final = gateway.request( + "POST", _path("chat"), _body("chat", model, (_openai_tool(_TOOL, _SCHEMA_AS_SENT),)), key="sk-not-a-key" + ) + assert response.status_code == 401, response.text + assert wire.drain() == () + + +def test_a_bedrock_rejection_of_an_unflagged_model_reaches_the_caller(gateway: Gateway) -> None: + with wire_server(_rejecting_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/converse/{_CLAUDE}") + response: Final = gateway.request( + "POST", _path("chat"), _body("chat", model, (_openai_tool(_TOOL, _SCHEMA_AS_SENT),)) + ) + assert response.status_code == 400, response.text + assert _BEDROCK_REJECTION in response.text, response.text + assert _only_schema(wire) == _WIRE_AS_SENT + + +@pytest.mark.timeout(120) +def test_the_worst_case_lookaround_input_scans_in_linear_time(gateway: Gateway) -> None: + pattern: Final = "(?<" * (2 * 1024 * 1024 // 3) + schema: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"collection": {"type": "string", "pattern": pattern}}, + } + liveliness: Final[list[tuple[float, int]]] = [] + stop: Final = threading.Event() + + def poll() -> None: + while not stop.is_set(): + liveliness.append(_timed_liveliness(gateway)) + + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}") + poller: Final = threading.Thread(target=poll) + poller.start() + started: Final = time.perf_counter() + response: Final = gateway.request("POST", _path("chat"), _body("chat", model, (_openai_tool(_TOOL, schema),))) + elapsed: Final = time.perf_counter() - started + stop.set() + poller.join() + _assert_tool_call_relayed("chat", response) + assert elapsed < 30, elapsed + assert liveliness and max(latency for latency, _ in liveliness) < 5, liveliness + assert {status for _, status in liveliness} == {200}, liveliness + assert len(wire.drain()) == 1 + + +def _timed_liveliness(gateway: Gateway) -> tuple[float, int]: + started: Final = time.perf_counter() + probe: Final = gateway.client.get("/health/liveliness") + return time.perf_counter() - started, probe.status_code + + +def _model_id(gateway: Gateway, name: str) -> str: + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list), entries + (identity,) = ( + string_value(object_value(object_value(entry)["model_info"])["id"]) + for entry in entries + if object_value(entry)["model_name"] == name + ) + return identity + + +def _settled_schema(gateway: Gateway, wire: Wire, model: str, expected: Mapping[str, JsonValue]) -> None: + tool: Final = _openai_tool(_TOOL, _SCHEMA_AS_SENT) + eventually( + lambda: tuple(_schema_sent_through(gateway, wire, "chat", model, tool) for _ in range(8)), + lambda schemas: all(schema == expected for schema in schemas), + seconds=90, + ) + + +def _patch_flag(gateway: Gateway, identity: str, flag: bool) -> None: + patched: Final = gateway.request( + "PATCH", f"/model/{identity}/update", {"model_info": {"supports_regex_lookaround": flag}} + ) + assert patched.status_code == 200, patched.text + + +@pytest.mark.timeout(300) +def test_updating_the_flag_on_a_live_deployment_takes_effect_without_a_restart(gateway: Gateway) -> None: + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, f"bedrock/{_KIMI}", model_info={"supports_regex_lookaround": True}) + _settled_schema(gateway, wire, model, _WIRE_AS_SENT) + identity: Final = _model_id(gateway, model) + _patch_flag(gateway, identity, False) + _settled_schema(gateway, wire, model, _WIRE_LOOKAROUND_FREE) + _patch_flag(gateway, identity, True) + _settled_schema(gateway, wire, model, _WIRE_AS_SENT) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 0375ff14852..7415c74226d 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -30,6 +30,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) _ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' +_ARTIFACT_DATA_ID_PATTERN: Final = r"^(?!\.\.?(?:\/|$))[A-Za-z0-9_\-.~:@+]{1,200}$" def test_get_format_from_file_id(): @@ -1620,39 +1621,74 @@ class TestToolWithSanitizedParameters: assert tool_with_sanitized_parameters(tool, flatten_combinators_and_drop_non_python_regex_patterns) is tool + def test_sanitizes_the_input_schema_of_an_anthropic_tool(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_lookaround_regex_patterns, + tool_with_sanitized_parameters, + ) + + tool = { + "name": "ArtifactData", + "description": "Read a shared database", + "input_schema": { + "type": "object", + "properties": {"doc_id": {"type": "string", "pattern": _ARTIFACT_DATA_ID_PATTERN}}, + }, + } + + result = tool_with_sanitized_parameters(tool, drop_lookaround_regex_patterns) + + assert result == { + "name": "ArtifactData", + "description": "Read a shared database", + "input_schema": {"type": "object", "properties": {"doc_id": {"type": "string"}}}, + } + assert tool["input_schema"]["properties"]["doc_id"]["pattern"] == _ARTIFACT_DATA_ID_PATTERN + + def test_returns_the_same_anthropic_tool_when_its_schema_has_nothing_to_drop(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_lookaround_regex_patterns, + tool_with_sanitized_parameters, + ) + + tool = {"name": "Read", "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}} + + assert tool_with_sanitized_parameters(tool, drop_lookaround_regex_patterns) is tool + + +def _regex_schema(pattern): + return { + "type": "object", + "properties": { + "field": {"type": "string", "pattern": pattern}, + "writes": { + "type": "array", + "items": {"properties": {"doc_id": {"type": "string", "pattern": pattern}}}, + }, + "query": {"anyOf": [{"type": "string", "pattern": pattern}, {"type": "null"}]}, + "pair": {"type": "array", "prefixItems": [{"type": "string", "pattern": pattern}]}, + "extra": {"type": "object", "additionalProperties": {"type": "string", "pattern": pattern}}, + "tagged": { + "type": "object", + "patternProperties": {pattern: {"type": "string"}, "^x_": {"type": "integer"}}, + }, + }, + "$defs": {"segment": {"type": "string", "pattern": pattern}}, + "required": ["field"], + } + class TestDropNonPythonRegexPatterns: """Claude Code's Artifact tool declares ECMA-262 ``\\p{..}`` escapes that OpenAI's validator, which compiles ``pattern`` values and ``patternProperties`` keys with Python ``re``, refuses as "not a 'regex'".""" - def _schema(self, pattern): - return { - "type": "object", - "properties": { - "field": {"type": "string", "pattern": pattern}, - "writes": { - "type": "array", - "items": {"properties": {"doc_id": {"type": "string", "pattern": pattern}}}, - }, - "query": {"anyOf": [{"type": "string", "pattern": pattern}, {"type": "null"}]}, - "pair": {"type": "array", "prefixItems": [{"type": "string", "pattern": pattern}]}, - "extra": {"type": "object", "additionalProperties": {"type": "string", "pattern": pattern}}, - "tagged": { - "type": "object", - "patternProperties": {pattern: {"type": "string"}, "^x_": {"type": "integer"}}, - }, - }, - "$defs": {"segment": {"type": "string", "pattern": pattern}}, - "required": ["field"], - } - def test_drops_every_regex_python_re_rejects_from_every_schema_position(self): from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_non_python_regex_patterns, ) - schema = self._schema(_ARTIFACT_FIELD_PATTERN) + schema = _regex_schema(_ARTIFACT_FIELD_PATTERN) result = drop_non_python_regex_patterns(schema) @@ -1666,14 +1702,14 @@ class TestDropNonPythonRegexPatterns: assert properties["tagged"]["patternProperties"] == {"^x_": {"type": "integer"}} assert result["$defs"]["segment"] == {"type": "string"} assert result["required"] == ["field"] - assert schema == self._schema(_ARTIFACT_FIELD_PATTERN) + assert schema == _regex_schema(_ARTIFACT_FIELD_PATTERN) def test_keeps_regexes_python_re_compiles_and_returns_the_same_object(self): from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_non_python_regex_patterns, ) - schema = self._schema(r'^(?!__.*__$)[^"\\./[\]]{1,200}$') + schema = _regex_schema(r'^(?!__.*__$)[^"\\./[\]]{1,200}$') assert drop_non_python_regex_patterns(schema) is schema @@ -1737,6 +1773,137 @@ class TestDropNonPythonRegexPatterns: assert drop_non_python_regex_patterns(schema) is schema +class TestDropLookaroundRegexPatterns: + """Kimi K3 and Grok 4.6/4.7 on Bedrock Converse reject every tool schema regex that + uses a lookaround assertion, Claude Code's ``ArtifactData`` ``pattern`` included.""" + + @pytest.mark.parametrize( + "pattern", + [r"^(?!x).*$", r"^(?=.*a).*$", r"^.*(?\w+)$", r"^(?i)abc$", r"^[^\p{Cc}\p{Cf}]{1,200}$"], + ids=["plain", "non-capturing-group", "named-group", "inline-flag", "non-python-without-lookaround"], + ) + def test_keeps_regexes_without_lookaround_and_returns_the_same_object(self, pattern): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_lookaround_regex_patterns, + ) + + schema = _regex_schema(pattern) + + assert drop_lookaround_regex_patterns(schema) is schema + + def test_lookaround_inside_data_positions_is_not_a_regex(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_lookaround_regex_patterns, + ) + + schema = { + "type": "object", + "properties": { + "pattern": {"type": "string"}, + "template": {"type": "object", "default": {"pattern": _ARTIFACT_DATA_ID_PATTERN}}, + "hint": {"type": "string", "description": "ids match " + _ARTIFACT_DATA_ID_PATTERN}, + }, + "required": ["pattern"], + } + + assert drop_lookaround_regex_patterns(schema) is schema + + +@pytest.mark.parametrize( + ("dropper", "patterns"), + [ + ("drop_non_python_regex_patterns", (_ARTIFACT_FIELD_PATTERN, r"^\p{L}+$")), + ("drop_lookaround_regex_patterns", (_ARTIFACT_DATA_ID_PATTERN, r"^(?=.*[a-z])\w+$")), + ], + ids=["non-python", "lookaround"], +) +class TestDroppedPatternPropertiesKeepTheirNamesAllowed: + """Dropping a ``patternProperties`` key from an object closed by ``additionalProperties: + false`` must not ban the names that key allowed: its value schema takes over as the + object's ``additionalProperties``.""" + + @staticmethod + def _drop(dropper): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + drop_lookaround_regex_patterns, + drop_non_python_regex_patterns, + ) + + return { + "drop_non_python_regex_patterns": drop_non_python_regex_patterns, + "drop_lookaround_regex_patterns": drop_lookaround_regex_patterns, + }[dropper] + + def test_closed_object_takes_the_dropped_value_schema(self, dropper, patterns): + schema = { + "type": "object", + "patternProperties": {patterns[0]: {"type": "string", "pattern": patterns[0]}}, + "additionalProperties": False, + } + + assert self._drop(dropper)(schema) == { + "type": "object", + "patternProperties": {}, + "additionalProperties": {"type": "string"}, + } + + def test_closed_object_losing_two_entries_accepts_either_value_schema(self, dropper, patterns): + schema = { + "type": "object", + "patternProperties": { + patterns[0]: {"type": "string"}, + patterns[1]: {"type": "integer"}, + "^x_": {"type": "boolean"}, + }, + "additionalProperties": False, + } + + assert self._drop(dropper)(schema) == { + "type": "object", + "patternProperties": {"^x_": {"type": "boolean"}}, + "additionalProperties": {"anyOf": [{"type": "string"}, {"type": "integer"}]}, + } + + def test_object_with_its_own_additional_properties_schema_keeps_it(self, dropper, patterns): + schema = { + "type": "object", + "patternProperties": {patterns[0]: {"type": "string"}}, + "additionalProperties": {"type": "integer"}, + } + + assert self._drop(dropper)(schema) == { + "type": "object", + "patternProperties": {}, + "additionalProperties": {"type": "integer"}, + } + + class TestRequestContainsImageContent: """One detector for every dialect that reaches pre-routing hooks untranslated.""" diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index f6f98e3b9bd..0c0be5d868f 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -637,6 +637,142 @@ def test_output_config_effort_forwarded_into_additional_request_fields(model): assert additional.get("output_config") == {"effort": "high"} +_ARTIFACT_DATA_ID_PATTERN: Final = r"^(?!\.\.?(?:\/|$))[A-Za-z0-9_\-.~:@+]{1,200}$" +_ARTIFACT_DATA_INPUT_SCHEMA: Final = { + "type": "object", + "properties": { + "collection": {"type": "string", "pattern": _ARTIFACT_DATA_ID_PATTERN, "description": "Collection"}, + "doc_id": {"type": "string", "pattern": _ARTIFACT_DATA_ID_PATTERN}, + "writes": { + "type": "array", + "items": { + "type": "object", + "properties": {"doc_id": {"type": "string", "pattern": _ARTIFACT_DATA_ID_PATTERN}}, + }, + }, + "limit": {"type": "integer", "minimum": 1}, + }, + "required": ["collection"], +} +_ARTIFACT_DATA_ANTHROPIC_TOOL: Final = { + "name": "ArtifactData", + "description": "Read a shared database", + "input_schema": _ARTIFACT_DATA_INPUT_SCHEMA, +} +_ARTIFACT_DATA_OPENAI_TOOL: Final = { + "type": "function", + "function": { + "name": "ArtifactData", + "description": "Read a shared database", + "parameters": _ARTIFACT_DATA_INPUT_SCHEMA, + }, +} +_LOOKAROUND_FREE_PROPERTIES: Final = { + "collection": {"type": "string", "description": "Collection"}, + "doc_id": {"type": "string"}, + "writes": {"type": "array", "items": {"type": "object", "properties": {"doc_id": {"type": "string"}}}}, + "limit": {"type": "integer", "minimum": 1}, +} + + +def _converse_tools(model, tools, litellm_params=None): + request = AmazonConverseConfig()._transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={"tools": copy.deepcopy(tools)}, + litellm_params=litellm_params or {}, + headers={}, + ) + return request["toolConfig"]["tools"] + + +def _tool_schema_properties(model, tool, litellm_params=None): + return _converse_tools(model, [tool], litellm_params)[0]["toolSpec"]["inputSchema"]["json"]["properties"] + + +@pytest.mark.parametrize( + "tool", [_ARTIFACT_DATA_ANTHROPIC_TOOL, _ARTIFACT_DATA_OPENAI_TOOL], ids=["anthropic-shape", "openai-shape"] +) +@pytest.mark.parametrize( + "model", + [ + "global.moonshotai.kimi-k3", + "us.moonshotai.kimi-k3", + "moonshotai.kimi-k3", + "us-east-1/us.moonshotai.kimi-k3", + "us.xai.grok-4.6", + "us-gov.xai.grok-4.6", + "global.xai.grok-4.7", + "xai.grok-4.7", + ], +) +def test_transform_request_drops_lookaround_regex_for_models_the_cost_map_flags(tool, model): + """Kimi K3 and Grok 4.6/4.7 refuse the whole request over a lookaround in a tool schema regex.""" + tools = _converse_tools(model, [tool]) + + json_schema = tools[0]["toolSpec"]["inputSchema"]["json"] + assert json_schema["properties"] == _LOOKAROUND_FREE_PROPERTIES + assert json_schema["required"] == ["collection"] + + +@pytest.mark.parametrize( + "model", + [ + "us.anthropic.claude-sonnet-4-6", + "us.amazon.nova-pro-v1:0", + "us.meta.llama4-maverick-17b-instruct-v1:0", + "us.openai.gpt-5.6-sol", + ], +) +def test_transform_request_keeps_lookaround_regex_for_models_that_accept_it(model): + assert _tool_schema_properties(model, _ARTIFACT_DATA_ANTHROPIC_TOOL) == _ARTIFACT_DATA_INPUT_SCHEMA["properties"] + + +@pytest.mark.parametrize( + "model", + [ + "us.amazon.nova-lite-v1:0", + "us.moonshotai.kimi-k4", + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + ], +) +def test_transform_request_drops_lookaround_regex_when_the_deployment_model_info_opts_in(model): + """A deployment's ``model_info`` flag covers a model the cost map does not know, an inference profile included.""" + properties = _tool_schema_properties( + model, _ARTIFACT_DATA_ANTHROPIC_TOOL, {"model_info": {"supports_regex_lookaround": False}} + ) + + assert properties == _LOOKAROUND_FREE_PROPERTIES + + +def test_transform_request_keeps_lookaround_regex_when_the_deployment_model_info_opts_out(): + properties = _tool_schema_properties( + "global.moonshotai.kimi-k3", _ARTIFACT_DATA_ANTHROPIC_TOOL, {"model_info": {"supports_regex_lookaround": True}} + ) + + assert properties["doc_id"]["pattern"] == _ARTIFACT_DATA_ID_PATTERN + + +def test_transform_request_resolves_an_inference_profile_through_its_base_model(): + properties = _tool_schema_properties( + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + _ARTIFACT_DATA_ANTHROPIC_TOOL, + {"base_model": "bedrock/global.moonshotai.kimi-k3"}, + ) + + assert properties == _LOOKAROUND_FREE_PROPERTIES + + +def test_transform_request_drops_lookaround_regex_around_pre_formatted_tool_blocks(): + """Blocks that arrive already in Bedrock shape, like Nova's grounding ``systemTool``, pass through as sent.""" + grounding: Final = {"systemTool": {"name": "nova_grounding"}} + + tools = _converse_tools("global.moonshotai.kimi-k3", [_ARTIFACT_DATA_OPENAI_TOOL, grounding]) + + assert tools[0]["toolSpec"]["inputSchema"]["json"]["properties"] == _LOOKAROUND_FREE_PROPERTIES + assert tools[1] == grounding + + def test_reasoning_effort_requests_summarized_display_converse(): """Regression LIT-5714: adaptive thinking synthesized from reasoning_effort must request the summarized display, otherwise the provider returns a blank thinking diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 74839831ca1..ff8c91cae70 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -514,6 +514,34 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing(): ) +def test_regex_lookaround_flag_stays_on_the_deployment_that_set_it() -> None: + """A deployment's ``supports_regex_lookaround`` override must not land on the shared + ``{provider}/{model}`` key, or every sibling deployment of that model would inherit it.""" + backend_model = "bedrock/us.xai.grok-4.6" + deploy_id = "grok-deploy-keep-regex" + + builtin_flag = litellm.get_model_info(model=backend_model).get("supports_regex_lookaround") + model_keys = { + deploy_id: litellm.model_cost.get(deploy_id), + backend_model: copy.deepcopy(litellm.model_cost.get(backend_model)), + } + try: + Router( + model_list=[ + { + "model_name": "grok-keep-regex", + "litellm_params": {"model": backend_model}, + "model_info": {"id": deploy_id, "supports_regex_lookaround": not builtin_flag}, + } + ], + ) + + assert litellm.model_cost[deploy_id]["supports_regex_lookaround"] is (not builtin_flag) + assert litellm.get_model_info(model=backend_model).get("supports_regex_lookaround") is builtin_flag + finally: + _restore_model_cost_entries(model_keys) + + def test_should_store_full_pricing_under_deployment_model_id(): """ Per-deployment pricing (including zero) should be stored and diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index ab7ecfcda05..62ac8f2395a 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -996,6 +996,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "enum": ["low", "medium", "high", "max", "xhigh"], }, "bedrock_converse_supports_strict_tools": {"type": "boolean"}, + "supports_regex_lookaround": {"type": "boolean"}, "tpm": {"type": "number"}, "supported_endpoints": { "type": "array",