From 566b9026deaa801a58df7e240789ae59cdce05e7 Mon Sep 17 00:00:00 2001 From: Selim Acerbas <91225118+selimacerbas@users.noreply.github.com> Date: Wed, 3 Jun 2026 23:46:10 +0200 Subject: [PATCH] fix(pydantic): use current_context() in RestateModelWrapper.request_stream request_stream() referenced `self._context`, which is never assigned anywhere in restate.ext.pydantic. The sibling request() and every other method use the current_context() extension point (added in #148) -- request_stream was simply missed in that migration. Any agent run that uses `event_stream_handler` (the only way restate.ext.pydantic surfaces per-turn model events) therefore raises AttributeError as soon as the streaming model path executes: File ".../restate/ext/pydantic/_model.py", line 116, in request_stream response = await self._context.run_typed("Model stream call", ...) File ".../pydantic_ai/models/wrapper.py", in __getattr__ return getattr(self.wrapped, item) AttributeError: 'OpenAIChatModel' object has no attribute '_context' (self._context -> WrapperModel.__getattr__ -> delegates to the wrapped model -> no such attribute.) Mirror request(): fetch current_context(), raise UserError when it is None, then use it for run_typed("Model stream call", ...). current_context and UserError are already imported in this module. --- python/restate/ext/pydantic/_model.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/restate/ext/pydantic/_model.py b/python/restate/ext/pydantic/_model.py index c20bdc5..074a09d 100644 --- a/python/restate/ext/pydantic/_model.py +++ b/python/restate/ext/pydantic/_model.py @@ -112,8 +112,13 @@ async def request_stream_run(): pass return streamed_response.get() + context = current_context() + if context is None: + raise UserError( + "A model cannot be used without a Restate context. Make sure to run it within an agent or a run context." + ) try: - response = await self._context.run_typed("Model stream call", request_stream_run, self._options) + response = await context.run_typed("Model stream call", request_stream_run, self._options) yield RestateStreamedResponse(model_request_parameters, response) except SdkInternalBaseException as e: raise Exception("Internal error during model stream call") from e