diff --git a/src/google/adk/models/gemini_llm_connection.py b/src/google/adk/models/gemini_llm_connection.py index cc380f229a..3356e45715 100644 --- a/src/google/adk/models/gemini_llm_connection.py +++ b/src/google/adk/models/gemini_llm_connection.py @@ -498,6 +498,64 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]: live_session_id=live_session_id, ) self._output_transcription_text = '' + # in case of empty content or parts, we still surface it + # in case it's an interrupted message, we merge the previous partial + # text. Other we don't merge. because content can be none when model + # safety threshold is triggered + if message.server_content.interrupted: + if text: + yield self.__build_full_text_response( + text, + is_thought, + last_grounding_metadata, + interrupted=True, + ) + text = '' + is_thought = False + last_grounding_metadata = None + else: + yield LlmResponse( + interrupted=message.server_content.interrupted, + grounding_metadata=last_grounding_metadata, + model_version=self._model_version, + live_session_id=live_session_id, + ) + last_grounding_metadata = None + if message.tool_call: + logger.debug('Received tool call: %s', message.tool_call) + if text: + yield self.__build_full_text_response( + text, is_thought, last_grounding_metadata + ) + text = '' + is_thought = False + last_grounding_metadata = None + tool_call_parts.extend([ + types.Part(function_call=function_call) + for function_call in message.tool_call.function_calls or [] + ]) + if not self._is_gemini_3_x_live: + if tool_call_metadata is None: + tool_call_metadata = last_grounding_metadata + # 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_x_live and tool_call_parts: + logger.debug( + 'Yielding tool_call_parts immediately for Gemini 3.x live tool' + ' call' + ) + yield LlmResponse( + content=types.Content(role='model', parts=tool_call_parts), + grounding_metadata=last_grounding_metadata, + model_version=self._model_version, + live_session_id=live_session_id, + ) + tool_call_parts = [] + last_grounding_metadata = None + if message.server_content: if message.server_content.turn_complete: # Capture final grounding metadata before last_grounding_metadata is cleared in the next block. final_grounding_metadata = ( @@ -564,63 +622,6 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]: ) last_grounding_metadata = None # Reset after yielding break - # in case of empty content or parts, we still surface it - # in case it's an interrupted message, we merge the previous partial - # text. Other we don't merge. because content can be none when model - # safety threshold is triggered - if message.server_content.interrupted: - if text: - yield self.__build_full_text_response( - text, - is_thought, - last_grounding_metadata, - interrupted=True, - ) - text = '' - is_thought = False - last_grounding_metadata = None - else: - yield LlmResponse( - interrupted=message.server_content.interrupted, - grounding_metadata=last_grounding_metadata, - model_version=self._model_version, - live_session_id=live_session_id, - ) - last_grounding_metadata = None - if message.tool_call: - logger.debug('Received tool call: %s', message.tool_call) - if text: - yield self.__build_full_text_response( - text, is_thought, last_grounding_metadata - ) - text = '' - is_thought = False - last_grounding_metadata = None - tool_call_parts.extend([ - types.Part(function_call=function_call) - for function_call in message.tool_call.function_calls or [] - ]) - if not self._is_gemini_3_x_live: - if tool_call_metadata is None: - tool_call_metadata = last_grounding_metadata - # 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_x_live and tool_call_parts: - logger.debug( - 'Yielding tool_call_parts immediately for Gemini 3.x live tool' - ' call' - ) - yield LlmResponse( - content=types.Content(role='model', parts=tool_call_parts), - grounding_metadata=last_grounding_metadata, - model_version=self._model_version, - live_session_id=live_session_id, - ) - tool_call_parts = [] - last_grounding_metadata = None if message.session_resumption_update: logger.debug('Received session resumption message: %s', message) yield ( diff --git a/tests/unittests/models/test_gemini_llm_connection.py b/tests/unittests/models/test_gemini_llm_connection.py index 3f141af3b9..5314ea7c4a 100644 --- a/tests/unittests/models/test_gemini_llm_connection.py +++ b/tests/unittests/models/test_gemini_llm_connection.py @@ -2377,3 +2377,70 @@ async def mock_receive_generator(): assert len(responses) == 1 assert responses[0].voice_activity == mock_vad + + +@pytest.mark.asyncio +async def test_receive_tool_call_and_turn_complete_in_same_message( + gemini_connection, mock_gemini_session +): + """Test receive correctly handles tool_call and turn_complete in the same frame.""" + mock_msg = mock.create_autospec(types.LiveServerMessage, instance=True) + mock_msg.usage_metadata = None + mock_msg.session_resumption_update = None + mock_msg.go_away = None + mock_msg.voice_activity = None + + # Set up server_content with text and turn_complete + part = types.Part.from_text(text='I will call the tool.') + mock_model_turn = mock.create_autospec(types.Content, instance=True) + mock_model_turn.parts = [part] + + mock_server_content = mock.create_autospec( + types.LiveServerContent, instance=True + ) + mock_server_content.model_turn = mock_model_turn + mock_server_content.grounding_metadata = None + mock_server_content.turn_complete = True + mock_server_content.interrupted = False + mock_server_content.input_transcription = None + mock_server_content.output_transcription = None + mock_server_content.generation_complete = False + mock_msg.server_content = mock_server_content + + # Set up tool_call + function_call = types.FunctionCall( + name='my_tool', + args={'arg': 'value1'}, + ) + mock_tool_call = mock.create_autospec(types.LiveServerToolCall, instance=True) + mock_tool_call.function_calls = [function_call] + mock_msg.tool_call = mock_tool_call + + async def mock_receive_generator(): + yield mock_msg + + mock_gemini_session.receive = mock.Mock(return_value=mock_receive_generator()) + + # Disable 3.x live streaming logic to test default turn-complete buffering + gemini_connection._is_gemini_3_x_live = False + + received_texts = [] + received_tools = [] + turn_complete_seen = False + + async for response in gemini_connection.receive(): + if response.content: + for part in response.content.parts: + if getattr(part, 'text', None): + received_texts.append(part.text) + if getattr(part, 'function_call', None): + received_tools.append(part.function_call.name) + if getattr(response, 'turn_complete', False): + turn_complete_seen = True + + # Ensure the text was received + assert 'I will call the tool.' in received_texts + # Ensure the tool call was received (not swallowed) + assert 'my_tool' in received_tools + # Ensure turn complete was registered + assert turn_complete_seen is True