fix: Allow LLM request to override the model used in the generate content async method in LiteLLM

Merge https://github.com/google/adk-python/pull/3066

Close #3065

Co-authored-by: Raman Mangla <ramanmangla@google.com>
PiperOrigin-RevId: 825880794
This commit is contained in:
Ayam
2025-10-29 23:34:33 -07:00
committed by Copybara-Service
co-authored by Raman Mangla
parent 3814d8b80f
commit ce8f674a28
2 changed files with 52 additions and 4 deletions
+47
View File
@@ -549,6 +549,53 @@ async def test_generate_content_async(mock_acompletion, lite_llm_instance):
)
@pytest.mark.asyncio
async def test_generate_content_async_with_model_override(
mock_acompletion, lite_llm_instance
):
llm_request = LlmRequest(
model="overridden_model",
contents=[
types.Content(
role="user", parts=[types.Part.from_text(text="Test prompt")]
)
],
)
async for response in lite_llm_instance.generate_content_async(llm_request):
assert response.content.role == "model"
assert response.content.parts[0].text == "Test response"
mock_acompletion.assert_called_once()
_, kwargs = mock_acompletion.call_args
assert kwargs["model"] == "overridden_model"
assert kwargs["messages"][0]["role"] == "user"
assert kwargs["messages"][0]["content"] == "Test prompt"
@pytest.mark.asyncio
async def test_generate_content_async_without_model_override(
mock_acompletion, lite_llm_instance
):
llm_request = LlmRequest(
model=None,
contents=[
types.Content(
role="user", parts=[types.Part.from_text(text="Test prompt")]
)
],
)
async for response in lite_llm_instance.generate_content_async(llm_request):
assert response.content.role == "model"
mock_acompletion.assert_called_once()
_, kwargs = mock_acompletion.call_args
assert kwargs["model"] == "test_model"
@pytest.mark.asyncio
async def test_generate_content_async_adds_fallback_user_message(
mock_acompletion, lite_llm_instance