diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 51e104d0a94..f298d6f4a12 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -699,7 +699,7 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False): def _convert_consts_to_enums( - schema: dict[str, object], # mutable-ok: response schema is normalized in place + schema: dict[str, object], # mutable-ok: Gemini schema is normalized in place depth: int = 0, ) -> None: """ diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 69fe5678de9..5d3db82f94b 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -86,6 +86,7 @@ from ..common_utils import ( VertexAIError, _build_json_schema, _build_vertex_schema, + _convert_consts_to_enums, supports_response_json_schema, ) from ..vertex_llm_base import VertexBase @@ -594,16 +595,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): for tool in value: openai_function_object: ChatCompletionToolParamFunctionChunk | None = None if "function" in tool: # tools list - _openai_function_object = ChatCompletionToolParamFunctionChunk(**tool["function"]) - - if ( - "parameters" in _openai_function_object - and _openai_function_object["parameters"] is not None - and isinstance(_openai_function_object["parameters"], dict) - ): # OPENAI accepts JSON Schema, Google accepts OpenAPI schema. - _openai_function_object["parameters"] = _build_vertex_schema(_openai_function_object["parameters"]) - - openai_function_object = _openai_function_object + openai_function_object = ChatCompletionToolParamFunctionChunk(**tool["function"]) elif "name" in tool: # functions list openai_function_object = ChatCompletionToolParamFunctionChunk(**tool) @@ -663,13 +655,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) _description = openai_function_object.get("description", None) _parameters = openai_function_object.get("parameters", None) - if isinstance(_parameters, str) and len(_parameters) == 0: - _parameters = { - "type": "object", - } if _description is not None: gtool_func_declaration["description"] = _description - if _parameters is not None: + if isinstance(_parameters, dict): + _convert_consts_to_enums(_parameters) + gtool_func_declaration["parameters"] = _build_vertex_schema(_parameters) + elif isinstance(_parameters, str) and len(_parameters) == 0: + gtool_func_declaration["parameters"] = {"type": "object"} + elif _parameters is not None: gtool_func_declaration["parameters"] = _parameters gtool_func_declarations.append(gtool_func_declaration) else: diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index d2788408e09..359e6de539a 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2,7 +2,7 @@ import asyncio import json import re from copy import deepcopy -from typing import Final, List, cast +from typing import Final, List, Literal, cast from unittest.mock import MagicMock, patch import pytest @@ -1219,6 +1219,34 @@ def test_vertex_ai_map_tools(): assert tools == new_tools +@pytest.mark.parametrize("legacy_functions", [False, True]) +def test_gemini_map_tool_converts_pydantic_consts_to_enums(legacy_functions: bool): + class Operation(BaseModel): + kind: Literal["create"] + resource_name: str + + class ToolInput(BaseModel): + status: Literal["pending"] + operation: Operation + + function = { + "name": "perform_action", + "description": "Perform an action", + "parameters": ToolInput.model_json_schema(), + } + tools_input = [function] if legacy_functions else [{"type": "function", "function": function}] + + tools = VertexGeminiConfig()._map_function(value=tools_input, optional_params={}) + parameters = tools[0]["function_declarations"][0]["parameters"] + + assert parameters["properties"]["status"]["enum"] == ["pending"] + assert parameters["properties"]["operation"]["properties"]["kind"]["enum"] == ["create"] + assert parameters["properties"]["operation"]["properties"]["resource_name"]["type"] == "string" + assert parameters["properties"]["operation"]["required"] == ["kind", "resource_name"] + assert "const" not in json.dumps(parameters) + assert "$ref" not in json.dumps(parameters) + + def test_vertex_ai_map_tool_with_anyof(): """ Related issue: https://github.com/BerriAI/litellm/issues/11164