Skip to content

Commit 0b1cff2

Browse files
xuanyang15copybara-github
authored andcommitted
feat: Enable PROGRESSIVE_SSE_STREAMING feature by default
Related: google#3705 Co-authored-by: Xuan Yang <xygoogle@google.com> PiperOrigin-RevId: 846793027
1 parent 71b3289 commit 0b1cff2

File tree

3 files changed

+61
-89
lines changed

3 files changed

+61
-89
lines changed

src/google/adk/features/_feature_registry.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ class FeatureConfig:
8989
FeatureStage.WIP, default_on=False
9090
),
9191
FeatureName.PROGRESSIVE_SSE_STREAMING: FeatureConfig(
92-
FeatureStage.WIP, default_on=False
92+
FeatureStage.EXPERIMENTAL, default_on=True
9393
),
9494
FeatureName.PUBSUB_TOOLSET: FeatureConfig(
9595
FeatureStage.EXPERIMENTAL, default_on=True

tests/unittests/flows/llm_flows/test_progressive_sse_streaming.py

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,6 @@
2626
from google.adk.runners import InMemoryRunner
2727
from google.adk.utils.streaming_utils import StreamingResponseAggregator
2828
from google.genai import types
29-
import pytest
30-
31-
32-
@pytest.fixture(autouse=True)
33-
def reset_env(monkeypatch):
34-
monkeypatch.setenv("ADK_ENABLE_PROGRESSIVE_SSE_STREAMING", "1")
35-
yield
36-
monkeypatch.delenv("ADK_ENABLE_PROGRESSIVE_SSE_STREAMING")
3729

3830

3931
def get_weather(location: str) -> dict[str, Any]:

tests/unittests/models/test_google_llm.py

Lines changed: 60 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)