NeoAnd/adk-python
1
0
Fork 0
mirror of https://github.com/EvolutionAPI/adk-python.git synced 2025-07-13 15:14:50 -06:00
Code Issues Actions Packages Projects Releases Wiki Activity

chore: autoformat the test_litellm.py

Browse Source 
PiperOrigin-RevId: 764937761
This commit is contained in:
Xiang (Sean) Zhou 2025-05-29 16:48:33 -07:00 committed by Copybara-Service
parent 2b41824e46
commit ae7d3a79c8
1 changed files with 62 additions and 66 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

128
tests/unittests/models/test_litellm.py
View File

@ -13,6 +13,7 @@
# limitations under the License. # limitations under the License.
import json
from unittest.mock import AsyncMock from unittest.mock import AsyncMock
from unittest.mock import Mock from unittest.mock import Mock
@ -38,7 +39,6 @@ from litellm.types.utils import Delta
from litellm.types.utils import ModelResponse from litellm.types.utils import ModelResponse
from litellm.types.utils import StreamingChoices from litellm.types.utils import StreamingChoices
import pytest import pytest
import json
LLM_REQUEST_WITH_FUNCTION_DECLARATION = LlmRequest( LLM_REQUEST_WITH_FUNCTION_DECLARATION = LlmRequest(
contents=[ contents=[
@ -1190,74 +1190,70 @@ async def test_generate_content_async_stream_with_usage_metadata(
async def test_generate_content_async_multiple_function_calls( async def test_generate_content_async_multiple_function_calls(
mock_completion, lite_llm_instance mock_completion, lite_llm_instance
): ):
"""Test handling of multiple function calls with different indices in streaming mode. """Test handling of multiple function calls with different indices in streaming mode.
This test verifies that: This test verifies that:
1. Multiple function calls with different indices are handled correctly 1. Multiple function calls with different indices are handled correctly
2. Arguments and names are properly accumulated for each function call 2. Arguments and names are properly accumulated for each function call
3. The final response contains all function calls with correct indices 3. The final response contains all function calls with correct indices
""" """
mock_completion.return_value = MULTIPLE_FUNCTION_CALLS_STREAM mock_completion.return_value = MULTIPLE_FUNCTION_CALLS_STREAM
llm_request = LlmRequest( llm_request = LlmRequest(
contents=[ contents=[
types.Content( types.Content(
role="user", role="user",
parts=[types.Part.from_text(text="Test multiple function calls")], parts=[types.Part.from_text(text="Test multiple function calls")],
) )
], ],
config=types.GenerateContentConfig( config=types.GenerateContentConfig(
tools=[ tools=[
types.Tool( types.Tool(
function_declarations=[ function_declarations=[
types.FunctionDeclaration( types.FunctionDeclaration(
name="function_1", name="function_1",
description="First test function", description="First test function",
parameters=types.Schema( parameters=types.Schema(
type=types.Type.OBJECT, type=types.Type.OBJECT,
properties={ properties={
"arg": types.Schema(type=types.Type.STRING), "arg": types.Schema(type=types.Type.STRING),
}, },
), ),
), ),
types.FunctionDeclaration( types.FunctionDeclaration(
name="function_2", name="function_2",
description="Second test function", description="Second test function",
parameters=types.Schema( parameters=types.Schema(
type=types.Type.OBJECT, type=types.Type.OBJECT,
properties={ properties={
"arg": types.Schema(type=types.Type.STRING), "arg": types.Schema(type=types.Type.STRING),
}, },
), ),
), ),
] ]
) )
], ],
), ),
) )
responses = [] responses = []
async for response in lite_llm_instance.generate_content_async( async for response in lite_llm_instance.generate_content_async(
llm_request, stream=True llm_request, stream=True
): ):
responses.append(response) responses.append(response)
# Verify we got the final response with both function calls # Verify we got the final response with both function calls
assert len(responses) > 0 assert len(responses) > 0
final_response = responses[-1] final_response = responses[-1]
assert final_response.content.role == "model" assert final_response.content.role == "model"
assert len(final_response.content.parts) == 2 assert len(final_response.content.parts) == 2
# Verify first function call # Verify first function call
assert final_response.content.parts[0].function_call.name == "function_1" assert final_response.content.parts[0].function_call.name == "function_1"
assert final_response.content.parts[0].function_call.id == "call_1" assert final_response.content.parts[0].function_call.id == "call_1"
assert final_response.content.parts[0].function_call.args == { assert final_response.content.parts[0].function_call.args == {"arg": "value1"}
"arg": "value1"
}
# Verify second function call # Verify second function call
assert final_response.content.parts[1].function_call.name == "function_2" assert final_response.content.parts[1].function_call.name == "function_2"
assert final_response.content.parts[1].function_call.id == "call_2" assert final_response.content.parts[1].function_call.id == "call_2"
assert final_response.content.parts[1].function_call.args == { assert final_response.content.parts[1].function_call.args == {"arg": "value2"}
"arg": "value2"
}