[Feat] VertexAI - Support googlemap grounding in vertex ai (#15179)

* add VertexToolName

* test_vertex_tool_params

* fix: working maps grounding

* test_gemini_google_maps_tool_simple

* test_vertex_ai_map_google_maps_tool_with_location

* fix  # noqa: PLR0915

* _extract_google_maps_retrieval_config

* fixes for linting

* docs: **Google Maps**
This commit is contained in:
Ishaan Jaff
2025-10-03 16:07:51 -07:00
committed by GitHub
parent 4415b195d1
commit 10d6d72ae3
7 changed files with 431 additions and 48 deletions
+157
View File
@@ -621,6 +621,163 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \
#### **Google Maps**
Use Google Maps to provide location-based context to your Gemini models.
[**Relevant Vertex AI Docs**](https://ai.google.dev/gemini-api/docs/grounding#google-maps)
<Tabs>
<TabItem value="sdk" label="SDK">
**Basic Usage - Enable Widget Only**
```python showLineNumbers
from litellm import completion
## SETUP ENVIRONMENT
# !gcloud auth application-default login - run this to add vertex credentials to your env
tools = [{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}] # 👈 ADD GOOGLE MAPS
resp = litellm.completion(
model="vertex_ai/gemini-2.0-flash",
messages=[{"role": "user", "content": "What restaurants are nearby?"}],
tools=tools,
)
print(resp)
```
**With Location Data**
You can specify a location to ground the model's responses with location-specific information:
```python showLineNumbers
from litellm import completion
## SETUP ENVIRONMENT
# !gcloud auth application-default login - run this to add vertex credentials to your env
tools = [{
"googleMaps": {
"enableWidget": "ENABLE_WIDGET",
"latitude": 37.7749, # San Francisco latitude
"longitude": -122.4194, # San Francisco longitude
"languageCode": "en_US" # Optional: language for results
}
}] # 👈 ADD GOOGLE MAPS WITH LOCATION
resp = litellm.completion(
model="vertex_ai/gemini-2.0-flash",
messages=[{"role": "user", "content": "What restaurants are nearby?"}],
tools=tools,
)
print(resp)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
<Tabs>
<TabItem value="openai" label="OpenAI Python SDK">
**Basic Usage - Enable Widget Only**
```python showLineNumbers
from openai import OpenAI
client = OpenAI(
api_key="sk-1234", # pass litellm proxy key, if you're using virtual keys
base_url="http://0.0.0.0:4000/v1/" # point to litellm proxy
)
response = client.chat.completions.create(
model="gemini-2.0-flash",
messages=[{"role": "user", "content": "What restaurants are nearby?"}],
tools=[{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}],
)
print(response)
```
**With Location Data**
```python showLineNumbers
from openai import OpenAI
client = OpenAI(
api_key="sk-1234", # pass litellm proxy key, if you're using virtual keys
base_url="http://0.0.0.0:4000/v1/" # point to litellm proxy
)
response = client.chat.completions.create(
model="gemini-2.0-flash",
messages=[{"role": "user", "content": "What restaurants are nearby?"}],
tools=[{
"googleMaps": {
"enableWidget": "ENABLE_WIDGET",
"latitude": 37.7749, # San Francisco latitude
"longitude": -122.4194, # San Francisco longitude
"languageCode": "en_US" # Optional: language for results
}
}],
)
print(response)
```
</TabItem>
<TabItem value="curl" label="cURL">
**Basic Usage - Enable Widget Only**
```bash showLineNumbers
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-1234" \
-d '{
"model": "gemini-2.0-flash",
"messages": [
{"role": "user", "content": "What restaurants are nearby?"}
],
"tools": [
{
"googleMaps": {"enableWidget": "ENABLE_WIDGET"}
}
]
}'
```
**With Location Data**
```bash showLineNumbers
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-1234" \
-d '{
"model": "gemini-2.0-flash",
"messages": [
{"role": "user", "content": "What restaurants are nearby?"}
],
"tools": [
{
"googleMaps": {
"enableWidget": "ENABLE_WIDGET",
"latitude": 37.7749,
"longitude": -122.4194,
"languageCode": "en_US"
}
}
]
}'
```
</TabItem>
</Tabs>
</TabItem>
</Tabs>
#### **Moving from Vertex AI SDK to LiteLLM (GROUNDING)**
@@ -3,10 +3,10 @@ This file contains the transformation logic for the Gemini realtime API.
"""
import json
from litellm._uuid import uuid
from typing import Any, Dict, List, Optional, Union, cast
from litellm import verbose_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
@@ -186,9 +186,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
)
vertex_gemini_config = VertexGeminiConfig()
vertex_gemini_config._map_function(value)
optional_params["generationConfig"]["tools"] = (
vertex_gemini_config._map_function(value)
vertex_gemini_config._map_function(
value=value, optional_params=optional_params
)
)
elif key == "input_audio_transcription" and value is not None:
optional_params["inputAudioTranscription"] = {}
@@ -68,6 +68,7 @@ from litellm.types.llms.vertex_ai import (
ToolConfig,
Tools,
UsageMetadata,
VertexToolName,
)
from litellm.types.utils import (
ChatCompletionAudioResponse,
@@ -276,42 +277,106 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"""
return Tools(googleSearch={})
def _map_function(self, value: List[dict]) -> List[Tools]: # noqa: PLR0915
def _extract_google_maps_retrieval_config(
self, google_maps_config: dict
) -> Tuple[dict, Optional[dict]]:
"""
Extract location configuration from googleMaps tool for Vertex AI toolConfig.
Supports two interface styles:
1. Nested (recommended): {"enableWidget": "...", "retrievalConfig": {"latitude": ..., "longitude": ...}}
2. Flat (backward compat): {"enableWidget": "...", "latitude": ..., "longitude": ...}
Args:
google_maps_config: The googleMaps tool configuration from LiteLLM
Returns:
Tuple of (cleaned_google_maps_config, retrieval_config):
- cleaned_google_maps_config: googleMaps config without location fields
- retrieval_config: Location config for toolConfig.retrievalConfig or None
"""
retrieval_config = None
latitude = google_maps_config.get("latitude")
longitude = google_maps_config.get("longitude")
language_code = google_maps_config.get("languageCode")
if latitude is not None and longitude is not None:
retrieval_config = {
"latLng": {
"latitude": latitude,
"longitude": longitude,
}
}
if language_code is not None:
retrieval_config["languageCode"] = language_code
# Remove location fields from tool definition
cleaned_config = {
k: v
for k, v in google_maps_config.items()
if k not in ["latitude", "longitude", "languageCode"]
}
return cleaned_config, retrieval_config
def get_tool_value(
self,
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
def _map_function( # noqa: PLR0915
self, value: List[dict], optional_params: dict
) -> List[Tools]:
"""
Map OpenAI-style tools/functions to Vertex AI format.
Args:
value: List of tool definitions
optional_params: Request-scoped parameters to store retrieval config
Returns:
List of mapped tools in Vertex AI format
Side effects:
May add 'toolConfig' with 'retrievalConfig' to optional_params if
googleMaps tools contain location data
"""
gtool_func_declarations = []
googleSearch: Optional[dict] = None
googleSearchRetrieval: Optional[dict] = None
enterpriseWebSearch: Optional[dict] = None
urlContext: Optional[dict] = None
code_execution: Optional[dict] = None
googleMaps: Optional[dict] = None
google_maps_retrieval_config: Optional[dict] = None
# remove 'additionalProperties' from tools
value = _remove_additional_properties(value)
# 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
@@ -341,17 +406,27 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
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"
tool_name == "codeExecution" or tool_name == VertexToolName.CODE_EXECUTION.value
): # 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 tool_name and tool_name == "urlContext":
urlContext = get_tool_value(tool, "urlContext")
code_execution = self.get_tool_value(tool, "codeExecution")
elif tool_name and tool_name == VertexToolName.GOOGLE_SEARCH.value:
googleSearch = self.get_tool_value(tool, VertexToolName.GOOGLE_SEARCH.value)
elif tool_name and tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value:
googleSearchRetrieval = self.get_tool_value(tool, VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value)
elif tool_name and tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value:
enterpriseWebSearch = self.get_tool_value(tool, VertexToolName.ENTERPRISE_WEB_SEARCH.value)
elif tool_name and tool_name == VertexToolName.URL_CONTEXT.value:
urlContext = self.get_tool_value(tool, VertexToolName.URL_CONTEXT.value)
elif tool_name and (
tool_name == VertexToolName.GOOGLE_MAPS.value or tool_name == "google_maps"
):
google_maps_value = self.get_tool_value(tool, VertexToolName.GOOGLE_MAPS.value)
# Extract and transform location configuration for toolConfig
if google_maps_value is not None:
googleMaps, google_maps_retrieval_config = self._extract_google_maps_retrieval_config(
google_maps_config=google_maps_value
)
elif openai_function_object is not None:
gtool_func_declaration = FunctionDeclaration(
name=openai_function_object["name"],
@@ -377,15 +452,24 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
function_declarations=gtool_func_declarations,
)
if googleSearch is not None:
_tools["googleSearch"] = googleSearch
_tools[VertexToolName.GOOGLE_SEARCH.value] = googleSearch
if googleSearchRetrieval is not None:
_tools["googleSearchRetrieval"] = googleSearchRetrieval
_tools[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = googleSearchRetrieval
if enterpriseWebSearch is not None:
_tools["enterpriseWebSearch"] = enterpriseWebSearch
_tools[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = enterpriseWebSearch
if code_execution is not None:
_tools["code_execution"] = code_execution
_tools[VertexToolName.CODE_EXECUTION.value] = code_execution
if urlContext is not None:
_tools["url_context"] = urlContext
_tools[VertexToolName.URL_CONTEXT.value] = urlContext
if googleMaps is not None:
_tools[VertexToolName.GOOGLE_MAPS.value] = googleMaps
# Add retrieval config to toolConfig if googleMaps has location data
if google_maps_retrieval_config is not None:
if "toolConfig" not in optional_params:
optional_params["toolConfig"] = {}
optional_params["toolConfig"]["retrievalConfig"] = google_maps_retrieval_config
return [_tools]
def _map_response_schema(self, value: dict) -> dict:
@@ -606,8 +690,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
and isinstance(value, list)
and value
):
# Pass optional_params so _map_function can add toolConfig if needed
mapped_tools = self._map_function(
value=value, optional_params=optional_params
)
optional_params = self._add_tools_to_optional_params(
optional_params, self._map_function(value=value)
optional_params, mapped_tools
)
elif param == "tool_choice" and (
isinstance(value, str) or isinstance(value, dict)
+11
View File
@@ -209,6 +209,16 @@ class GenerationConfig(TypedDict, total=False):
thinkingConfig: GeminiThinkingConfig
class VertexToolName(str, Enum):
"""Enum for Vertex AI tool field names."""
GOOGLE_SEARCH = "googleSearch"
GOOGLE_SEARCH_RETRIEVAL = "googleSearchRetrieval"
ENTERPRISE_WEB_SEARCH = "enterpriseWebSearch"
URL_CONTEXT = "url_context"
CODE_EXECUTION = "code_execution"
GOOGLE_MAPS = "googleMaps"
class Tools(TypedDict, total=False):
function_declarations: List[FunctionDeclaration]
googleSearch: dict
@@ -216,6 +226,7 @@ class Tools(TypedDict, total=False):
enterpriseWebSearch: dict
url_context: dict
code_execution: dict
googleMaps: dict
retrieval: Retrieval
@@ -3846,3 +3846,33 @@ def test_gemini_grounding_on_streaming():
vertex_ai_grounding_metadata_shows_up = True
print(chunk)
assert vertex_ai_grounding_metadata_shows_up
def test_gemini_google_maps_tool_simple():
"""
Test googleMaps tool with just enableWidget parameter.
"""
load_vertex_ai_credentials()
litellm._turn_on_debug()
tools = [{"googleMaps": {"enableWidget": True}}]
tools_with_location = [{"googleMaps": {"enableWidget": True, "latitude": 37.7749, "longitude": -122.4194, "languageCode": "en_US"}}]
try:
for tools in [tools, tools_with_location]:
response = completion(
model="vertex_ai/gemini-2.0-flash",
messages=[
{
"role": "user",
"content": "What restaurants are nearby?",
}
],
tools=tools,
)
print(f"Response: {response.model_dump_json(indent=4)}")
assert response.choices[0].message.content is not None
except litellm.RateLimitError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@@ -444,12 +444,14 @@ def test_vertex_ai_map_thinking_param_with_budget_tokens_0():
def test_vertex_ai_map_tools():
v = VertexGeminiConfig()
tools = v._map_function(value=[{"code_execution": {}}])
optional_params = {}
tools = v._map_function(value=[{"code_execution": {}}], optional_params=optional_params)
assert len(tools) == 1
assert tools[0]["code_execution"] == {}
print(tools)
new_tools = v._map_function(value=[{"codeExecution": {}}])
new_optional_params = {}
new_tools = v._map_function(value=[{"codeExecution": {}}], optional_params=new_optional_params)
assert len(new_tools) == 1
print("new_tools", new_tools)
assert new_tools[0]["code_execution"] == {}
@@ -465,6 +467,7 @@ def test_vertex_ai_map_tool_with_anyof():
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()
optional_params = {}
value = [
{
"type": "function",
@@ -488,7 +491,7 @@ def test_vertex_ai_map_tool_with_anyof():
},
}
]
tools = v._map_function(value=value)
tools = v._map_function(value=value, optional_params=optional_params)
assert tools[0]["function_declarations"][0]["parameters"]["properties"][
"base_branch"
@@ -496,6 +499,7 @@ def test_vertex_ai_map_tool_with_anyof():
"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']}"
new_optional_params = {}
new_value = [
{
"type": "function",
@@ -518,7 +522,7 @@ def test_vertex_ai_map_tool_with_anyof():
},
}
]
new_tools = v._map_function(value=new_value)
new_tools = v._map_function(value=new_value, optional_params=new_optional_params)
assert new_tools[0]["function_declarations"][0]["parameters"]["properties"][
"base_branch"
@@ -1056,3 +1060,94 @@ def test_vertex_ai_code_line_length():
# Verify it contains the expected UUID format
assert 'uuid.uuid4().hex[:28]' in id_line, f"Line should contain shortened UUID format: {id_line}"
def test_vertex_ai_map_google_maps_tool_simple():
"""
Test googleMaps tool transformation without location data.
Input:
value=[{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}]
optional_params={}
Expected Output:
tools=[{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}]
optional_params={} (unchanged)
"""
v = VertexGeminiConfig()
optional_params = {}
tools = v._map_function(
value=[{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}],
optional_params=optional_params
)
assert len(tools) == 1
assert "googleMaps" in tools[0]
assert tools[0]["googleMaps"]["enableWidget"] == "ENABLE_WIDGET"
assert "toolConfig" not in optional_params
def test_vertex_ai_map_google_maps_tool_with_location():
"""
Test googleMaps tool transformation with location data.
Verifies latitude/longitude/languageCode are extracted to toolConfig.retrievalConfig.
Input:
value=[{
"googleMaps": {
"enableWidget": "ENABLE_WIDGET",
"latitude": 37.7749,
"longitude": -122.4194,
"languageCode": "en_US"
}
}]
optional_params={}
Expected Output:
tools=[{
"googleMaps": {"enableWidget": "ENABLE_WIDGET"}
}]
optional_params={
"toolConfig": {
"retrievalConfig": {
"latLng": {
"latitude": 37.7749,
"longitude": -122.4194
},
"languageCode": "en_US"
}
}
}
"""
v = VertexGeminiConfig()
optional_params = {}
tools = v._map_function(
value=[{
"googleMaps": {
"enableWidget": "ENABLE_WIDGET",
"latitude": 37.7749,
"longitude": -122.4194,
"languageCode": "en_US"
}
}],
optional_params=optional_params
)
assert len(tools) == 1
assert "googleMaps" in tools[0]
google_maps_tool = tools[0]["googleMaps"]
assert google_maps_tool["enableWidget"] == "ENABLE_WIDGET"
assert "latitude" not in google_maps_tool
assert "longitude" not in google_maps_tool
assert "languageCode" not in google_maps_tool
assert "toolConfig" in optional_params
assert "retrievalConfig" in optional_params["toolConfig"]
retrieval_config = optional_params["toolConfig"]["retrievalConfig"]
assert retrieval_config["latLng"]["latitude"] == 37.7749
assert retrieval_config["latLng"]["longitude"] == -122.4194
assert retrieval_config["languageCode"] == "en_US"
@@ -145,6 +145,7 @@ def test_build_vertex_schema():
([{"googleSearchRetrieval": {}}], "googleSearchRetrieval"),
([{"enterpriseWebSearch": {}}], "enterpriseWebSearch"),
([{"code_execution": {}}], "code_execution"),
([{"googleMaps": {}}], "googleMaps"),
],
)
def test_vertex_tool_params(tools, key):