diff --git a/google/genai/_gaos/google_genai.py b/google/genai/_gaos/google_genai.py index 93fd0b1d9..95c9fdf7c 100644 --- a/google/genai/_gaos/google_genai.py +++ b/google/genai/_gaos/google_genai.py @@ -1382,17 +1382,18 @@ def _output_properties_from_steps(steps: list[Any]) -> dict[str, Any]: collecting = False for step in reversed(steps): - if _get_value(step, 'type') == 'user_input': + step_type = _get_value(step, 'type') + if step_type == 'user_input': break - if _get_value(step, 'type') != 'model_output': + if step_type == 'thought': + continue + if step_type != 'model_output': if collecting: break continue content = _get_value(step, 'content') if not isinstance(content, list): - if collecting: - break continue should_stop = False diff --git a/google/genai/_gaos/types/interactions/interaction.py b/google/genai/_gaos/types/interactions/interaction.py index 06f14a07b..29be08c22 100644 --- a/google/genai/_gaos/types/interactions/interaction.py +++ b/google/genai/_gaos/types/interactions/interaction.py @@ -468,14 +468,14 @@ def _populate_output_helpers(self): step_type = getattr(step, "type", None) if step_type == "user_input": break + if step_type == "thought": + continue if step_type != "model_output": if collecting: break continue content = getattr(step, "content", None) if not isinstance(content, list): - if collecting: - break continue should_stop = False for item in reversed(content): diff --git a/google/genai/tests/interactions/test_integration.py b/google/genai/tests/interactions/test_integration.py index b33733f4e..27787e420 100644 --- a/google/genai/tests/interactions/test_integration.py +++ b/google/genai/tests/interactions/test_integration.py @@ -153,7 +153,6 @@ async def test_async_client_genai_default_timeout(): assert async_sdk_client.timeout != httpx.Timeout(5.0) - @pytest.mark.filterwarnings("error") def test_unrecognized_model_serialization(): from ..._gaos.types.interactions.createmodelinteraction import CreateModelInteraction @@ -231,3 +230,111 @@ def test_interaction_lenient_parsing(): assert isinstance(parsed, gaos_interaction.Interaction) assert parsed.id == "x" + +def test_output_text_concatenates_multi_step_model_response(): + """Tests that output_text concatenates text across consecutive model_output and thought steps.""" + from ..._gaos.types.interactions.interaction import Interaction # pylint: disable=g-import-not-at-top + + raw_steps = [ + { + "type": "thought", + "content": [{"type": "text", "text": "Initial reasoning..."}], + }, + { + "type": "model_output", + "content": [{"type": "text", "text": "First part. "}], + }, + { + "type": "thought", + "content": [{"type": "text", "text": "Intermediate reasoning..."}], + }, + { + "type": "model_output", + "content": [{"type": "text", "text": "Second part. "}], + }, + { + "type": "model_output", + }, + { + "type": "model_output", + "content": [{"type": "text", "text": "Third part."}], + }, + { + "type": "thought", + "content": [{"type": "text", "text": "Trailing thought"}], + }, + ] + + interaction = Interaction.model_validate({ + "id": "v1_multi_step_123", + "status": "completed", + "steps": raw_steps, + }) + assert interaction.output_text == "First part. Second part. Third part." + + wrapped = gaos_google_genai._add_output_properties_if_interaction({ # pylint: disable=protected-access + "id": "v1_multi_step_123", + "status": "completed", + "steps": raw_steps, + }) + assert wrapped["output_text"] == "First part. Second part. Third part." + + +def test_output_text_stops_at_tool_and_user_boundaries(): + """Tests that output_text stops at tool call and user input boundaries.""" + from ..._gaos.types.interactions.interaction import Interaction # pylint: disable=g-import-not-at-top + + steps_with_tool = [ + { + "type": "model_output", + "content": [{"type": "text", "text": "Before tool call. "}], + }, + { + "type": "function_call", + "id": "call_1", + "name": "get_weather", + "arguments": {"city": "Mountain View"}, + }, + { + "type": "function_result", + "call_id": "call_1", + "name": "get_weather", + "result": "Sunny", + }, + { + "type": "thought", + "content": [{"type": "text", "text": "Post-tool thought 1"}], + }, + { + "type": "model_output", + "content": [{"type": "text", "text": "Final part 1. "}], + }, + { + "type": "thought", + "content": [{"type": "text", "text": "Post-tool thought 2"}], + }, + { + "type": "model_output", + "content": [{"type": "text", "text": "Final part 2."}], + }, + { + "type": "function_call", + "id": "call_trailing", + "name": "noop", + "arguments": {}, + }, + ] + + interaction = Interaction.model_validate({ + "id": "v1_tool_boundary_123", + "status": "completed", + "steps": steps_with_tool, + }) + assert interaction.output_text == "Final part 1. Final part 2." + + wrapped = gaos_google_genai._add_output_properties_if_interaction({ # pylint: disable=protected-access + "id": "v1_tool_boundary_123", + "status": "completed", + "steps": steps_with_tool, + }) + assert wrapped["output_text"] == "Final part 1. Final part 2."