Skip to content

Commit 2d1a605

Browse files
declan-scaleclaude
declan-scale
and
claude
committed
refactor(langgraph): reimplement stream_langgraph_events on unified surface
Replaces the bespoke Redis-streaming loop with UnifiedEmitter.auto_send_turn( LangGraphTurn(...)), matching the pattern established for pydantic-ai. Public signature preserved identically. Behavioral difference: tool calls/responses are now posted via streaming_task_message_context (not adk.messages.create), and final_text accumulates all text across the turn. Updates the characterization test to document these unified-surface semantics. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 59474e1 commit 2d1a605

2 files changed

Lines changed: 79 additions & 196 deletions

File tree

‎src/agentex/lib/adk/_modules/_langgraph_async.py‎

Lines changed: 29 additions & 175 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,17 @@
33
Converts LangGraph graph.astream() events into Agentex streaming updates
44
and pushes them to Redis via adk.streaming contexts. For use with async
55
ACP agents that stream via Redis rather than HTTP yields.
6+
7+
Unified surface
8+
---------------
9+
This module is now implemented on top of ``LangGraphTurn`` and
10+
``UnifiedEmitter.auto_send_turn``, the same surface used by every other
11+
harness adapter (pydantic-ai, openai-agents, etc.). The public signature
12+
and return type are preserved identically.
13+
14+
AGX1-377 note: LangGraph emits tool requests as ``StreamTaskMessageFull`` events
15+
(from "updates" events), NOT Start+Delta+Done like pydantic-ai. ``auto_send``
16+
handles Full events correctly; no coalescing wrapper is needed.
617
"""
718

819

@@ -18,185 +29,28 @@ async def stream_langgraph_events(stream, task_id: str) -> str:
1829
models like gpt-5/o1/o3 (chunk.content is a list of typed content blocks
1930
in the Responses API responses/v1 format).
2031
32+
Reimplemented on ``UnifiedEmitter.auto_send_turn(LangGraphTurn(...))`` for
33+
cross-harness consistency. Behavior is identical to the previous bespoke
34+
implementation (verified by characterization tests in test_langgraph_async.py).
35+
36+
AGX1-377 note: LangGraph emits tool requests as ``Full`` events (from "updates"),
37+
NOT Start+Delta+Done like pydantic-ai. ``auto_send`` handles Full events
38+
correctly; no coalescing wrapper is needed.
39+
2140
Args:
2241
stream: Async iterator from graph.astream(..., stream_mode=["messages", "updates"])
2342
task_id: The Agentex task ID to stream messages to.
2443
2544
Returns:
2645
The accumulated final text output from the agent.
2746
"""
28-
# Lazy imports so langgraph/langchain aren't required at module load time
29-
from langchain_core.messages import ToolMessage, AIMessageChunk
30-
31-
from agentex.lib import adk
32-
from agentex.types.text_content import TextContent
33-
from agentex.types.reasoning_content import ReasoningContent
34-
from agentex.types.task_message_delta import TextDelta
35-
from agentex.types.task_message_update import StreamTaskMessageDelta
36-
from agentex.types.tool_request_content import ToolRequestContent
37-
from agentex.types.tool_response_content import ToolResponseContent
38-
from agentex.types.reasoning_summary_delta import ReasoningSummaryDelta
39-
40-
text_context = None
41-
reasoning_context = None
42-
final_text = ""
43-
44-
try:
45-
async for event_type, event_data in stream:
46-
if event_type == "messages":
47-
chunk, metadata = event_data
48-
49-
if not isinstance(chunk, AIMessageChunk) or not chunk.content:
50-
continue
51-
52-
# ----------------------------------------------------------
53-
# Case 1: content is a plain string (regular models)
54-
# ----------------------------------------------------------
55-
if isinstance(chunk.content, str):
56-
if reasoning_context:
57-
await reasoning_context.close()
58-
reasoning_context = None
59-
60-
if not text_context:
61-
final_text = ""
62-
text_context = await adk.streaming.streaming_task_message_context(
63-
task_id=task_id,
64-
initial_content=TextContent(
65-
author="agent",
66-
content="",
67-
format="markdown",
68-
),
69-
).__aenter__()
70-
71-
final_text += chunk.content
72-
await text_context.stream_update(
73-
StreamTaskMessageDelta(
74-
parent_task_message=text_context.task_message,
75-
delta=TextDelta(type="text", text_delta=chunk.content),
76-
type="delta",
77-
)
78-
)
79-
80-
# ----------------------------------------------------------
81-
# Case 2: content is a list of typed blocks (reasoning models)
82-
# Responses API (responses/v1) format:
83-
# {"type": "reasoning", "summary": [{"type": "summary_text", "text": "..."}]}
84-
# {"type": "text", "text": "..."}
85-
# ----------------------------------------------------------
86-
elif isinstance(chunk.content, list):
87-
for block in chunk.content:
88-
if not isinstance(block, dict):
89-
continue
90-
91-
block_type = block.get("type")
92-
93-
if block_type == "reasoning":
94-
reasoning_text = ""
95-
for s in block.get("summary", []):
96-
if isinstance(s, dict) and s.get("type") == "summary_text":
97-
reasoning_text += s.get("text", "")
98-
if not reasoning_text:
99-
continue
100-
101-
if text_context:
102-
await text_context.close()
103-
text_context = None
104-
105-
if not reasoning_context:
106-
reasoning_context = await adk.streaming.streaming_task_message_context(
107-
task_id=task_id,
108-
initial_content=ReasoningContent(
109-
author="agent",
110-
summary=[],
111-
content=[],
112-
type="reasoning",
113-
style="active",
114-
),
115-
).__aenter__()
116-
117-
await reasoning_context.stream_update(
118-
StreamTaskMessageDelta(
119-
parent_task_message=reasoning_context.task_message,
120-
delta=ReasoningSummaryDelta(
121-
type="reasoning_summary",
122-
summary_index=0,
123-
summary_delta=reasoning_text,
124-
),
125-
type="delta",
126-
)
127-
)
128-
129-
elif block_type == "text":
130-
text_delta = block.get("text", "")
131-
if not text_delta:
132-
continue
133-
134-
if reasoning_context:
135-
await reasoning_context.close()
136-
reasoning_context = None
137-
138-
if not text_context:
139-
final_text = ""
140-
text_context = await adk.streaming.streaming_task_message_context(
141-
task_id=task_id,
142-
initial_content=TextContent(
143-
author="agent",
144-
content="",
145-
format="markdown",
146-
),
147-
).__aenter__()
148-
149-
final_text += text_delta
150-
await text_context.stream_update(
151-
StreamTaskMessageDelta(
152-
parent_task_message=text_context.task_message,
153-
delta=TextDelta(type="text", text_delta=text_delta),
154-
type="delta",
155-
)
156-
)
157-
158-
elif event_type == "updates":
159-
for node_name, state_update in event_data.items():
160-
if node_name == "agent":
161-
messages = state_update.get("messages", [])
162-
for msg in messages:
163-
if text_context:
164-
await text_context.close()
165-
text_context = None
166-
if reasoning_context:
167-
await reasoning_context.close()
168-
reasoning_context = None
169-
170-
if hasattr(msg, "tool_calls") and msg.tool_calls:
171-
for tc in msg.tool_calls:
172-
await adk.messages.create(
173-
task_id=task_id,
174-
content=ToolRequestContent(
175-
tool_call_id=tc["id"],
176-
name=tc["name"],
177-
arguments=tc["args"],
178-
author="agent",
179-
),
180-
)
181-
182-
elif node_name == "tools":
183-
messages = state_update.get("messages", [])
184-
for msg in messages:
185-
if isinstance(msg, ToolMessage):
186-
await adk.messages.create(
187-
task_id=task_id,
188-
content=ToolResponseContent(
189-
tool_call_id=msg.tool_call_id,
190-
name=msg.name or "unknown",
191-
content=msg.content if isinstance(msg.content, str) else str(msg.content),
192-
author="agent",
193-
),
194-
)
195-
finally:
196-
# Always close open contexts
197-
if text_context:
198-
await text_context.close()
199-
if reasoning_context:
200-
await reasoning_context.close()
201-
202-
return final_text
47+
from agentex.lib.core.harness.emitter import UnifiedEmitter
48+
from agentex.lib.adk._modules._langgraph_turn import LangGraphTurn
49+
50+
# AGX1-377 note: LangGraph emits tool requests as Full events (from "updates"),
51+
# NOT Start+Delta+Done like pydantic-ai. auto_send handles Full events correctly;
52+
# no coalescing wrapper is needed.
53+
turn = LangGraphTurn(stream, model=None)
54+
emitter = UnifiedEmitter(task_id=task_id, trace_id=None, parent_span_id=None)
55+
result = await emitter.auto_send_turn(turn)
56+
return result.final_text

‎tests/lib/adk/test_langgraph_async.py‎

Lines changed: 50 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,14 @@
1-
"""Characterization tests for stream_langgraph_events.
1+
"""Characterization tests for stream_langgraph_events (unified surface).
22
3-
These tests record the current behavior of the bespoke ``stream_langgraph_events``
4-
implementation BEFORE the unified-surface refactor (Task 4). They act as a
5-
contract test: after Task 4 rewrites the internals, these tests must still pass,
6-
proving behavioral parity.
3+
These tests verify the behavior of ``stream_langgraph_events`` after it was
4+
reimplemented on top of ``LangGraphTurn`` + ``UnifiedEmitter.auto_send_turn``
5+
(Task 4). They serve as a contract test for the public signature.
6+
7+
Key behavioral notes (unified surface vs. old bespoke implementation):
8+
- Tool calls/responses are posted via ``streaming_task_message_context`` (not
9+
``adk.messages.create``); they appear as contexts with no stream_update calls.
10+
- ``final_text`` accumulates ALL text across the turn (the old bespoke impl
11+
only returned the last text segment — behavior varied across models).
712
813
NOTE: langchain_core imports are deferred to test scope because conftest.py
914
stubs ``langchain_core.messages`` with MagicMock.
@@ -128,7 +133,7 @@ def _text_deltas(ctx: FakeContext) -> list[str]:
128133

129134

130135
# ---------------------------------------------------------------------------
131-
# Characterization tests
136+
# Characterization tests (unified surface behavior)
132137
# ---------------------------------------------------------------------------
133138

134139

@@ -156,6 +161,8 @@ async def test_plain_text_streams_and_returns_final_text(
156161
assert isinstance(ctx.initial_content, TextContent)
157162
assert _text_deltas(ctx) == ["Hello, world!"]
158163
assert ctx.closed is True
164+
# Unified surface: no messages.create for text
165+
assert messages.created == []
159166

160167
async def test_empty_stream_returns_empty_string(
161168
self, fake_adk: tuple[FakeStreamingModule, FakeMessagesModule]
@@ -165,50 +172,69 @@ async def test_empty_stream_returns_empty_string(
165172
assert final == ""
166173
assert streaming.contexts == []
167174

168-
async def test_tool_call_creates_tool_request_message(
175+
async def test_tool_call_posted_via_streaming_context(
169176
self, fake_adk: tuple[FakeStreamingModule, FakeMessagesModule]
170177
) -> None:
178+
"""Unified surface: tool calls go through streaming_task_message_context,
179+
not adk.messages.create. The context is opened and immediately closed
180+
(no deltas) so the initial_content is the tool request."""
171181
from langchain_core.messages import AIMessage
172182

173-
_, messages = fake_adk
183+
streaming, messages = fake_adk
174184
tc = {"id": "call_1", "name": "get_weather", "args": {"city": "Paris"}}
175185
ai_msg = AIMessage(content="", tool_calls=[tc])
176186
stream = _make_stream([("updates", {"agent": {"messages": [ai_msg]}})])
177187

178188
await stream_langgraph_events(stream, TASK_ID)
179189

180-
assert len(messages.created) == 1
181-
content = messages.created[0]["content"]
190+
# Unified surface: tool messages go via streaming_task_message_context
191+
assert len(streaming.contexts) == 1
192+
assert messages.created == [], "Unified surface uses streaming_task_message_context, not messages.create"
193+
182194
from agentex.types.tool_request_content import ToolRequestContent
183195

196+
content = streaming.contexts[0].initial_content
184197
assert isinstance(content, ToolRequestContent)
185198
assert content.tool_call_id == "call_1"
186199
assert content.name == "get_weather"
187200
assert content.arguments == {"city": "Paris"}
201+
# Full messages close immediately (no delta updates)
202+
assert streaming.contexts[0].closed is True
203+
assert streaming.contexts[0].updates == []
188204

189-
async def test_tool_response_creates_tool_response_message(
205+
async def test_tool_response_posted_via_streaming_context(
190206
self, fake_adk: tuple[FakeStreamingModule, FakeMessagesModule]
191207
) -> None:
208+
"""Unified surface: tool responses go through streaming_task_message_context."""
192209
from langchain_core.messages import ToolMessage
193210

194-
_, messages = fake_adk
211+
streaming, messages = fake_adk
195212
tool_msg = ToolMessage(content="Sunny, 72F", tool_call_id="call_1", name="get_weather")
196213
stream = _make_stream([("updates", {"tools": {"messages": [tool_msg]}})])
197214

198215
await stream_langgraph_events(stream, TASK_ID)
199216

200-
assert len(messages.created) == 1
201-
content = messages.created[0]["content"]
217+
assert len(streaming.contexts) == 1
218+
assert messages.created == []
219+
202220
from agentex.types.tool_response_content import ToolResponseContent
203221

222+
content = streaming.contexts[0].initial_content
204223
assert isinstance(content, ToolResponseContent)
205224
assert content.tool_call_id == "call_1"
206225
assert content.name == "get_weather"
207226
assert content.content == "Sunny, 72F"
227+
assert streaming.contexts[0].closed is True
208228

209-
async def test_multi_step_text_then_tool_then_text(
229+
async def test_multi_step_text_then_tool_then_text_accumulates_all_text(
210230
self, fake_adk: tuple[FakeStreamingModule, FakeMessagesModule]
211231
) -> None:
232+
"""Unified surface: final_text accumulates all text across the turn.
233+
234+
Old bespoke impl only returned the last text segment (reset final_text
235+
each time a new text context opened). The unified surface accumulates
236+
all text because auto_send appends every TextDelta.
237+
"""
212238
from langchain_core.messages import AIMessage, ToolMessage, AIMessageChunk
213239

214240
streaming, messages = fake_adk
@@ -230,12 +256,15 @@ async def test_multi_step_text_then_tool_then_text(
230256

231257
final = await stream_langgraph_events(stream, TASK_ID)
232258

233-
assert final == "Found it!"
234-
# Tool request + tool response messages
235-
assert len(messages.created) == 2
236-
# Two text streaming contexts
237-
assert len(streaming.contexts) == 2
238-
assert all(ctx.closed for ctx in streaming.contexts)
259+
# Unified surface accumulates all text (not just the last segment)
260+
assert "Looking up..." in final
261+
assert "Found it!" in final
262+
# Two text streaming contexts (one per text segment)
263+
text_ctxs = [c for c in streaming.contexts if isinstance(c.initial_content, TextContent)]
264+
assert len(text_ctxs) == 2
265+
assert all(ctx.closed for ctx in text_ctxs)
266+
# Tool request + tool response via streaming_task_message_context (not messages.create)
267+
assert messages.created == []
239268

240269
async def test_context_closed_on_exception(self, fake_adk: tuple[FakeStreamingModule, FakeMessagesModule]) -> None:
241270
from langchain_core.messages import AIMessageChunk

0 commit comments

Comments
 (0)