diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 14aebcaabaf..51e104d0a94 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -698,6 +698,44 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False): return parameters +def _convert_consts_to_enums( + schema: dict[str, object], # mutable-ok: response schema is normalized in place + depth: int = 0, +) -> None: + """ + Converts 'const' to 'enum' only in schema fields (intentionally skips examples, etc) + """ + if depth > DEFAULT_MAX_RECURSE_DEPTH: + return + + if "const" in schema: + const_value: Final = schema.pop("const") # rebind-ok: removes the unsupported keyword + enum_values: Final = [const_value] # mutable-ok: JSON Schema enum requires an array + schema["enum"] = enum_values # rebind-ok: replaces const with its singleton enum + + for schema_map in (schema.get("$defs"), schema.get("properties")): + if not isinstance(schema_map, dict): + continue + for mapped_schema in schema_map.values(): + if isinstance(mapped_schema, dict): + _convert_consts_to_enums(mapped_schema, depth + 1) + + for direct_schema in (schema.get("additionalProperties"), schema.get("items")): + if isinstance(direct_schema, dict): + _convert_consts_to_enums(direct_schema, depth + 1) + + for schema_list in ( + schema.get("prefixItems"), + schema.get("anyOf"), + schema.get("oneOf"), + ): + if not isinstance(schema_list, list): + continue + for listed_schema in schema_list: + if isinstance(listed_schema, dict): + _convert_consts_to_enums(listed_schema, depth + 1) + + def _build_json_schema(parameters: dict) -> dict: """ Build a JSON Schema for use with Gemini's responseJsonSchema parameter. @@ -707,6 +745,7 @@ def _build_json_schema(parameters: dict) -> dict: - Does NOT add propertyOrdering - Does NOT filter fields (allows additionalProperties) - Preserves $defs/$ref (Gemini 2.0+ supports JSON Schema references natively) + - Converts const values to equivalent single-value enums Parameters: parameters: dict - the JSON schema to process @@ -714,13 +753,7 @@ def _build_json_schema(parameters: dict) -> dict: Returns: dict - the processed schema in standard JSON Schema format """ - # Gemini 2.0+ with responseJsonSchema accepts standard JSON Schema as-is, - # including $ref, $defs, anyOf, etc. No transformations needed — the - # OpenAPI-specific fixes (unpack_defs, add_object_type, convert_anyof, etc.) - # are only required for responseSchema (Gemini 1.5) and can break valid - # JSON Schema by adding conflicting fields to $ref nodes. - # See: https://blog.google/technology/developers/gemini-api-structured-outputs/ - + _convert_consts_to_enums(parameters) return parameters diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index e9f87ba6cae..9998ec0c8e2 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -27,6 +27,7 @@ IGNORE_FUNCTIONS = [ "_extract_fields_recursive", # max depth set. "_remove_json_schema_refs", # max depth set., "_convert_schema_types", # max depth set., + "_convert_consts_to_enums", # max depth set. "_fix_enum_empty_strings", # max depth set., "get_access_token", # max depth set., "_redact_base64", # max depth set. diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index c206fcec420..a99ffa29a1f 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1,6 +1,8 @@ +from typing import Literal, Annotated, Union from unittest.mock import patch import pytest +from pydantic import BaseModel, Field from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -298,6 +300,149 @@ def test_build_vertex_schema(): assert _build_vertex_schema(parameters) == expected_output +def test_build_json_schema_converts_nested_consts_to_enums_and_preserves_refs(): + from litellm.llms.vertex_ai.common_utils import _build_json_schema + + parameters = { + "$defs": { + "CreateAction": { + "properties": { + "kind": {"const": "create", "type": "string"}, + "mode": {"enum": ["fast", "safe"], "type": "string"}, + }, + "required": ["kind"], + "type": "object", + }, + "DeleteAction": { + "properties": { + "kind": {"const": "delete", "type": "string"}, + }, + "required": ["kind"], + "type": "object", + }, + }, + "properties": { + "actions": { + "items": { + "anyOf": [ + {"$ref": "#/$defs/CreateAction"}, + {"$ref": "#/$defs/DeleteAction"}, + ] + }, + "type": "array", + }, + "status": {"const": "pending", "type": "string"}, + }, + "type": "object", + } + + result = _build_json_schema(parameters) + + assert result == { + "$defs": { + "CreateAction": { + "properties": { + "kind": {"enum": ["create"], "type": "string"}, + "mode": {"enum": ["fast", "safe"], "type": "string"}, + }, + "required": ["kind"], + "type": "object", + }, + "DeleteAction": { + "properties": { + "kind": {"enum": ["delete"], "type": "string"}, + }, + "required": ["kind"], + "type": "object", + }, + }, + "properties": { + "actions": { + "items": { + "anyOf": [ + {"$ref": "#/$defs/CreateAction"}, + {"$ref": "#/$defs/DeleteAction"}, + ] + }, + "type": "array", + }, + "status": {"enum": ["pending"], "type": "string"}, + }, + "type": "object", + } + + +def test_build_json_schema_converts_nested_consts_to_enums_and_preserves_refs_with_pydantic(): + from litellm.llms.vertex_ai.common_utils import _build_json_schema + + class CreateAction(BaseModel): + kind: Literal["create"] + resource_name: str + + class DeleteAction(BaseModel): + kind: Literal["delete"] + resource_id: str + + Action = Annotated[ + Union[CreateAction, DeleteAction], + Field(discriminator="kind"), + ] + + class ActionResponse(BaseModel): + status: Literal["pending"] + actions: list[Action] + + pydantic_schema = ActionResponse.model_json_schema() + + result = _build_json_schema(pydantic_schema) + + assert result == { + "$defs": { + "CreateAction": { + "properties": { + "kind": {"enum": ["create"], "title": "Kind", "type": "string"}, + "resource_name": {"title": "Resource Name", "type": "string"}, + }, + "required": ["kind", "resource_name"], + "title": "CreateAction", + "type": "object", + }, + "DeleteAction": { + "properties": { + "kind": {"enum": ["delete"], "title": "Kind", "type": "string"}, + "resource_id": {"title": "Resource Id", "type": "string"}, + }, + "required": ["kind", "resource_id"], + "title": "DeleteAction", + "type": "object", + } + }, + "properties": { + "status": {"enum": ["pending"], "title": "Status", "type": "string"}, + "actions": { + "items": { + "discriminator": { + "mapping": { + "create": "#/$defs/CreateAction", + "delete": "#/$defs/DeleteAction", + }, + "propertyName": "kind", + }, + "oneOf": [ + {"$ref": "#/$defs/CreateAction"}, + {"$ref": "#/$defs/DeleteAction"}, + ], + }, + "title": "Actions", + "type": "array", + }, + }, + "required": ["status", "actions"], + "title": "ActionResponse", + "type": "object", + } + + def test_process_items_with_excessive_nesting(): """Test process_items with excessive nesting > max levels +1 deep.""" # generate a schema with excessive nesting