scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

VertexAI - codeExecution tool support + anyOf handling (#11195)

Browse source 
* 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
This commit is contained in:
Krish Dholakia 2025-05-27 21:23:14 -07:00 • committed by GitHub
parent 40a4e97d8e
commit 7072466775
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 215 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
docs/my-website/docs/providers/vertex.md
View file

@ -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 \
</TabItem>
</Tabs>
#### **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).
<Tabs>
@ -491,6 +495,53 @@ curl http://localhost:4000/v1/chat/completions \
</TabItem>
</Tabs>
#### **Code Execution**
<Tabs>
<TabItem value="sdk" label="SDK">
```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)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
```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": {}}]
}
'
```
</TabItem>
</Tabs>
#### **Moving from Vertex AI SDK to LiteLLM (GROUNDING)**

30
litellm/llms/vertex_ai/common_utils.py
View file

@ -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

52
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -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=(

39
tests/local_testing/test_amazing_vertex_completion.py
View file

@ -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
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)

55
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -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']}"