From 29270a36a5afce9930b224904d83870313c30d0e Mon Sep 17 00:00:00 2001 From: milan-berri Date: Tue, 2 Jun 2026 00:28:31 +0300 Subject: [PATCH] fix(anthropic, fireworks): inline legacy $ref defs in tool schemas (#28646) Tools sourced from MCP servers and OpenAPI-derived gateways (AWS AgentCore + Google Workspace, DevRev MCP, etc.) frequently carry JSON Schemas backed by legacy ``definitions`` (draft-04) or OpenAPI ``components.schemas`` instead of ``$defs`` (JSON Schema 2020-12). Anthropic and Fireworks only resolve ``$defs``. Their tool-schema filters silently drop the unrecognised def blocks while keeping the ``$ref`` pointers, so the upstream rejects the request: - Anthropic: ``tools.0.input_schema: Invalid tool schema, $ref is not supported`` - Fireworks: ``Error resolving schema reference '#/definitions/...'`` (PointerToNowhere) Add ``unpack_legacy_defs(schema, *, copy=False)`` next to the existing ``unpack_defs`` -- a single helper that pops draft-04 ``definitions`` and OpenAPI ``components.schemas`` and feeds them through ``unpack_defs`` in place. ``$defs`` is left untouched (resolved natively). ``copy=True`` deep-copies first when there is actually work to do, used by Anthropic so the caller's tool dict is preserved. Anthropic ``_map_tool_helper`` calls ``unpack_legacy_defs(_, copy=True)``; Fireworks ``_transform_tools`` calls ``unpack_legacy_defs(params)`` in place. Refs: https://github.com/BerriAI/litellm/issues/26692 Co-authored-by: Cursor --- .../prompt_templates/common_utils.py | 117 +++++++++- litellm/llms/anthropic/chat/transformation.py | 5 + .../llms/fireworks_ai/chat/transformation.py | 10 +- ...ore_utils_prompt_templates_common_utils.py | 175 +++++++++++++++ .../test_anthropic_chat_transformation.py | 201 ++++++++++++++++++ .../test_fireworks_ai_chat_transformation.py | 167 +++++++++++++++ 6 files changed, 672 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index b44d21368f..fe34731759 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -850,7 +850,47 @@ def extract_file_data(file_data: FileTypes) -> ExtractedFileData: # --------------------------------------------------------------------------- -def unpack_defs(schema: dict, defs: dict) -> None: +def _estimate_json_bytes(obj: Any) -> int: + """Estimate the JSON-serialised byte size of ``obj`` without materialising + JSON. Walks iteratively (no recursion stack risk). + + String length is read via ``len()`` (O(1) on Python ``str``) so a target + containing a 100MB description costs ~one walk step, not a 100MB + serialisation. Escape sequences are not counted exactly, so this is an + approximation -- but always within a small constant factor of the real + serialised size, which is what a schema-bomb budget needs. + """ + total = 0 + stack: list = [obj] + while stack: + x = stack.pop() + if isinstance(x, dict): + total += 2 # `{}` + for k, v in x.items(): + total += len(str(k)) + 4 # `"k":,` + stack.append(v) + elif isinstance(x, list): + total += 2 # `[]` + total += max(0, len(x) - 1) # commas between items + stack.extend(x) + elif isinstance(x, str): + total += len(x) + 2 + elif isinstance(x, bool): # bool subclasses int -- check first + total += 4 if x else 5 + elif x is None: + total += 4 + elif isinstance(x, (int, float)): + total += 24 # generous upper bound for stringified numbers + else: + total += 24 + return total + + +def unpack_defs( + schema: dict, + defs: dict, + max_inlined_bytes: Optional[int] = None, +) -> None: """Expand *all* ``$ref`` entries pointing into ``$defs`` / ``definitions``. This utility walks the entire schema tree (dicts and lists) so it naturally @@ -860,6 +900,15 @@ def unpack_defs(schema: dict, defs: dict) -> None: It mutates *schema* in-place and does **not** return anything. The helper keeps memory overhead low by resolving nodes as it encounters them rather than materialising a fully dereferenced copy first. + + ``max_inlined_bytes`` caps the cumulative JSON-byte size of every target + that has been inlined and is checked *before* each ``copy.deepcopy``, so + an oversized expansion is rejected without first materialising it. A byte + bound is the universal measure of expansion -- it simultaneously caps + ref-count fan-out, node-count amplification, and scalar-byte amplification + (a target containing a large string, ``const``, or ``enum`` entry). + Defaults to ``None`` (unbounded) so existing callers are unaffected; + raises ``ValueError`` on overflow. """ import copy @@ -879,6 +928,7 @@ def unpack_defs(schema: dict, defs: dict) -> None: queue: deque[ tuple[Any, Union[dict, list, None], Union[str, int, None], dict, set] ] = deque([(schema, None, None, root_defs, set())]) + inlined_bytes = 0 while queue: node, parent, key, active_defs, ref_chain = queue.popleft() @@ -899,6 +949,16 @@ def unpack_defs(schema: dict, defs: dict) -> None: if target_schema is None: continue + if max_inlined_bytes is not None: + inlined_bytes += _estimate_json_bytes(target_schema) + if inlined_bytes > max_inlined_bytes: + raise ValueError( + f"unpack_defs: inlined schema exceeded the " + f"{max_inlined_bytes:,}-byte budget. Refusing to " + f"deep-copy further to prevent schema-bomb " + f"resource exhaustion." + ) + # Merge defs from the target to capture nested definitions child_defs = { **active_defs, @@ -946,6 +1006,61 @@ def unpack_defs(schema: dict, defs: dict) -> None: queue.append((item, node, idx, active_defs, ref_chain)) +def _has_legacy_defs(schema: object) -> bool: + if not isinstance(schema, dict): + return False + components = schema.get("components") + return "definitions" in schema or ( + isinstance(components, dict) and isinstance(components.get("schemas"), dict) + ) + + +# Schema-bomb budget for ``unpack_legacy_defs``: cap the cumulative JSON-byte +# size of every inlined target. A byte cap is the universal measure of +# expansion -- it simultaneously bounds ref-count fan-out, node-count +# amplification, and scalar-byte amplification (large ``description`` / +# ``const`` / ``enum`` values). Real-world MCP / OpenAPI-derived tool schemas +# inline well under 1MB; 10MB sits two orders of magnitude above that, well +# below memory-pressure territory, and rejects request-supplied bombs before +# the proxy materialises them. +_LEGACY_DEFS_MAX_INLINED_BYTES = 10_000_000 + + +def unpack_legacy_defs( + schema: dict, + *, + copy: bool = False, + max_inlined_bytes: int = _LEGACY_DEFS_MAX_INLINED_BYTES, +) -> dict: + """Inline ``$ref``s backed by draft-04 ``definitions`` / OpenAPI + ``components.schemas``. ``$defs`` is left untouched. + + Anthropic and Fireworks tool-schema resolvers only recognise ``$defs``; + legacy / OpenAPI def blocks are otherwise silently dropped and leave + dangling pointers. See https://github.com/BerriAI/litellm/issues/26692. + + Mutates ``schema`` in place and returns it. Pass ``copy=True`` to deep-copy + first (only when there is actually work to do). ``max_inlined_bytes`` + bounds the cumulative JSON-byte size of inlined targets so request-supplied + schemas cannot expand into a schema-bomb before reaching the upstream + provider -- raises ``ValueError`` on overflow. + """ + if not _has_legacy_defs(schema): + return schema + if copy: + import copy as _copy + + schema = _copy.deepcopy(schema) + # On key collision, ``definitions`` wins over ``components.schemas`` -- + # ``unpack_defs`` keys refs by last path segment so a single name can only + # resolve to one body, and ``definitions`` is the JSON-Schema-native + # namespace. + defs = schema.pop("components", {}).get("schemas") or {} + defs.update(schema.pop("definitions", None) or {}) + unpack_defs(schema, defs, max_inlined_bytes=max_inlined_bytes) + return schema + + def _get_image_mime_type_from_url(url: str) -> Optional[str]: """ Get mime type for common image URLs diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 57609cfcd2..4f15d1b3ce 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -29,6 +29,7 @@ from litellm.constants import ( RESPONSE_FORMAT_TOOL_NAME, ) from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_legacy_defs from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.anthropic import ( @@ -680,6 +681,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if "properties" not in _input_schema: _input_schema["properties"] = {} + # Inline legacy / OpenAPI $refs before the allow-list filter strips + # their backing def blocks (https://github.com/BerriAI/litellm/issues/26692). + _input_schema = unpack_legacy_defs(_input_schema, copy=True) + _allowed_properties = set(AnthropicInputSchema.__annotations__.keys()) input_schema_filtered = { k: v for k, v in _input_schema.items() if k in _allowed_properties diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index d39adf0b6f..9e9d300b58 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_legacy_defs from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) @@ -216,8 +217,13 @@ class FireworksAIConfig(OpenAIGPTConfig): self, tools: List[OpenAIChatCompletionToolParam] ) -> List[OpenAIChatCompletionToolParam]: for tool in tools: - if tool.get("type") == "function": - tool["function"].pop("strict", None) + if tool.get("type") != "function": + continue + function = tool["function"] + function.pop("strict", None) + params = function.get("parameters") + if isinstance(params, dict): + unpack_legacy_defs(params) return tools def _transform_messages_helper( diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 2aaeaefce5..1b1db634ed 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -546,3 +546,178 @@ class TestExtractFileDataBareStr: extracted = extract_file_data(("foo.txt", b"raw bytes content")) assert extracted.get("filename") == "foo.txt" assert extracted.get("content") == b"raw bytes content" + + +class TestUnpackLegacyDefs: + """Cover the public ``unpack_legacy_defs`` helper directly so the no-op + branches (non-dict input, schema with no legacy/OpenAPI defs) are exercised + without needing a provider-specific entry point. + """ + + @pytest.mark.parametrize( + "value", + [None, [], "string-not-a-dict", 42, 1.5, True, set(), tuple()], + ) + def test_non_dict_returns_unchanged_no_op(self, value): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + unpack_legacy_defs, + ) + + # Should never raise; returns the input unchanged. + assert unpack_legacy_defs(value) is value + assert unpack_legacy_defs(value, copy=True) is value + + def test_dict_without_legacy_defs_is_no_op(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + unpack_legacy_defs, + ) + + schema = { + "type": "object", + "properties": {"a": {"$ref": "#/$defs/A"}}, + "$defs": {"A": {"type": "string"}}, + } + snapshot = json.loads(json.dumps(schema)) + + # No `definitions` and no `components.schemas` -> early return, no work. + out = unpack_legacy_defs(schema) + assert out is schema + assert schema == snapshot, "schema mutated despite no legacy defs" + + def test_components_with_no_schemas_block_is_no_op(self): + """``components`` without a ``schemas`` sub-key must not be popped.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + unpack_legacy_defs, + ) + + schema = { + "type": "object", + "properties": {"a": {"type": "string"}}, + "components": {"securitySchemes": {"foo": "bar"}}, + } + snapshot = json.loads(json.dumps(schema)) + + unpack_legacy_defs(schema) + assert schema == snapshot, "components without schemas was incorrectly popped" + + def test_legitimate_schema_within_budget_succeeds(self): + """A flat schema with many distinct ``$ref``s into small targets must + inline cleanly under the default budget -- the budget rejects bombs, + not legitimately-shaped schemas. + """ + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + unpack_legacy_defs, + ) + + n = 200 + schema = { + "type": "object", + "properties": {f"f{i}": {"$ref": f"#/definitions/T{i}"} for i in range(n)}, + "definitions": {f"T{i}": {"type": "string"} for i in range(n)}, + } + + out = unpack_legacy_defs(schema) + assert "definitions" not in out + for i in range(n): + assert out["properties"][f"f{i}"] == {"type": "string"} + + # Schema-bomb amplification vectors. ``max_inlined_bytes`` is the universal + # measure of expansion: every other dimension (ref count, node count, + # scalar size) reduces to bytes-on-the-wire, so a single byte budget + # closes all three vectors at once. + + def test_rejects_fan_out_bomb(self): + """Each level multiplies refs (cycle detection only stops re-entry + along the *same* path). Must trip the byte budget.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + unpack_legacy_defs, + ) + + depth, fanout = 12, 2 # 2**12 = 4096 leaves + definitions = { + f"L{i}": { + "type": "object", + "properties": { + f"x{j}": {"$ref": f"#/definitions/L{i + 1}"} for j in range(fanout) + }, + } + for i in range(depth) + } + definitions[f"L{depth}"] = {"type": "string"} + schema = { + "type": "object", + "properties": {"root": {"$ref": "#/definitions/L0"}}, + "definitions": definitions, + } + + with pytest.raises(ValueError, match="byte budget"): + unpack_legacy_defs(schema, max_inlined_bytes=100_000) + + def test_rejects_target_amplification_bomb(self): + """Few refs each deep-copying one large target -- bounded total + expanded bytes catches it even though ref count is small.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + unpack_legacy_defs, + ) + + big = { + "type": "object", + "properties": {f"p{i}": {"type": "string"} for i in range(100)}, + } + schema = { + "type": "object", + "properties": {f"r{i}": {"$ref": "#/definitions/Big"} for i in range(50)}, + "definitions": {"Big": big}, + } + + with pytest.raises(ValueError, match="byte budget"): + unpack_legacy_defs(schema, max_inlined_bytes=10_000) + + def test_rejects_scalar_byte_amplification_bomb(self): + """Many ``$ref``s to a target containing one large scalar (e.g. a + long ``description``, ``const`` value, or ``enum`` entry). A + node-counter would treat this as 1 node per resolution and miss it; + a byte budget catches the actual wire-size amplification. + """ + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + unpack_legacy_defs, + ) + + big_description = "x" * 100_000 # 100KB string + schema = { + "type": "object", + "properties": {f"r{i}": {"$ref": "#/definitions/Big"} for i in range(50)}, + "definitions": { + "Big": {"type": "string", "description": big_description}, + }, + } + # 50 refs * ~100KB string == ~5MB cumulative; 1MB budget trips. + with pytest.raises(ValueError, match="byte budget"): + unpack_legacy_defs(schema, max_inlined_bytes=1_000_000) + + def test_budget_does_not_trip_for_legitimate_large_schema(self): + """An OpenAPI-derived tool with ~50 small targets must inline cleanly + under the default ``max_inlined_bytes`` budget.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + unpack_legacy_defs, + ) + + schema = { + "type": "object", + "properties": { + f"r{i}": {"$ref": f"#/components/schemas/T{i}"} for i in range(50) + }, + "components": { + "schemas": { + f"T{i}": { + "type": "object", + "properties": {f"p{j}": {"type": "string"} for j in range(5)}, + } + for i in range(50) + } + }, + } + + out = unpack_legacy_defs(schema) + assert "components" not in out + assert out["properties"]["r0"]["properties"]["p0"] == {"type": "string"} diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 687c5a2e73..d501ae0f79 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -4889,3 +4889,204 @@ def test_sanitize_tool_names_in_request_no_tools_is_noop(): forward, reverse = AnthropicConfig._sanitize_tool_names_in_request({"tools": []}) assert forward == {} assert reverse == {} + + +# ----------------------------------------------------------------------------- +# Regression tests for legacy / OpenAPI $ref defs in tool input_schema. +# +# Anthropic only resolves `$defs` (JSON Schema 2020-12). Tools coming from MCP +# servers (legacy `definitions`) or OpenAPI-derived gateways like AWS +# AgentCore (`components.schemas`) used to silently lose their def blocks +# while keeping dangling `$ref`s, causing upstream 400s. See +# https://github.com/BerriAI/litellm/issues/26692. +# ----------------------------------------------------------------------------- + + +def _assert_no_unresolved_refs(input_schema: dict) -> None: + import json + + blob = json.dumps(input_schema) + assert "$ref" not in blob, f"unresolved $ref in transformed input_schema: {blob}" + + +def test_map_tool_helper_inlines_components_schemas_refs(): + """OpenAPI `components.schemas` $refs (AgentCore-style) must be inlined.""" + config = AnthropicConfig() + tool = { + "type": "function", + "function": { + "name": "slides_presentations_create", + "description": "Create a Google Slides presentation", + "parameters": { + "type": "object", + "properties": { + "body": {"$ref": "#/components/schemas/Presentation"}, + }, + "required": ["body"], + "components": { + "schemas": { + "Presentation": { + "type": "object", + "properties": { + "title": {"type": "string"}, + "presentationId": {"type": "string"}, + }, + } + } + }, + }, + }, + } + + transformed, _ = config._map_tool_helper(tool) + + assert transformed is not None + schema = transformed["input_schema"] + _assert_no_unresolved_refs(schema) + assert schema["properties"]["body"] == { + "type": "object", + "properties": { + "title": {"type": "string"}, + "presentationId": {"type": "string"}, + }, + } + # The OpenAPI components block is not part of Anthropic's allow-list and + # must not be forwarded. + assert "components" not in schema + + +def test_map_tool_helper_inlines_legacy_definitions_refs(): + """Legacy draft-04 `definitions` $refs (DevRev MCP-style) must be inlined.""" + config = AnthropicConfig() + tool = { + "type": "function", + "function": { + "name": "create_thing", + "description": "Create a thing", + "parameters": { + "type": "object", + "properties": { + "thing": {"$ref": "#/definitions/Thing"}, + }, + "definitions": { + "Thing": { + "type": "object", + "properties": {"id": {"type": "string"}}, + } + }, + }, + }, + } + + transformed, _ = config._map_tool_helper(tool) + + assert transformed is not None + schema = transformed["input_schema"] + _assert_no_unresolved_refs(schema) + assert schema["properties"]["thing"] == { + "type": "object", + "properties": {"id": {"type": "string"}}, + } + assert "definitions" not in schema + + +def test_map_tool_helper_preserves_native_dollar_defs(): + """`$defs` is JSON Schema 2020-12 native; Anthropic resolves it itself. + + Re-implementation must not pop or unpack `$defs`. + """ + config = AnthropicConfig() + tool = { + "type": "function", + "function": { + "name": "native_defs_tool", + "description": "", + "parameters": { + "type": "object", + "properties": {"a": {"$ref": "#/$defs/A"}}, + "$defs": {"A": {"type": "string"}}, + }, + }, + } + + transformed, _ = config._map_tool_helper(tool) + + assert transformed is not None + schema = transformed["input_schema"] + assert schema["$defs"] == {"A": {"type": "string"}} + assert schema["properties"]["a"] == {"$ref": "#/$defs/A"} + + +def test_map_tool_helper_does_not_mutate_caller_dict(): + """Caller-supplied tool dict must not be mutated by the inlining step.""" + import copy + + config = AnthropicConfig() + tool = { + "type": "function", + "function": { + "name": "create_thing", + "description": "Create a thing", + "parameters": { + "type": "object", + "properties": {"thing": {"$ref": "#/definitions/Thing"}}, + "definitions": { + "Thing": { + "type": "object", + "properties": {"id": {"type": "string"}}, + } + }, + }, + }, + } + snapshot = copy.deepcopy(tool) + + config._map_tool_helper(tool) + + assert tool == snapshot, "caller's tool dict was mutated in place" + + +def test_map_tool_helper_collision_prefers_definitions_over_components_schemas(): + """If both `definitions.X` and `components.schemas.X` exist with the same + name, prefer the `definitions` body. ``unpack_defs`` keys refs by last path + segment so only one body can win; pick the JSON-Schema-native one. + + This locks in the residual limitation as a deliberate contract: a ref + written as ``#/components/schemas/X`` will *also* resolve to the + ``definitions`` body when both namespaces define ``X``. Cross-namespace + disambiguation would require teaching ``unpack_defs`` to key by full ref + path, which is out of scope here. + """ + config = AnthropicConfig() + tool = { + "type": "function", + "function": { + "name": "collision_tool", + "description": "", + "parameters": { + "type": "object", + "properties": { + "from_definitions": {"$ref": "#/definitions/Thing"}, + "from_components": {"$ref": "#/components/schemas/Thing"}, + }, + "definitions": { + "Thing": {"type": "string", "description": "from-definitions"}, + }, + "components": { + "schemas": { + "Thing": {"type": "integer", "description": "from-components"}, + } + }, + }, + }, + } + + transformed, _ = config._map_tool_helper(tool) + + assert transformed is not None + expected = {"type": "string", "description": "from-definitions"} + # Direct ref resolves to the `definitions` body (the documented winner). + assert transformed["input_schema"]["properties"]["from_definitions"] == expected + # Cross-namespace ref *also* resolves to the `definitions` body because + # ``unpack_defs`` keys by last path segment -- documented limitation. + assert transformed["input_schema"]["properties"]["from_components"] == expected diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index a29365544d..2061522fef 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -329,3 +329,170 @@ def test_transform_messages_helper_strips_thinking_blocks(): ) assert "thinking_blocks" not in out[1] assert out[1]["content"] == "I can help." + + +# ----------------------------------------------------------------------------- +# Regression tests for legacy / OpenAPI $ref defs in tool parameters. +# +# Fireworks (like Anthropic) only resolves `$defs` (JSON Schema 2020-12). Tools +# coming from MCP servers (legacy `definitions`) or OpenAPI-derived gateways +# such as AWS AgentCore (`components.schemas`) used to leave dangling `$ref` +# pointers, causing upstream "Error resolving schema reference" failures. See +# https://github.com/BerriAI/litellm/issues/26692. +# ----------------------------------------------------------------------------- + + +def _assert_no_unresolved_refs(parameters: dict) -> None: + blob = json.dumps(parameters) + assert "$ref" not in blob, f"unresolved $ref in transformed parameters: {blob}" + + +def test_transform_tools_inlines_components_schemas_refs(): + """OpenAPI `components.schemas` $refs (AgentCore-style) must be inlined.""" + config = FireworksAIConfig() + tools = [ + { + "type": "function", + "function": { + "name": "slides_presentations_create", + "description": "Create a Google Slides presentation", + "parameters": { + "type": "object", + "properties": { + "body": {"$ref": "#/components/schemas/Presentation"}, + }, + "required": ["body"], + "components": { + "schemas": { + "Presentation": { + "type": "object", + "properties": { + "title": {"type": "string"}, + "presentationId": {"type": "string"}, + }, + } + } + }, + }, + }, + } + ] + + out = config._transform_tools(tools) + + params = out[0]["function"]["parameters"] + _assert_no_unresolved_refs(params) + assert params["properties"]["body"] == { + "type": "object", + "properties": { + "title": {"type": "string"}, + "presentationId": {"type": "string"}, + }, + } + assert "components" not in params + + +def test_transform_tools_inlines_legacy_definitions_refs(): + """Legacy draft-04 `definitions` $refs must be inlined.""" + config = FireworksAIConfig() + tools = [ + { + "type": "function", + "function": { + "name": "create_thing", + "description": "Create a thing", + "parameters": { + "type": "object", + "properties": {"thing": {"$ref": "#/definitions/Thing"}}, + "definitions": { + "Thing": { + "type": "object", + "properties": {"id": {"type": "string"}}, + } + }, + }, + }, + } + ] + + out = config._transform_tools(tools) + + params = out[0]["function"]["parameters"] + _assert_no_unresolved_refs(params) + assert params["properties"]["thing"] == { + "type": "object", + "properties": {"id": {"type": "string"}}, + } + assert "definitions" not in params + + +def test_transform_tools_preserves_native_dollar_defs(): + """`$defs` is JSON Schema 2020-12 native; Fireworks resolves it itself.""" + config = FireworksAIConfig() + tools = [ + { + "type": "function", + "function": { + "name": "native_defs_tool", + "description": "", + "parameters": { + "type": "object", + "properties": {"a": {"$ref": "#/$defs/A"}}, + "$defs": {"A": {"type": "string"}}, + }, + }, + } + ] + + out = config._transform_tools(tools) + + params = out[0]["function"]["parameters"] + assert params["$defs"] == {"A": {"type": "string"}} + assert params["properties"]["a"] == {"$ref": "#/$defs/A"} + + +def test_transform_tools_skips_non_function_tools(): + """Non-``function`` tools (e.g. provider-native tool types) must pass + through ``_transform_tools`` untouched -- no ``strict`` pop, no $ref + inlining, no error. + """ + config = FireworksAIConfig() + non_function_tool = { + "type": "code_interpreter", + "code_interpreter": {"some": "config"}, + } + function_tool = { + "type": "function", + "function": { + "name": "create_thing", + "description": "Create a thing", + "parameters": { + "type": "object", + "properties": {"thing": {"$ref": "#/definitions/Thing"}}, + "definitions": { + "Thing": { + "type": "object", + "properties": {"id": {"type": "string"}}, + } + }, + }, + "strict": True, + }, + } + + out = config._transform_tools([non_function_tool, function_tool]) + + # Non-function tool is preserved verbatim. + assert out[0] == { + "type": "code_interpreter", + "code_interpreter": {"some": "config"}, + } + # Function tool still goes through both transformations: `strict` popped + # and the legacy $ref inlined. + assert "strict" not in out[1]["function"] + inlined = out[1]["function"]["parameters"] + assert "definitions" not in inlined + assert inlined["properties"]["thing"] == { + "type": "object", + "properties": {"id": {"type": "string"}}, + }