@@ -1116,10 +1116,16 @@ async def mock_coro():
11161116 )
11171117 ]
11181118
1119- # Should have only 1 response (no aggregated content generated)
1120- assert len (responses ) == 1
1121- # Verify it's a function call, not text
1119+ # With progressive SSE streaming enabled by default, we get 2 responses:
1120+ # 1. Partial response with function call
1121+ # 2. Final aggregated response with function call
1122+ assert len (responses ) == 2
1123+ # First response is partial
1124+ assert responses [0 ].partial is True
11221125 assert responses [0 ].content .parts [0 ].function_call is not None
1126+ # Second response is the final aggregated response
1127+ assert responses [1 ].partial is False
1128+ assert responses [1 ].content .parts [0 ].function_call is not None
11231129
11241130
11251131@pytest .mark .asyncio
@@ -1194,37 +1200,33 @@ async def mock_coro():
11941200 )
11951201 ]
11961202
1197- # Should have multiple responses:
1203+ # With progressive SSE streaming enabled, we get 4 responses:
11981204 # 1. Partial text "First text"
1199- # 2. Aggregated "First text" when function call interrupts
1200- # 3. Function call
1201- # 4. Partial text " second text"
1202- # 5. Final aggregated " second text"
1203- assert len (responses ) == 5
1205+ # 2. Partial function call
1206+ # 3. Partial text " second text"
1207+ # 4. Final aggregated response with all parts (text + FC + text)
1208+ assert len (responses ) == 4
12041209
12051210 # First partial text
12061211 assert responses [0 ].partial is True
12071212 assert responses [0 ].content .parts [0 ].text == "First text"
12081213
1209- # Aggregated first text (when function call interrupts)
1210- assert responses [1 ].content .parts [0 ].text == "First text"
1211- assert (
1212- responses [1 ].partial is None
1213- ) # Aggregated responses don't have partial flag
1214-
1215- # Function call
1216- assert responses [2 ].content .parts [0 ].function_call is not None
1217- assert responses [2 ].content .parts [0 ].function_call .name == "test_func"
1214+ # Partial function call
1215+ assert responses [1 ].partial is True
1216+ assert responses [1 ].content .parts [0 ].function_call is not None
1217+ assert responses [1 ].content .parts [0 ].function_call .name == "test_func"
12181218
1219- # Second partial text
1220- assert responses [3 ].partial is True
1221- assert responses [3 ].content .parts [0 ].text == " second text"
1219+ # Partial second text
1220+ assert responses [2 ].partial is True
1221+ assert responses [2 ].content .parts [0 ].text == " second text"
12221222
1223- # Final aggregated text with error info
1224- assert responses [4 ].content .parts [0 ].text == " second text"
1225- assert (
1226- responses [4 ].error_code is None
1227- ) # STOP finish reason should have None error_code
1223+ # Final aggregated response with all parts
1224+ assert responses [3 ].partial is False
1225+ assert len (responses [3 ].content .parts ) == 3
1226+ assert responses [3 ].content .parts [0 ].text == "First text"
1227+ assert responses [3 ].content .parts [1 ].function_call .name == "test_func"
1228+ assert responses [3 ].content .parts [2 ].text == " second text"
1229+ assert responses [3 ].error_code is None # STOP finish reason
12281230
12291231
12301232@pytest .mark .asyncio
@@ -1376,28 +1378,27 @@ async def mock_coro():
13761378 )
13771379 ]
13781380
1379- # Should properly separate thought and regular text across aggregations
1380- assert len (responses ) > 5 # Multiple partial + aggregated responses
1381+ # With progressive SSE streaming, we get 6 responses:
1382+ # 5 partial responses + 1 final aggregated response
1383+ assert len (responses ) == 6
13811384
1382- # Verify we get both thought and regular text parts in aggregated responses
1383- aggregated_responses = [
1384- r
1385- for r in responses
1386- if r .partial is None and r .content and len (r .content .parts ) > 1
1387- ]
1388- assert (
1389- len (aggregated_responses ) > 0
1390- ) # Should have at least one aggregated response with multiple parts
1385+ # All but the last should be partial
1386+ for i in range (5 ):
1387+ assert responses [i ].partial is True
13911388
1392- # Final aggregated response should have both thought and text
1389+ # Final aggregated response should have all parts
13931390 final_response = responses [- 1 ]
1394- assert (
1395- final_response .error_code is None
1396- ) # STOP finish reason should have None error_code
1397- assert len (final_response .content .parts ) == 2 # thought part + text part
1391+ assert final_response . partial is False
1392+ assert final_response .error_code is None # STOP finish reason
1393+ # Final response aggregates: thought + text + FC + thought + text
1394+ assert len (final_response .content .parts ) == 5
13981395 assert final_response .content .parts [0 ].thought is True
1399- assert "More thinking..." in final_response .content .parts [0 ].text
1400- assert final_response .content .parts [1 ].text == " and conclusion"
1396+ assert "Thinking..." in final_response .content .parts [0 ].text
1397+ assert final_response .content .parts [1 ].text == "Here's my answer"
1398+ assert final_response .content .parts [2 ].function_call .name == "lookup"
1399+ assert final_response .content .parts [3 ].thought is True
1400+ assert "More thinking..." in final_response .content .parts [3 ].text
1401+ assert final_response .content .parts [4 ].text == " and conclusion"
14011402
14021403
14031404@pytest .mark .asyncio
@@ -1491,44 +1492,23 @@ async def mock_coro():
14911492 )
14921493 ]
14931494
1494- # Find the aggregated text responses (non-partial, text-only)
1495- aggregated_text_responses = [
1496- r
1497- for r in responses
1498- if (
1499- r .partial is None
1500- and r .content
1501- and r .content .parts
1502- and r .content .parts [0 ].text
1503- and not r .content .parts [0 ].function_call
1504- )
1505- ]
1506-
1507- # Should have two separate text aggregations: "First chunk" and "Second chunk"
1508- assert len (aggregated_text_responses ) >= 2
1495+ # With progressive SSE streaming, we get 6 responses:
1496+ # 5 partial responses + 1 final aggregated response
1497+ assert len (responses ) == 6
15091498
1510- # First aggregation should contain "First chunk"
1511- first_aggregation = aggregated_text_responses [ 0 ]
1512- assert first_aggregation . content . parts [ 0 ]. text == "First chunk"
1499+ # All but the last should be partial
1500+ for i in range ( 5 ):
1501+ assert responses [ i ]. partial is True
15131502
1514- # Final aggregation should contain "Second chunk" and have error info
1515- final_aggregation = aggregated_text_responses [- 1 ]
1516- assert final_aggregation .content .parts [0 ].text == "Second chunk"
1517- assert (
1518- final_aggregation .error_code is None
1519- ) # STOP finish reason should have None error_code
1520-
1521- # Verify the function call is preserved between aggregations
1522- function_call_responses = [
1523- r
1524- for r in responses
1525- if (r .content and r .content .parts and r .content .parts [0 ].function_call )
1526- ]
1527- assert len (function_call_responses ) == 1
1528- assert (
1529- function_call_responses [0 ].content .parts [0 ].function_call .name
1530- == "divide"
1531- )
1503+ # Final response should be aggregated with all parts
1504+ final_response = responses [- 1 ]
1505+ assert final_response .partial is False
1506+ assert final_response .error_code is None # STOP finish reason
1507+ # Final response aggregates: text1 + text2 + FC + text3 + text4
1508+ assert len (final_response .content .parts ) == 3
1509+ assert final_response .content .parts [0 ].text == "First chunk"
1510+ assert final_response .content .parts [1 ].function_call .name == "divide"
1511+ assert final_response .content .parts [2 ].text == "Second chunk"
15321512
15331513
15341514@pytest .mark .asyncio
0 commit comments