From 10d6d72ae3b985bfeee48ced73b1be813aad26a7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 3 Oct 2025 16:07:51 -0700 Subject: [PATCH] [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** --- docs/my-website/docs/providers/vertex.md | 157 ++++++++++++++++ .../llms/gemini/realtime/transformation.py | 7 +- .../vertex_and_google_ai_studio_gemini.py | 170 +++++++++++++----- litellm/types/llms/vertex_ai.py | 11 ++ .../test_amazing_vertex_completion.py | 30 ++++ ...test_vertex_and_google_ai_studio_gemini.py | 103 ++++++++++- .../llms/vertex_ai/test_vertex.py | 1 + 7 files changed, 431 insertions(+), 48 deletions(-) diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 943823c638..2543c15d13 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -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) + + + + +**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) +``` + + + + + + + +**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) +``` + + + +**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" + } + } + ] + }' +``` + + + + + + #### **Moving from Vertex AI SDK to LiteLLM (GROUNDING)** diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index e1dd6f146f..62329358e4 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -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"] = {} 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 5871c35338..cc50bc9954 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 @@ -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) diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 854d37522c..f1f7ac2c66 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -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 diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 8262d43a0d..d533418117 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -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}") + 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 62a11bf676..5ae6cf3da2 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 @@ -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" diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py index 31d4fd1c19..39ed09f81b 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex.py @@ -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):