diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 1f49fa65a4..80fda09ab7 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -983,23 +983,20 @@ def get_author_for_event(llm_response: LlmResponse) -> str: async for event in agen: # Cache output audio chunks from model responses # TODO: support video data - if ( - run_config.save_live_blob - and event.content - and event.content.parts - and event.content.parts[0].inline_data - and event.content.parts[0].inline_data.mime_type - and event.content.parts[0].inline_data.mime_type.startswith( - 'audio/' - ) - ): - audio_blob = types.Blob( - data=event.content.parts[0].inline_data.data, - mime_type=event.content.parts[0].inline_data.mime_type, - ) - self.audio_cache_manager.cache_audio( - invocation_context, audio_blob, cache_type='output' - ) + if run_config.save_live_blob and event.content and event.content.parts: + for part in event.content.parts: + if ( + part.inline_data + and part.inline_data.mime_type + and part.inline_data.mime_type.startswith('audio/') + ): + audio_blob = types.Blob( + data=part.inline_data.data, + mime_type=part.inline_data.mime_type, + ) + self.audio_cache_manager.cache_audio( + invocation_context, audio_blob, cache_type='output' + ) yield event # Give opportunity for other tasks to run. diff --git a/src/google/adk/models/gemini_llm_connection.py b/src/google/adk/models/gemini_llm_connection.py index cc380f229a..321d249461 100644 --- a/src/google/adk/models/gemini_llm_connection.py +++ b/src/google/adk/models/gemini_llm_connection.py @@ -124,7 +124,7 @@ async def _send_content( complete the model turn. """ assert content.parts - if content.parts[0].function_response: + if any(p.function_response for p in content.parts): # All parts have to be function responses. function_responses = [ function_response @@ -380,24 +380,25 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]: llm_response.grounding_metadata = ( message.server_content.grounding_metadata ) - if content.parts[0].text: - current_is_thought = getattr(content.parts[0], 'thought', False) - if text and current_is_thought != is_thought: - yield self.__build_full_text_response(text, is_thought) + for part in content.parts: + if part.text: + current_is_thought = getattr(part, 'thought', False) + if text and current_is_thought != is_thought: + yield self.__build_full_text_response(text, is_thought) + text = '' + is_thought = False + + text += part.text + is_thought = current_is_thought + llm_response.partial = True + # don't yield the merged text event when receiving audio data + elif text and not part.inline_data: + yield self.__build_full_text_response( + text, is_thought, last_grounding_metadata + ) text = '' is_thought = False - - text += content.parts[0].text - is_thought = current_is_thought - llm_response.partial = True - # don't yield the merged text event when receiving audio data - elif text and not content.parts[0].inline_data: - yield self.__build_full_text_response( - text, is_thought, last_grounding_metadata - ) - text = '' - is_thought = False - last_grounding_metadata = None + last_grounding_metadata = None yield llm_response # Note: in some cases, tool_call may arrive before # generation_complete, causing transcription to appear after diff --git a/tests/unittests/models/test_gemini_llm_connection.py b/tests/unittests/models/test_gemini_llm_connection.py index 3f141af3b9..a1a155bc16 100644 --- a/tests/unittests/models/test_gemini_llm_connection.py +++ b/tests/unittests/models/test_gemini_llm_connection.py @@ -190,6 +190,29 @@ async def test_send_content_function_response( assert call_args['function_responses'] == [function_response] +@pytest.mark.asyncio +async def test_send_content_mixed_content_raises_value_error( + gemini_connection, mock_gemini_session +): + """Test send_content with mixed text and function response raises ValueError.""" + function_response = types.FunctionResponse( + name='test_function', response={'result': 'success'} + ) + content = types.Content( + role='user', + parts=[ + types.Part.from_text(text='Hello'), + types.Part(function_response=function_response), + ], + ) + + with pytest.raises( + ValueError, + match='Function-response content cannot mix function and non-function parts.', + ): + await gemini_connection.send_content(content) + + @pytest.mark.asyncio async def test_close(gemini_connection, mock_gemini_session): """Test close method.""" @@ -1589,6 +1612,60 @@ async def mock_receive_generator(): assert responses[5].turn_complete is True +@pytest.mark.asyncio +async def test_receive_multiplexed_thought_and_text( + gemini_connection, mock_gemini_session +): + """Test receive with multiplexed thought and text in a single chunk.""" + part1 = types.Part.from_text(text='Let me think.') + part1.thought = True + part2 = types.Part.from_text(text=' Hello.') + part2.thought = False + mock_content = types.Content( + role='model', + parts=[part1, part2], + ) + mock_server_content = mock.Mock() + mock_server_content.model_turn = mock_content + mock_server_content.interrupted = False + mock_server_content.input_transcription = None + mock_server_content.output_transcription = None + mock_server_content.turn_complete = True + mock_server_content.grounding_metadata = None + mock_server_content.turn_complete_reason = None + + mock_message = mock.AsyncMock() + mock_message.usage_metadata = None + mock_message.server_content = mock_server_content + mock_message.tool_call = None + mock_message.session_resumption_update = None + mock_message.go_away = None + mock_message.voice_activity = None + + async def mock_receive_generator(): + yield mock_message + + receive_mock = mock.Mock(return_value=mock_receive_generator()) + mock_gemini_session.receive = receive_mock + + responses = [resp async for resp in gemini_connection.receive()] + + assert len(responses) == 4 + + assert responses[0].content.parts[0].text == 'Let me think.' + assert responses[0].content.parts[0].thought is True + assert responses[0].partial is False + + assert responses[1].content == mock_content + assert responses[1].partial is True + + assert responses[2].content.parts[0].text == ' Hello.' + assert not getattr(responses[2].content.parts[0], 'thought', False) + assert responses[2].partial is False + + assert responses[3].turn_complete is True + + @pytest.mark.asyncio async def test_receive_video_content(gemini_connection, mock_gemini_session): """Test receive with video content."""