From e40e9136221169257ad4a20ccc2775f58df0a9bf Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 19:42:18 +0530 Subject: [PATCH] Fix vertex ai function calls --- .../test_amazing_vertex_completion.py | 149 ++++++++++-------- 1 file changed, 86 insertions(+), 63 deletions(-) diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 998f2beb4a1..8974275ca83 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -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(): }, } } - ] + ], }, ]