diff --git a/python/restate/ext/pydantic/_agent.py b/python/restate/ext/pydantic/_agent.py index cc7888c..754bd4e 100644 --- a/python/restate/ext/pydantic/_agent.py +++ b/python/restate/ext/pydantic/_agent.py @@ -168,7 +168,13 @@ async def wrapped_event_stream_handler( async def single_event(): yield event - await context.run_typed("run event", lambda: fn(ctx, single_event())) + async def run_event() -> None: + # An async action so ctx.run() awaits the handler; a sync lambda + # returning the coroutine would never run it and the un-awaited + # coroutine fails JSON serialization of the journal entry. + await fn(ctx, single_event()) + + await context.run_typed("run event", run_event) @property def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]: diff --git a/python/restate/ext/pydantic/_model.py b/python/restate/ext/pydantic/_model.py index fd162da..b61ff5c 100644 --- a/python/restate/ext/pydantic/_model.py +++ b/python/restate/ext/pydantic/_model.py @@ -120,6 +120,11 @@ async def request_stream_run(): ) try: response = await context.run_typed("Model stream call", request_stream_run, self._options) + # Arm the turnstile that orders journaled tool executions, exactly + # like the non-streaming request() does — without this the first + # tool call of a streamed run fails with KeyError in wait_for(). + ids = [c.tool_call_id for c in response.tool_calls] + current_state().turnstile = Turnstile(ids) yield RestateStreamedResponse(model_request_parameters, response) except SdkInternalBaseException as e: raise Exception("Internal error during model stream call") from e