Skip to content

Commit b382e01

Browse files
yyyu-googlecopybara-github
yyyu-google
authored and
copybara-github
committed
fix: do not record incomplete continuation stream turns in curated chat history
PiperOrigin-RevId: 996900716
1 parent bf9d6d7 commit b382e01

2 files changed

Lines changed: 183 additions & 10 deletions

File tree

‎google/genai/chats.py‎

Lines changed: 67 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,29 @@ def _validate_response(response: GenerateContentResponse) -> bool:
7575
return _validate_content(response.candidates[0].content)
7676

7777

78+
def _update_stream_finish_reason(
79+
chunk: GenerateContentResponse,
80+
finish_reason: Optional[types.FinishReason],
81+
hop_continuation_token: Optional[bytes],
82+
enable_continuation: bool,
83+
) -> tuple[Optional[types.FinishReason], Optional[bytes]]:
84+
if not chunk.candidates:
85+
return finish_reason, hop_continuation_token
86+
candidate = chunk.candidates[0]
87+
if candidate.continuation_token:
88+
hop_continuation_token = candidate.continuation_token
89+
if candidate.finish_reason:
90+
finish_reason = candidate.finish_reason
91+
if (
92+
enable_continuation
93+
and hop_continuation_token
94+
and _is_resumable_finish_reason(finish_reason)
95+
):
96+
finish_reason = None
97+
hop_continuation_token = None
98+
return finish_reason, hop_continuation_token
99+
100+
78101
def _extract_curated_history(
79102
comprehensive_history: list[Content],
80103
) -> list[Content]:
@@ -542,6 +565,10 @@ def send_message_stream(
542565

543566
if disable_afc:
544567
if isinstance(self._modules, Models):
568+
enable_continuation = _should_enable_automatic_continuation(
569+
parsed_config, default_enabled=True
570+
)
571+
hop_continuation_token = None
545572
for chunk in self._generate_content_stream_with_continuation(
546573
contents=contents_to_model, # type: ignore[arg-type]
547574
config=parsed_config,
@@ -550,8 +577,12 @@ def send_message_stream(
550577
is_valid = False
551578
if chunk.candidates and chunk.candidates[0].content:
552579
model_output.append(chunk.candidates[0].content)
553-
if chunk.candidates and chunk.candidates[0].finish_reason:
554-
finish_reason = chunk.candidates[0].finish_reason
580+
finish_reason, hop_continuation_token = _update_stream_finish_reason(
581+
chunk,
582+
finish_reason,
583+
hop_continuation_token,
584+
enable_continuation,
585+
)
555586
yield chunk
556587
self.record_history(
557588
user_input=user_input,
@@ -579,6 +610,9 @@ def send_message_stream(
579610
f"AFC is enabled with max remote calls: {remaining_remote_calls_afc}."
580611
)
581612
function_map = _extra_utils.get_function_map(parsed_config)
613+
enable_continuation = _should_enable_automatic_continuation(
614+
parsed_config, default_enabled=True
615+
)
582616
i = 0
583617
if isinstance(self._modules, Models):
584618
while remaining_remote_calls_afc > 0:
@@ -599,6 +633,7 @@ def send_message_stream(
599633

600634
model_output = []
601635
finish_reason = None
636+
hop_continuation_token = None
602637
is_valid = True
603638
func_response_parts = []
604639
chunk = None
@@ -622,8 +657,12 @@ def send_message_stream(
622657

623658
if chunk.candidates and chunk.candidates[0].content:
624659
model_output.append(chunk.candidates[0].content)
625-
if chunk.candidates and chunk.candidates[0].finish_reason:
626-
finish_reason = chunk.candidates[0].finish_reason
660+
finish_reason, hop_continuation_token = _update_stream_finish_reason(
661+
chunk,
662+
finish_reason,
663+
hop_continuation_token,
664+
enable_continuation,
665+
)
627666
yield chunk
628667

629668
if is_last_remote_call_afc:
@@ -643,7 +682,7 @@ def send_message_stream(
643682
self.record_history(
644683
user_input=user_input,
645684
model_output=model_output,
646-
is_valid=is_valid,
685+
is_valid=is_valid and finish_reason is not None,
647686
)
648687
user_input = func_response_content
649688

@@ -1103,6 +1142,10 @@ async def async_generator(): # type: ignore[no-untyped-def]
11031142
if disable_afc:
11041143
output_contents = []
11051144
finish_reason = None
1145+
hop_continuation_token = None
1146+
enable_continuation = _should_enable_automatic_continuation(
1147+
parsed_config, default_enabled=True
1148+
)
11061149
is_valid = True
11071150
chunk = None
11081151
async for chunk in self._generate_content_stream_with_continuation(
@@ -1113,8 +1156,12 @@ async def async_generator(): # type: ignore[no-untyped-def]
11131156
is_valid = False
11141157
if chunk.candidates and chunk.candidates[0].content:
11151158
output_contents.append(chunk.candidates[0].content)
1116-
if chunk.candidates and chunk.candidates[0].finish_reason:
1117-
finish_reason = chunk.candidates[0].finish_reason
1159+
finish_reason, hop_continuation_token = _update_stream_finish_reason(
1160+
chunk,
1161+
finish_reason,
1162+
hop_continuation_token,
1163+
enable_continuation,
1164+
)
11181165
yield chunk
11191166

11201167
if not output_contents or finish_reason is None:
@@ -1212,6 +1259,9 @@ async def async_generator(): # type: ignore[no-untyped-def]
12121259
mcp_to_genai_tool_adapters,
12131260
is_caller_method_async=True,
12141261
)
1262+
enable_continuation = _should_enable_automatic_continuation(
1263+
final_parsed_config, default_enabled=True
1264+
)
12151265

12161266
i = 0
12171267
model_output: list[types.Content] = []
@@ -1236,6 +1286,7 @@ async def async_generator(): # type: ignore[no-untyped-def]
12361286

12371287
model_output = []
12381288
finish_reason = None
1289+
hop_continuation_token = None
12391290
is_valid = True
12401291
func_response_parts = []
12411292
chunk = None
@@ -1261,8 +1312,14 @@ async def async_generator(): # type: ignore[no-untyped-def]
12611312

12621313
if chunk.candidates and chunk.candidates[0].content:
12631314
model_output.append(chunk.candidates[0].content)
1264-
if chunk.candidates and chunk.candidates[0].finish_reason:
1265-
finish_reason = chunk.candidates[0].finish_reason
1315+
finish_reason, hop_continuation_token = (
1316+
_update_stream_finish_reason(
1317+
chunk,
1318+
finish_reason,
1319+
hop_continuation_token,
1320+
enable_continuation,
1321+
)
1322+
)
12661323
yield chunk
12671324

12681325
if is_last_remote_call_afc:
@@ -1282,7 +1339,7 @@ async def async_generator(): # type: ignore[no-untyped-def]
12821339
self.record_history(
12831340
user_input=user_input,
12841341
model_output=model_output,
1285-
is_valid=is_valid,
1342+
is_valid=is_valid and finish_reason is not None,
12861343
)
12871344
user_input = func_response_content
12881345

‎google/genai/tests/chats/test_continuation_token.py‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1618,3 +1618,119 @@ async def test_async_chat_automatic_continuation_config_rules(mock_api_client):
16181618
)
16191619
assert mock_gc.call_count == 1
16201620
assert resp.text == 'Async Hop 1. '
1621+
1622+
1623+
def test_chat_send_message_stream_incomplete_continuation_not_recorded(
1624+
mock_api_client,
1625+
):
1626+
"""When a continuation hop is cut off without finish_reason, the turn is excluded from curated history."""
1627+
models_module = models.Models(mock_api_client)
1628+
chats_module = chats.Chats(modules=models_module)
1629+
chat = chats_module.create(model='gemini-2.5-pro')
1630+
1631+
hop1_chunks = [
1632+
types.GenerateContentResponse(
1633+
candidates=[
1634+
types.Candidate(
1635+
content=types.Content(
1636+
role='model', parts=[types.Part(text='Hop 1 part. ')]
1637+
),
1638+
continuation_token=b'tok_hop_1',
1639+
)
1640+
]
1641+
),
1642+
types.GenerateContentResponse(
1643+
candidates=[
1644+
types.Candidate(
1645+
content=types.Content(
1646+
role='model', parts=[types.Part(text='End of hop 1. ')]
1647+
),
1648+
finish_reason=types.FinishReason.CONTINUATION,
1649+
)
1650+
]
1651+
),
1652+
]
1653+
hop2_cutoff_chunks = [
1654+
types.GenerateContentResponse(
1655+
candidates=[
1656+
types.Candidate(
1657+
content=types.Content(
1658+
role='model',
1659+
parts=[types.Part(text='Hop 2 cut off mid-stream')],
1660+
),
1661+
finish_reason=None,
1662+
)
1663+
]
1664+
)
1665+
]
1666+
1667+
with mock.patch.object(
1668+
models.Models,
1669+
'generate_content_stream',
1670+
side_effect=[iter(hop1_chunks), iter(hop2_cutoff_chunks)],
1671+
) as mock_stream:
1672+
chunks = list(chat.send_message_stream('Write a long story'))
1673+
assert mock_stream.call_count == 2
1674+
assert [c.text for c in chunks] == [
1675+
'Hop 1 part. ',
1676+
'End of hop 1. ',
1677+
'Hop 2 cut off mid-stream',
1678+
]
1679+
assert chat.get_history(curated=True) == []
1680+
1681+
1682+
@pytest.mark.asyncio
1683+
async def test_async_chat_send_message_stream_incomplete_continuation_not_recorded(
1684+
mock_api_client,
1685+
):
1686+
"""When an async continuation hop is cut off without finish_reason, the turn is excluded from curated history."""
1687+
models_module = models.AsyncModels(mock_api_client)
1688+
chats_module = chats.AsyncChats(modules=models_module)
1689+
chat = chats_module.create(model='gemini-2.5-pro')
1690+
1691+
hop1_chunks = [
1692+
types.GenerateContentResponse(
1693+
candidates=[
1694+
types.Candidate(
1695+
content=types.Content(
1696+
role='model', parts=[types.Part(text='Async Hop 1. ')]
1697+
),
1698+
continuation_token=b'async_tok_hop_1',
1699+
finish_reason=types.FinishReason.CONTINUATION,
1700+
)
1701+
]
1702+
)
1703+
]
1704+
hop2_cutoff_chunks = [
1705+
types.GenerateContentResponse(
1706+
candidates=[
1707+
types.Candidate(
1708+
content=types.Content(
1709+
role='model',
1710+
parts=[types.Part(text='Async Hop 2 cut off')],
1711+
),
1712+
finish_reason=None,
1713+
)
1714+
]
1715+
)
1716+
]
1717+
1718+
async def _make_async_stream(chunk_list):
1719+
for c in chunk_list:
1720+
yield c
1721+
1722+
with mock.patch.object(
1723+
models.AsyncModels,
1724+
'generate_content_stream',
1725+
new_callable=mock.AsyncMock,
1726+
side_effect=[
1727+
_make_async_stream(hop1_chunks),
1728+
_make_async_stream(hop2_cutoff_chunks),
1729+
],
1730+
) as mock_async_stream:
1731+
chunks = []
1732+
async for chunk in await chat.send_message_stream('Write a long story'):
1733+
chunks.append(chunk)
1734+
assert mock_async_stream.call_count == 2
1735+
assert [c.text for c in chunks] == ['Async Hop 1. ', 'Async Hop 2 cut off']
1736+
assert chat.get_history(curated=True) == []

0 commit comments

Comments
 (0)