fix(pydantic): make streamed runs work with auto_wrap_tools tool calls - #224
Open
QuocThuanTruong wants to merge 1 commit into
Open
fix(pydantic): make streamed runs work with auto_wrap_tools tool calls#224QuocThuanTruong wants to merge 1 commit into
QuocThuanTruong wants to merge 1 commit into
Conversation
Two bugs that stack on the first tool call of a streamed run
(RestateAgent with an event_stream_handler):
1. wrapped_event_stream_handler journaled each tool-side event via
run_typed("run event", lambda: fn(ctx, single_event())). The sync
lambda returns the async handler's coroutine un-awaited, so the
handler never runs and the journal entry fails with
"TypeError: Object of type coroutine is not JSON serializable".
Use an async action so ctx.run() awaits it.
2. request_stream() never armed the turnstile that orders journaled
tool executions — only the non-streaming request() did — so the
first tool call of a streamed run raised KeyError('call_...') in
Turnstile.wait_for. Arm it from the recorded response, mirroring
request().
Fixes restatedev#222
Fixes restatedev#223
|
All contributors have signed the CLA ✍️ ✅ |
Author
|
I have read the CLA Document and I hereby sign the CLA |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes the two bugs that break streamed runs with
auto_wrap_tools=Trueon the first tool call (RestateAgentwith anevent_stream_handler). Text-only streamed runs work today; the first tool call crashes the invocation.Fixes #222
Fixes #223
Bug 1 —
"run event"journal entry crashes withTypeError: Object of type coroutine is not JSON serializable(#222)wrapped_event_stream_handlerjournaled each tool-side event via:fn(the user'sevent_stream_handler) is async, but the sync lambda meanscreate_run_coroutinetakes the non-coroutine branch: the handler is never awaited (so it never runs) and the returned coroutine object is handed toserde.serialize, which raises. Replaced the lambda with an async action soctx.run()awaits it.Bug 2 —
KeyError('call_...')fromTurnstile.wait_foron streamed tool calls (#223)Tool executions in
RestateContextRunToolSet.call_toolare ordered by aTurnstilebuilt from the model response's tool-call ids — but only the non-streamingRestateModelWrapper.request()armed it.request_stream()left the default emptyTurnstile([]), so the first tool call of a streamed run failed the id lookup. Nowrequest_stream()arms it from the recorded response after the journaled"Model stream call"step, mirroringrequest()(deterministic on replay, since the response comes from the journal).Testing
Verified on a production-shaped app (FastAPI + Virtual Object per conversation, pydantic-ai 2.4.0, Python 3.12, Restate server 1.7.2):
TypeError(and, with bug 1 worked around,KeyError('call_...')); invocation retries until paused.Model stream call→Calling <tool>× N →Model stream call→ output) and the live event stream reaches the handler for every step.