fix: Update custom gemini llm connection logic to be used for all 3_x models, not just 3.1

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 940680554

Author: wukath

src/google/adk/flows/llm_flows/basic.py

@@ -120,14 +120,14 @@ def _build_basic_request(
       getattr(getattr(agent, 'canonical_live_model', None), 'model', None)
       or llm_request.model
   )
-  is_gemini_31 = model_name_utils.is_gemini_3_1_flash_live(active_model_name)
+  is_gemini_3_x = model_name_utils._is_gemini_3_x_live(active_model_name)
   llm_request.live_connect_config.enable_affective_dialog = (
       None
-      if is_gemini_31
+      if is_gemini_3_x
       else invocation_context.run_config.enable_affective_dialog
   )
   llm_request.live_connect_config.proactivity = (
-      None if is_gemini_31 else invocation_context.run_config.proactivity
+      None if is_gemini_3_x else invocation_context.run_config.proactivity
   )
   llm_request.live_connect_config.session_resumption = (
       invocation_context.run_config.session_resumption

src/google/adk/models/gemini_llm_connection.py

@@ -50,7 +50,7 @@ def __init__(
     self._output_transcription_text: str = ''
     self._api_backend = api_backend
     self._model_version = model_version
-    self._is_gemini_3_1_flash_live = model_name_utils.is_gemini_3_1_flash_live(
+    self._is_gemini_3_x_live = model_name_utils._is_gemini_3_x_live(
         model_version
     )
     self._is_gemini_3_5_live_translate = (
@@ -115,11 +115,11 @@ async def send_content(self, content: types.Content):
     else:
       logger.debug('Sending LLM new content %s', content)
       if (
-          self._is_gemini_3_1_flash_live
+          self._is_gemini_3_x_live
           and len(content.parts) == 1
           and content.parts[0].text
       ):
-        logger.debug('Using send_realtime_input for Gemini 3.1 text input')
+        logger.debug('Using send_realtime_input for Gemini 3.x Live text input')
         await self._gemini_session.send_realtime_input(
             text=content.parts[0].text
         )
@@ -140,7 +140,7 @@ async def send_realtime(self, input: RealtimeInput):
     if isinstance(input, types.Blob):
       # The blob is binary and is very large. So let's not log it.
       logger.debug('Sending LLM Blob.')
-      if self._is_gemini_3_1_flash_live or self._is_gemini_3_5_live_translate:
+      if self._is_gemini_3_x_live or self._is_gemini_3_5_live_translate:
         if input.mime_type and input.mime_type.startswith('audio/'):
           await self._gemini_session.send_realtime_input(audio=input)
         elif input.mime_type and input.mime_type.startswith('image/'):
@@ -362,9 +362,9 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
           # generation_complete, causing transcription to appear after
           # tool_call in the session log.
           if message.server_content.input_transcription:
-            # Gemini 3.1 Flash Live only sends a single final input
+            # Gemini 3.x Live only sends a single final input
             # transcription
-            if self._is_gemini_3_1_flash_live:
+            if self._is_gemini_3_x_live:
               if message.server_content.input_transcription.text:
                 yield LlmResponse(
                     input_transcription=types.Transcription(
@@ -464,7 +464,7 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
                 or last_grounding_metadata
                 or (
                     types.GroundingMetadata()
-                    if self._is_gemini_3_1_flash_live
+                    if self._is_gemini_3_x_live
                     else None
                 )
             )
@@ -512,7 +512,7 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
                 or last_grounding_metadata
                 or (
                     types.GroundingMetadata()
-                    if self._is_gemini_3_1_flash_live
+                    if self._is_gemini_3_x_live
                     else None
                 ),
                 model_version=self._model_version,
@@ -559,17 +559,17 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
               types.Part(function_call=function_call)
               for function_call in message.tool_call.function_calls
           ])
-          if not self._is_gemini_3_1_flash_live:
+          if not self._is_gemini_3_x_live:
             if tool_call_metadata is None:
               tool_call_metadata = last_grounding_metadata
-          # Gemini 3.1 does not emit turn_complete until it receives the
+          # Gemini 3.x Live does not emit turn_complete until it receives the
           # tool response, so yield tool calls immediately to avoid
           # deadlocking the conversation. Other models (e.g. 2.5-pro,
           # native-audio) send turn_complete after tool calls, so buffer
           # and merge them into a single response at turn_complete.
-          if self._is_gemini_3_1_flash_live and tool_call_parts:
+          if self._is_gemini_3_x_live and tool_call_parts:
             logger.debug(
-                'Yielding tool_call_parts immediately for Gemini 3.1 live tool'
+                'Yielding tool_call_parts immediately for Gemini 3.x live tool'
                 ' call'
             )
             yield LlmResponse(

src/google/adk/utils/model_name_utils.py

@@ -184,19 +184,23 @@ def _is_gemini_eap_model(model_string: Optional[str]) -> bool:
   )
 
 
-def is_gemini_3_1_flash_live(model_string: Optional[str]) -> bool:
-  """Check if the model is a Gemini 3.1 Flash Live model.
+def _is_gemini_3_x_live(model_string: Optional[str]) -> bool:
+  """Check if the model is a Gemini 3.x Live model.
 
   Args:
     model_string: The model name
 
   Returns:
-    True if it's a Gemini 3.1 Flash Live model, False otherwise
+    True if it's a Gemini 3.x Live model, False otherwise
   """
   if not model_string:
     return False
   model_name = extract_model_name(model_string)
-  return model_name.startswith('gemini-3.1-flash-live')
+  return (
+      model_name.startswith('gemini-3.')
+      and '-live' in model_name
+      and not is_gemini_3_5_live_translate(model_string)
+  )
 
 
 def is_gemini_3_5_live_translate(model_string: Optional[str]) -> bool:

tests/unittests/flows/llm_flows/test_basic_processor.py

@@ -210,13 +210,13 @@ async def test_skips_output_schema_for_task_mode(self):
     assert llm_request.config.response_schema is None
 
   @pytest.mark.asyncio
-  async def test_disables_affective_dialog_and_proactivity_for_gemini_3_1_live(
+  async def test_disables_affective_dialog_and_proactivity_for_gemini_3_x_live(
       self,
   ):
-    """Gemini 3.1 Live does not support affective_dialog/proactivity."""
+    """Gemini 3.x Live does not support affective_dialog/proactivity."""
     agent = LlmAgent(
         name='test_agent',
-        model='gemini-3.1-flash-live-preview',
+        model='gemini-3.5-flash-lite-live-preview',
     )
     invocation_context = await _create_invocation_context(agent)
     invocation_context.run_config = RunConfig(
@@ -233,10 +233,10 @@ async def test_disables_affective_dialog_and_proactivity_for_gemini_3_1_live(
     assert llm_request.live_connect_config.proactivity is None
 
   @pytest.mark.asyncio
-  async def test_keeps_affective_dialog_and_proactivity_for_non_gemini_3_1(
+  async def test_keeps_affective_dialog_and_proactivity_for_non_gemini_3_x_live(
       self,
   ):
-    """Non-3.1 live models keep the configured affective_dialog/proactivity."""
+    """Non-3.x live models keep the configured affective_dialog/proactivity."""
     agent = LlmAgent(
         name='test_agent',
         model='gemini-2.5-flash-live',

tests/unittests/utils/test_model_name_utils.py

@@ -15,10 +15,10 @@
 """Tests for model name utility functions."""
 
 from google.adk.models.llm_request import LlmRequest
+from google.adk.utils.model_name_utils import _is_gemini_3_x_live
 from google.adk.utils.model_name_utils import _is_managed_agent
 from google.adk.utils.model_name_utils import extract_model_name
 from google.adk.utils.model_name_utils import is_gemini_1_model
-from google.adk.utils.model_name_utils import is_gemini_3_1_flash_live
 from google.adk.utils.model_name_utils import is_gemini_3_5_live_translate
 from google.adk.utils.model_name_utils import is_gemini_eap_or_2_or_above
 from google.adk.utils.model_name_utils import is_gemini_model
@@ -387,31 +387,33 @@ def test_true_enables_check_bypass(self, monkeypatch):
     assert is_gemini_model_id_check_disabled() is True
 
 
-class TestIsGemini31FlashLive:
-  """Test the is_gemini_3_1_flash_live function."""
+class TestIsGemini3XLive:
+  """Test the _is_gemini_3_x_live function."""
 
-  def test_is_gemini_3_1_flash_live_simple_name(self):
+  def test_is_gemini_3_x_live_simple_name(self):
     """Test with simple model name format."""
-    assert is_gemini_3_1_flash_live('gemini-3.1-flash-live') is True
-    assert is_gemini_3_1_flash_live('gemini-3.1-flash-live-preview') is True
-    assert is_gemini_3_1_flash_live('gemini-3.1-pro-live') is False
-    assert is_gemini_3_1_flash_live('gemini-2.5-flash-live') is False
-
-  def test_is_gemini_3_1_flash_live_path_based_name(self):
+    assert _is_gemini_3_x_live('gemini-3.1-flash-live') is True
+    assert _is_gemini_3_x_live('gemini-3.1-flash-live-preview') is True
+    assert _is_gemini_3_x_live('gemini-3.5-flash-lite-live-preview') is True
+    assert _is_gemini_3_x_live('gemini-3.5-live-translate') is False
+    assert _is_gemini_3_x_live('gemini-3.1-pro') is False
+    assert _is_gemini_3_x_live('gemini-2.5-flash-live') is False
+
+  def test_is_gemini_3_x_live_path_based_name(self):
     """Test with path-based format (Vertex AI etc.)."""
     vertex_path = 'projects/123/locations/us-central1/publishers/google/models/gemini-3.1-flash-live'
-    assert is_gemini_3_1_flash_live(vertex_path) is True
+    assert _is_gemini_3_x_live(vertex_path) is True
 
-    vertex_path_preview = 'projects/123/locations/us-central1/publishers/google/models/gemini-3.1-flash-live-preview'
-    assert is_gemini_3_1_flash_live(vertex_path_preview) is True
+    vertex_path_preview = 'projects/123/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite-live-preview'
+    assert _is_gemini_3_x_live(vertex_path_preview) is True
 
     non_live_path = 'projects/123/locations/us-central1/publishers/google/models/gemini-3.1-flash'
-    assert is_gemini_3_1_flash_live(non_live_path) is False
+    assert _is_gemini_3_x_live(non_live_path) is False
 
-  def test_is_gemini_3_1_flash_live_edge_cases(self):
+  def test_is_gemini_3_x_live_edge_cases(self):
     """Test edge cases."""
-    assert is_gemini_3_1_flash_live(None) is False
-    assert is_gemini_3_1_flash_live('') is False
+    assert _is_gemini_3_x_live(None) is False
+    assert _is_gemini_3_x_live('') is False
 
 
 class TestIsGemini35LiveTranslate: