Skip to content

Commit a9cabe7

Browse files
AtulJoshi1206copybara-github
authored andcommitted
fix: keep the streamed thought signature on its own part
Merge #7588 PiperOrigin-RevId: 1010761841
1 parent 2a0bffa commit a9cabe7

2 files changed

Lines changed: 276 additions & 9 deletions

File tree

‎src/google/adk/models/interactions_utils.py‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1055,12 +1055,10 @@ def _handle_thought_signature(
10551055
signature = delta.signature
10561056
if not signature:
10571057
return None
1058-
for part in reversed(state.parts):
1059-
if part.thought:
1060-
part.thought_signature = base64.b64decode(signature)
1061-
return None
1062-
# A signature often arrives with no preceding thought summary, leaving no
1063-
# part to attach it to. Keep it on its own part rather than dropping it.
1058+
# Keep every signature on a part of its own, as the non-streaming conversion
1059+
# does. On the step's thought summary part it would be re-encoded as a text
1060+
# step, which drops the signature, and on an earlier step's signature part it
1061+
# would overwrite that step's signature.
10641062
state.parts.append(
10651063
types.Part(thought=True, thought_signature=base64.b64decode(signature))
10661064
)

‎tests/unittests/models/test_interactions_utils.py‎

Lines changed: 272 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,18 @@
2323
import json
2424
import logging
2525
import typing
26+
from unittest.mock import AsyncMock
2627
from unittest.mock import MagicMock
2728

29+
from google.adk.agents._streaming_mode import StreamingMode
30+
from google.adk.agents.llm_agent import LlmAgent
31+
from google.adk.agents.run_config import RunConfig
2832
from google.adk.models import interactions_utils
33+
from google.adk.models.google_llm import Gemini
2934
from google.adk.models.llm_request import LlmRequest
3035
from google.adk.models.llm_response import LlmResponse
36+
from google.adk.runners import InMemoryRunner
37+
from google.genai import Client
3138
from google.genai import interactions
3239
from google.genai import types
3340
from google.genai.interactions import CodeExecutionResultStep
@@ -37,6 +44,7 @@
3744
from google.genai.interactions import Interaction
3845
from google.genai.interactions import InteractionCompletedEvent
3946
from google.genai.interactions import InteractionCreatedEvent
47+
from google.genai.interactions import InteractionSSEEvent
4048
from google.genai.interactions import InteractionSseEventInteraction
4149
from google.genai.interactions import ModelOutputStep
4250
from google.genai.interactions import StepDelta
@@ -1911,8 +1919,15 @@ def test_thought_summary_delta_ignores_non_text_content(self):
19111919
assert result is None
19121920
assert state.parts == []
19131921

1914-
def test_thought_signature_delta_attaches_to_last_thought(self):
1915-
"""thought_signature mutates the last thought part and emits no event."""
1922+
def test_thought_signature_delta_after_a_summary_gets_its_own_part(
1923+
self,
1924+
) -> None:
1925+
"""The signature follows the summary on a text-free part of its own.
1926+
1927+
A part that also carries text is re-encoded as a text step, which has no
1928+
signature field, so a signature written onto the summary part would be
1929+
dropped on the next request.
1930+
"""
19161931
state = interactions_utils._StreamState()
19171932
interactions_utils.convert_interaction_event_to_llm_response(
19181933
StepDelta(
@@ -1936,8 +1951,14 @@ def test_thought_signature_delta_attaches_to_last_thought(self):
19361951
state,
19371952
interaction_id='int_ts',
19381953
)
1954+
19391955
assert result is None
1940-
assert state.parts[-1].thought_signature == b'sig-bytes'
1956+
assert len(state.parts) == 2
1957+
assert state.parts[0].text == 'reasoning'
1958+
assert state.parts[0].thought_signature is None
1959+
assert state.parts[1].thought is True
1960+
assert state.parts[1].text is None
1961+
assert state.parts[1].thought_signature == b'sig-bytes'
19411962

19421963
def test_thought_signature_delta_alone_becomes_its_own_part(self):
19431964
"""A signature with no preceding thought summary still reaches the stream.
@@ -2012,6 +2033,67 @@ def test_thought_signature_delta_without_signature_adds_no_part(self):
20122033
assert result is None
20132034
assert state.parts == []
20142035

2036+
def test_thought_signature_deltas_of_two_thought_steps_are_both_kept(
2037+
self,
2038+
) -> None:
2039+
"""Each thought step keeps its own signature, in step order.
2040+
2041+
The non-streaming conversion gives every thought step a part of its own,
2042+
so a later step's signature must not overwrite an earlier one's.
2043+
"""
2044+
state = interactions_utils._StreamState()
2045+
events: list[InteractionSSEEvent] = [
2046+
StepDelta(
2047+
event_type='step.delta',
2048+
index=0,
2049+
delta={
2050+
'type': 'thought_signature',
2051+
'signature': base64.b64encode(b'first-sig').decode('utf-8'),
2052+
},
2053+
),
2054+
StepStart(
2055+
event_type='step.start',
2056+
index=1,
2057+
step=FunctionCallStep(
2058+
type='function_call',
2059+
id='call_1',
2060+
name='get_weather',
2061+
arguments={},
2062+
),
2063+
),
2064+
StepStop(event_type='step.stop', index=1),
2065+
StepDelta(
2066+
event_type='step.delta',
2067+
index=2,
2068+
delta={
2069+
'type': 'thought_signature',
2070+
'signature': base64.b64encode(b'second-sig').decode('utf-8'),
2071+
},
2072+
),
2073+
StepStart(
2074+
event_type='step.start',
2075+
index=3,
2076+
step=FunctionCallStep(
2077+
type='function_call',
2078+
id='call_2',
2079+
name='get_time',
2080+
arguments={},
2081+
),
2082+
),
2083+
StepStop(event_type='step.stop', index=3),
2084+
]
2085+
for event in events:
2086+
interactions_utils.convert_interaction_event_to_llm_response(
2087+
event, state, interaction_id='int_ts'
2088+
)
2089+
2090+
assert [part.thought_signature for part in state.parts] == [
2091+
b'first-sig',
2092+
None,
2093+
b'second-sig',
2094+
None,
2095+
]
2096+
20152097
def test_audio_delta_with_data(self):
20162098
"""audio delta becomes an inline_data part via the shared media handler."""
20172099
event = StepDelta(
@@ -2931,6 +3013,193 @@ def test_arguments_delta_for_finalized_call_logs_warning(self, caplog):
29313013
assert state.parts[0].function_call.args == {'city': 'Paris'}
29323014

29333015

3016+
class TestThoughtSignatureRoundTrip:
3017+
"""A tool-calling turn's thought signature reaches the request that follows.
3018+
3019+
With thought summaries on, the summary streams ahead of the signature. Kept
3020+
on the summary part, the signature is lost on the ``previous_interaction_id``
3021+
path: the summary is echoed back as ``model_output`` text and the signature
3022+
is dropped. A stateless Gemini 3 request carrying the tool result without
3023+
the signature is rejected.
3024+
"""
3025+
3026+
_SIGNATURE = base64.b64encode(b'sig-bytes').decode('utf-8')
3027+
_SUMMARY = 'I should check the weather.'
3028+
3029+
def _streamed_events(self) -> list[InteractionSSEEvent]:
3030+
return [
3031+
StepDelta(
3032+
event_type='step.delta',
3033+
index=0,
3034+
delta={
3035+
'type': 'thought_summary',
3036+
'content': {'type': 'text', 'text': self._SUMMARY},
3037+
},
3038+
),
3039+
StepDelta(
3040+
event_type='step.delta',
3041+
index=0,
3042+
delta={'type': 'thought_signature', 'signature': self._SIGNATURE},
3043+
),
3044+
StepStart(
3045+
event_type='step.start',
3046+
index=1,
3047+
step=FunctionCallStep(
3048+
type='function_call',
3049+
id='call_1',
3050+
name='get_weather',
3051+
arguments={},
3052+
),
3053+
),
3054+
StepDelta(
3055+
event_type='step.delta',
3056+
index=1,
3057+
delta={'type': 'arguments_delta', 'arguments': '{"city": "Tokyo"}'},
3058+
),
3059+
StepStop(event_type='step.stop', index=1),
3060+
InteractionCompletedEvent(
3061+
event_type='interaction.completed',
3062+
interaction=InteractionSseEventInteraction(
3063+
id='int_1', status='requires_action', steps=[]
3064+
),
3065+
),
3066+
]
3067+
3068+
def _streamed_turn(self) -> types.Content:
3069+
state = interactions_utils._StreamState()
3070+
responses = [
3071+
interactions_utils.convert_interaction_event_to_llm_response(
3072+
event, state, interaction_id='int_1'
3073+
)
3074+
for event in self._streamed_events()
3075+
]
3076+
final_response = responses[-1]
3077+
assert final_response is not None
3078+
assert final_response.content is not None
3079+
return final_response.content
3080+
3081+
def _unary_turn(self) -> types.Content:
3082+
now = datetime.now(timezone.utc).isoformat()
3083+
interaction = Interaction(
3084+
id='int_1',
3085+
created=now,
3086+
updated=now,
3087+
status='requires_action',
3088+
steps=[
3089+
ThoughtStep(
3090+
type='thought',
3091+
signature=self._SIGNATURE,
3092+
summary=[TextContent(type='text', text=self._SUMMARY)],
3093+
),
3094+
FunctionCallStep(
3095+
type='function_call',
3096+
id='call_1',
3097+
name='get_weather',
3098+
arguments={'city': 'Tokyo'},
3099+
),
3100+
],
3101+
)
3102+
response = interactions_utils.convert_interaction_to_llm_response(
3103+
interaction
3104+
)
3105+
assert response.content is not None
3106+
return response.content
3107+
3108+
@pytest.mark.parametrize('streamed', [True, False], ids=['streamed', 'unary'])
3109+
def test_signature_is_sent_back_with_the_tool_result(
3110+
self, streamed: bool
3111+
) -> None:
3112+
"""The follow-up request echoes the signature and not the summary."""
3113+
model_turn = self._streamed_turn() if streamed else self._unary_turn()
3114+
contents = [
3115+
types.Content(role='user', parts=[types.Part(text='Weather?')]),
3116+
model_turn,
3117+
types.Content(
3118+
role='user',
3119+
parts=[
3120+
types.Part(
3121+
function_response=types.FunctionResponse(
3122+
id='call_1',
3123+
name='get_weather',
3124+
response={'weather': 'sunny'},
3125+
)
3126+
)
3127+
],
3128+
),
3129+
]
3130+
3131+
steps = interactions_utils._convert_contents_to_steps(
3132+
interactions_utils._get_latest_user_contents(contents)
3133+
)
3134+
3135+
assert [step['type'] for step in steps] == ['thought', 'function_result']
3136+
assert steps[0] == {'type': 'thought', 'signature': self._SIGNATURE}
3137+
3138+
async def test_runner_sends_the_streamed_signature_back(self) -> None:
3139+
"""An agent streaming over the interactions API echoes the signature."""
3140+
client = MagicMock(spec=Client)
3141+
client.vertexai = False
3142+
client.aio.interactions.create = AsyncMock(
3143+
side_effect=[
3144+
_MockAsyncIterator([*self._streamed_events()]),
3145+
_MockAsyncIterator([
3146+
StepDelta(
3147+
event_type='step.delta',
3148+
index=0,
3149+
delta={'type': 'text', 'text': 'It is sunny.'},
3150+
),
3151+
InteractionCompletedEvent(
3152+
event_type='interaction.completed',
3153+
interaction=InteractionSseEventInteraction(
3154+
id='int_2', status='completed', steps=[]
3155+
),
3156+
),
3157+
]),
3158+
]
3159+
)
3160+
3161+
def get_weather(city: str) -> dict[str, str]:
3162+
"""Returns the weather in a city."""
3163+
del city
3164+
return {'weather': 'sunny'}
3165+
3166+
runner = InMemoryRunner(
3167+
agent=LlmAgent(
3168+
name='weather_agent',
3169+
model=Gemini(
3170+
model='gemini-3-flash-preview',
3171+
use_interactions_api=True,
3172+
client=client,
3173+
),
3174+
tools=[get_weather],
3175+
),
3176+
app_name='weather_app',
3177+
)
3178+
session = await runner.session_service.create_session(
3179+
app_name='weather_app', user_id='user'
3180+
)
3181+
async for _ in runner.run_async(
3182+
user_id='user',
3183+
session_id=session.id,
3184+
new_message=types.Content(
3185+
role='user', parts=[types.Part(text='Weather in Tokyo?')]
3186+
),
3187+
run_config=RunConfig(streaming_mode=StreamingMode.SSE),
3188+
):
3189+
pass
3190+
3191+
follow_up = client.aio.interactions.create.call_args_list[1].kwargs
3192+
assert follow_up['previous_interaction_id'] == 'int_1'
3193+
assert [step['type'] for step in follow_up['input']] == [
3194+
'thought',
3195+
'function_result',
3196+
]
3197+
assert follow_up['input'][0] == {
3198+
'type': 'thought',
3199+
'signature': self._SIGNATURE,
3200+
}
3201+
3202+
29343203
@pytest.mark.parametrize(
29353204
('streamed_events_factory', 'expected_ids'),
29363205
[

0 commit comments

Comments
 (0)