ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix vertex ai function calls

Browse Source
This commit is contained in:
Sameer Kankute 2026-03-02 19:42:18 +05:30
parent 7e2f2a8ffa
commit e40e913622
1 changed files with 86 additions and 63 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

149
tests/local_testing/test_amazing_vertex_completion.py
View File

@ -2881,74 +2881,96 @@ def test_gemini_function_call_parameter_in_messages():
client = HTTPHandler(concurrent_limit=1)
with patch.object(client, "post", new=MagicMock()) as mock_client:
try:
response_stream = completion(
model="vertex_ai/gemini-1.5-pro",
messages=messages,
tools=tools,
tool_choice="auto",
client=client,
)
except Exception as e:
print(e)
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {}
mock_response.json.return_value = {
"candidates": [
{
"content": {"parts": [{"text": "test"}], "role": "model"},
"finishReason": "STOP",
}
],
"usageMetadata": {
"promptTokenCount": 0,
"candidatesTokenCount": 0,
"totalTokenCount": 0,
},
}
# mock_client.assert_any_call()
with patch(
"litellm.llms.vertex_ai.vertex_llm_base.VertexLLMBase._ensure_access_token",
return_value=({"Authorization": "Bearer fake"}, "test-project"),
):
with patch.object(client, "post", new=MagicMock()) as mock_client:
mock_client.return_value = mock_response
try:
completion(
model="vertex_ai/gemini-1.5-pro",
messages=messages,
tools=tools,
tool_choice="auto",
client=client,
)
except Exception as e:
print(e)
assert {
"contents": [
{
"role": "user",
"parts": [{"text": "search for weather in boston (use `search`)"}],
},
{
"role": "model",
"parts": [
{
"function_call": {
"name": "search",
"args": {"queries": ["weather in boston"]},
assert mock_client.called
assert {
"contents": [
{
"role": "user",
"parts": [{"text": "search for weather in boston (use `search`)"}],
},
{
"role": "model",
"parts": [
{
"function_call": {
"name": "search",
"args": {"queries": ["weather in boston"]},
}
}
}
],
},
{
"parts": [
{
"function_response": {
],
},
{
"role": "user",
"parts": [
{
"function_response": {
"name": "search",
"response": {
"content": "The current weather in Boston is 22°F."
},
}
]
],
},
],
"system_instruction": {"parts": [{"text": "Use search for most queries."}]},
"tools": [
{
"function_declarations": [
{
"name": "search",
"response": {
"content": "The current weather in Boston is 22°F."
"description": "Executes searches.",
"parameters": {
"type": "object",
"properties": {
"queries": {
"type": "array",
"description": "A list of queries to search for.",
"items": {"type": "string"},
}
},
"required": ["queries"],
},
}
}
]
},
],
"system_instruction": {"parts": [{"text": "Use search for most queries."}]},
"tools": [
{
"function_declarations": [
{
"name": "search",
"description": "Executes searches.",
"parameters": {
"type": "object",
"properties": {
"queries": {
"type": "array",
"description": "A list of queries to search for.",
"items": {"type": "string"},
}
},
"required": ["queries"],
},
}
]
}
],
"toolConfig": {"functionCallingConfig": {"mode": "AUTO"}},
} == mock_client.call_args.kwargs["json"]
]
}
],
"toolConfig": {"functionCallingConfig": {"mode": "AUTO"}},
} == mock_client.call_args.kwargs["json"]
def test_gemini_function_call_parameter_in_messages_2():
@ -2995,6 +3017,7 @@ def test_gemini_function_call_parameter_in_messages_2():
],
},
{
"role": "user",
"parts": [
{
"function_response": {
@ -3004,7 +3027,7 @@ def test_gemini_function_call_parameter_in_messages_2():
},
}
}
]
],
},
]