From 70724667754fdd7ff396393627b0097dd1f9dab2 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Tue, 27 May 2025 21:23:14 -0700 Subject: [PATCH] VertexAI - `codeExecution` tool support + anyOf handling (#11195) * fix(vertex_and_google_ai_studio_gemini.py): handle both camel case and underscores in the tool for vertex ai code execution support vertex ai code execution * docs(vertex.md): add code execution example to vertex ai * fix(vertex_ai/common_utils.py): when anyof in field, just select anyof - don't include other k,v pairs - vertex throws error Fixes https://github.com/BerriAI/litellm/issues/11164 * fix(common_utils.py): add title field inside anyof - to retain some description Addresses https://github.com/BerriAI/litellm/issues/11164#issuecomment-2914728385 --- docs/my-website/docs/providers/vertex.md | 53 +++++++++++++++++- litellm/llms/vertex_ai/common_utils.py | 30 +++++++++- .../vertex_and_google_ai_studio_gemini.py | 52 ++++++++++++++---- .../test_amazing_vertex_completion.py | 39 ++++++++++++- ...test_vertex_and_google_ai_studio_gemini.py | 55 +++++++++++++++++++ 5 files changed, 215 insertions(+), 14 deletions(-) diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index e8130741732..dae68f4fc25 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -347,7 +347,9 @@ Return a `list[Recipe]` completion(model="vertex_ai/gemini-1.5-flash-preview-0514", messages=messages, response_format={ "type": "json_object" }) ``` -### **Grounding - Web Search** +### **Google Hosted Tools (Web Search, Code Execution, etc.)** + +#### **Web Search** Add Google Search Result grounding to vertex ai calls. @@ -422,6 +424,8 @@ curl http://localhost:4000/v1/chat/completions \ +#### **Enterprise Web Search** + You can also use the `enterpriseWebSearch` tool for an [enterprise compliant search](https://cloud.google.com/vertex-ai/generative-ai/docs/grounding/web-grounding-enterprise). @@ -491,6 +495,53 @@ curl http://localhost:4000/v1/chat/completions \ +#### **Code Execution** + + + + + + +```python showLineNumbers +from litellm import completion +import os + +## SETUP ENVIRONMENT +# !gcloud auth application-default login - run this to add vertex credentials to your env + + +tools = [{"codeExecution": {}}] # 👈 ADD CODE EXECUTION + +response = completion( + model="vertex_ai/gemini-2.0-flash", + messages=[{"role": "user", "content": "What is the weather in San Francisco?"}], + tools=tools, +) + +print(response) +``` + + + + +```bash showLineNumbers +curl -X POST 'http://0.0.0.0:4000/chat/completions' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer sk-1234' \ +-d '{ + "model": "gemini-2.0-flash", + "messages": [{"role": "user", "content": "What is the weather in San Francisco?"}], + "tools": [{"codeExecution": {}}] +} +' +``` + + + + + + + #### **Moving from Vertex AI SDK to LiteLLM (GROUNDING)** diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index f96848c6d56..8324af73a0f 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -84,7 +84,7 @@ def _get_vertex_url( endpoint = "generateContent" if stream is True: endpoint = "streamGenerateContent" - if vertex_location== "global": + if vertex_location == "global": url = f"https://aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/global/publishers/google/models/{model}:{endpoint}?alt=sse" else: url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}?alt=sse" @@ -212,6 +212,33 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False): return parameters +def _filter_anyof_fields(schema_dict: Dict[str, Any]) -> Dict[str, Any]: + """ + When anyof is present, only keep the anyof field and its contents - otherwise VertexAI will throw an error - https://github.com/BerriAI/litellm/issues/11164 + Filter out other fields in the same dict. + + E.g. {"anyOf": [{"type": "string"}, {"type": "null"}], "default": "test"} -> {"anyOf": [{"type": "string"}, {"type": "null"}]} + + Case 2: If additional metadata is present, try to keep it + E.g. {"anyOf": [{"type": "string"}, {"type": "null"}], "default": "test", "title": "test"} -> {"anyOf": [{"type": "string", "title": "test"}, {"type": "null", "title": "test"}]} + """ + title = schema_dict.get("title", None) + + if isinstance(schema_dict, dict) and schema_dict.get("anyOf"): + any_of = schema_dict["anyOf"] + if ( + title + and isinstance(any_of, list) + and all(isinstance(item, dict) for item in any_of) + ): + for item in any_of: + item["title"] = title + return {"anyOf": any_of} + else: + return schema_dict + return schema_dict + + def process_items(schema, depth=0): if depth > DEFAULT_MAX_RECURSE_DEPTH: raise ValueError( @@ -277,6 +304,7 @@ def filter_schema_fields( return schema_dict result = {} + schema_dict = _filter_anyof_fields(schema_dict) for key, value in schema_dict.items(): if key not in valid_fields: continue 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 b347d775ed8..b52e3798a89 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 @@ -262,6 +262,30 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # remove 'strict' from tools value = _remove_strict_from_schema(value) + def get_tool_value(tool: dict, tool_name: str) -> Optional[dict]: + """ + Helper function to get tool value handling both camelCase and underscore_case variants + + Args: + tool (dict): The tool dictionary + tool_name (str): The base tool name (e.g. "codeExecution") + + Returns: + Optional[dict]: The tool value if found, None otherwise + """ + # Convert camelCase to underscore_case + underscore_name = "".join( + ["_" + c.lower() if c.isupper() else c for c in tool_name] + ).lstrip("_") + # Try both camelCase and underscore_case variants + + if tool.get(tool_name) is not None: + return tool.get(tool_name) + elif tool.get(underscore_name) is not None: + return tool.get(underscore_name) + else: + return None + for tool in value: openai_function_object: Optional[ ChatCompletionToolParamFunctionChunk @@ -284,15 +308,17 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif "name" in tool: # functions list openai_function_object = ChatCompletionToolParamFunctionChunk(**tool) # type: ignore - # check if grounding - if tool.get("googleSearch", None) is not None: - googleSearch = tool["googleSearch"] - elif tool.get("googleSearchRetrieval", None) is not None: - googleSearchRetrieval = tool["googleSearchRetrieval"] - elif tool.get("enterpriseWebSearch", None) is not None: - enterpriseWebSearch = tool["enterpriseWebSearch"] - elif tool.get("code_execution", None) is not None: - code_execution = tool["code_execution"] + tool_name = list(tool.keys())[0] if len(tool.keys()) == 1 else None + if tool_name and ( + tool_name == "codeExecution" or tool_name == "code_execution" + ): # code_execution maintained for backwards compatibility + code_execution = get_tool_value(tool, "codeExecution") + elif tool_name and tool_name == "googleSearch": + googleSearch = get_tool_value(tool, "googleSearch") + elif tool_name and tool_name == "googleSearchRetrieval": + googleSearchRetrieval = get_tool_value(tool, "googleSearchRetrieval") + elif tool_name and tool_name == "enterpriseWebSearch": + enterpriseWebSearch = get_tool_value(tool, "enterpriseWebSearch") elif openai_function_object is not None: gtool_func_declaration = FunctionDeclaration( name=openai_function_object["name"], @@ -466,8 +492,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): optional_params["tool_choice"] = _tool_choice_value elif param == "parallel_tool_calls": if value is False: - tools = non_default_params.get("tools", non_default_params.get("functions")) - num_function_declarations = len(tools) if isinstance(tools, list) else 0 + tools = non_default_params.get( + "tools", non_default_params.get("functions") + ) + num_function_declarations = ( + len(tools) if isinstance(tools, list) else 0 + ) if num_function_declarations > 1: raise litellm.utils.UnsupportedParamsError( message=( diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index e35ab0a1348..e36b57e5570 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -3743,4 +3743,41 @@ def test_vertex_ai_llama_tool_calling(): assert response.choices[0].message.tool_calls is not None assert response.choices[0].finish_reason == "tool_calls" - assert response._hidden_params["response_cost"] > 0 \ No newline at end of file + assert response._hidden_params["response_cost"] > 0 + + +def test_vertex_schema_test(): + load_vertex_ai_credentials() + litellm._turn_on_debug() + + def tool_call(text: str | None) -> str: + return text or "No text provided" + + + tool = { + "type": "function", + "function": { + "name": "git_create_branch", + "description": "Creates a new branch from an optional base branch", + "parameters": { + "type": "object", + "properties": { + "repo_path": {"title": "Repo Path", "type": "string"}, + "branch_name": {"title": "Branch Name", "type": "string"}, + "base_branch": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": None, "title": "Base Branch"}, + }, + "required": ["repo_path", "branch_name"], + "title": "GitCreateBranch", + } + } + } + + + response = litellm.completion( + model="vertex_ai/gemini-2.5-flash-preview-05-20", + messages=[{"role": "user", "content": "call the tool"}], + tools=[tool], + tool_choice="required", + ) + + print(response) 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 455d9fb9d12..24705bb6ea5 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 @@ -375,3 +375,58 @@ def test_vertex_ai_usage_metadata_response_token_count(): assert result.prompt_tokens_details.audio_tokens is None assert result.prompt_tokens_details.cached_tokens is None assert result.completion_tokens_details.text_tokens == 74 + + +def test_vertex_ai_map_tools(): + v = VertexGeminiConfig() + tools = v._map_function(value=[{"code_execution": {}}]) + assert len(tools) == 1 + assert tools[0]["code_execution"] == {} + print(tools) + + new_tools = v._map_function(value=[{"codeExecution": {}}]) + assert len(new_tools) == 1 + print("new_tools", new_tools) + assert new_tools[0]["code_execution"] == {} + print(new_tools) + + assert tools == new_tools + + +def test_vertex_ai_map_tool_with_anyof(): + """ + Related issue: https://github.com/BerriAI/litellm/issues/11164 + + Ensure if anyof is present, only the anyof field and its contents are kept - otherwise VertexAI will throw an error - https://github.com/BerriAI/litellm/issues/11164 + """ + v = VertexGeminiConfig() + value = [ + { + "type": "function", + "function": { + "name": "git_create_branch", + "description": "Creates a new branch from an optional base branch", + "parameters": { + "type": "object", + "properties": { + "repo_path": {"title": "Repo Path", "type": "string"}, + "branch_name": {"title": "Branch Name", "type": "string"}, + "base_branch": { + "anyOf": [{"type": "string"}, {"type": "null"}], + "default": None, + "title": "Base Branch", + }, + }, + "required": ["repo_path", "branch_name"], + "title": "GitCreateBranch", + }, + }, + } + ] + tools = v._map_function(value=value) + + assert tools[0]["function_declarations"][0]["parameters"]["properties"][ + "base_branch" + ] == { + "anyOf": [{"type": "string", "nullable": True, "title": "Base Branch"}] + }, f"Expected only anyOf field and its contents to be kept, but got {tools[0]['function_declarations'][0]['parameters']['properties']['base_branch']}"