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

Merge #3066

Close #3065

Co-authored-by: Raman Mangla <[email protected]>
PiperOrigin-RevId: 825880794

Author: 2 people

src/google/adk/models/lite_llm.py

@@ -384,8 +384,8 @@ def _to_litellm_role(role: Optional[str]) -> Literal["user", "assistant"]:
 
 
 def _schema_to_dict(schema: types.Schema) -> dict:
-  """
-  Recursively converts a types.Schema to a pure-python dict
+  """Recursively converts a types.Schema to a pure-python dict
+
   with all enum values written as lower-case strings.
 
   Args:
@@ -631,7 +631,8 @@ def _get_completion_inputs(
     llm_request: The LlmRequest to convert.
 
   Returns:
-    The litellm inputs (message list, tool dictionary, response format and generation params).
+    The litellm inputs (message list, tool dictionary, response format and
+    generation params).
   """
   # 1. Construct messages
   messages: List[Message] = []
@@ -905,7 +906,7 @@ async def generate_content_async(
       tools = None
 
     completion_args = {
-        "model": self.model,
+        "model": llm_request.model or self.model,
         "messages": messages,
         "tools": tools,
         "response_format": response_format,

tests/unittests/models/test_litellm.py

@@ -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