mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-23 00:25:00 +00:00
[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:
@@ -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)
|
||||
|
||||
@@ -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}")
|
||||
|
||||
|
||||
+99
-4
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user