from __future__ import annotations from collections.abc import AsyncGenerator, Callable from contextlib import asynccontextmanager from typing import Any from prefect import task from prefect.context import FlowRunContext from pydantic_ai import ( ModelMessage, ModelResponse, ) from pydantic_ai._utils import fill_run_metadata from pydantic_ai.agent import EventStreamHandler from pydantic_ai.models import CompletedStreamedResponse, ModelRequestParameters, StreamedResponse from pydantic_ai.models.wrapper import WrapperModel from pydantic_ai.settings import ModelSettings from pydantic_ai.tools import RunContext from ._types import TaskConfig, default_task_config def _stamp_response_provenance(response: ModelResponse, messages: list[ModelMessage]) -> None: """Stamp the producing run's `run_id`/`conversation_id` on the response before Prefect persists it. The agent graph only fills these after the task returns, so without this the cached payload has them unset and a cache replay in a different conversation would be re-stamped as if it were produced there. Server-side state guards (e.g. OpenAI `openai_conversation_id='auto'`) rely on a replayed response keeping its original `conversation_id` to avoid continuing another conversation's provider-side state. """ if messages: # pragma: no branch final_request = messages[-1] fill_run_metadata(response, run_id=final_request.run_id, conversation_id=final_request.conversation_id) class PrefectModel(WrapperModel): """A wrapper for Model that integrates with Prefect, turning request and request_stream into Prefect tasks.""" def __init__( self, model: Any, *, task_config: TaskConfig, get_event_stream_handler: Callable[[], EventStreamHandler[Any] | None], ): super().__init__(model) self.task_config = default_task_config | (task_config or {}) # Resolve the effective event stream handler lazily inside the task so that a per-run # handler (set on a `ContextVar` by `PrefectAgent`) is picked up without rebuilding the model # and re-registering its Prefect tasks. self._get_event_stream_handler = get_event_stream_handler @task async def wrapped_request( messages: list[ModelMessage], model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters, ) -> ModelResponse: response = await super(PrefectModel, self).request(messages, model_settings, model_request_parameters) _stamp_response_provenance(response, messages) return response self._wrapped_request = wrapped_request @task async def request_stream_task( messages: list[ModelMessage], model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters, ctx: RunContext[Any] | None, ) -> ModelResponse: event_stream_handler = self._get_event_stream_handler() async with super(PrefectModel, self).request_stream( messages, model_settings, model_request_parameters, ctx ) as streamed_response: if event_stream_handler is not None: assert ctx is not None, ( 'A Prefect model cannot be used with `pydantic_ai.direct.model_request_stream()` as it requires a `run_context`. ' 'Set an `event_stream_handler` on the agent and use `agent.run()` instead.' ) await event_stream_handler(ctx, streamed_response) # Consume the entire stream async for _ in streamed_response: pass response = streamed_response.get() _stamp_response_provenance(response, messages) return response self._wrapped_request_stream = request_stream_task @task async def cancel_suspended_response_task(response: ModelResponse) -> None: await super(PrefectModel, self).cancel_suspended_response(response) self._wrapped_cancel_suspended_response = cancel_suspended_response_task async def request( self, messages: list[ModelMessage], model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters, ) -> ModelResponse: """Make a model request, wrapped as a Prefect task when in a flow.""" return await self._wrapped_request.with_options( name=f'Model Request: {self.wrapped.model_name}', **self.task_config )(messages, model_settings, model_request_parameters) async def cancel_suspended_response(self, response: ModelResponse) -> None: """Cancel a server-side suspended/background response, wrapped as a Prefect task. The teardown performs a raw HTTP call to the provider, so it runs as a task (durable, retried) rather than inline in the flow. """ await self._wrapped_cancel_suspended_response.with_options( name=f'Model Cancel Suspended Response: {self.wrapped.model_name}', **self.task_config )(response) @asynccontextmanager async def request_stream( self, messages: list[ModelMessage], model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters, run_context: RunContext[Any] | None = None, ) -> AsyncGenerator[StreamedResponse]: """Make a streaming model request. When inside a Prefect flow, the stream is consumed within a task and a non-streaming response is returned. When not in a flow, behaves normally. """ # Check if we're in a flow context flow_run_context = FlowRunContext.get() # If not in a flow, just call the wrapped request_stream method if flow_run_context is None: async with super().request_stream( messages, model_settings, model_request_parameters, run_context ) as streamed_response: yield streamed_response return # If in a flow, consume the stream in a task and return the final response response = await self._wrapped_request_stream.with_options( name=f'Model Request (Streaming): {self.wrapped.model_name}', **self.task_config )(messages, model_settings, model_request_parameters, run_context) # Without an `event_stream_handler`, the task drained and discarded the real stream's events # (e.g. `agent.iter` inside a flow, where the caller drives the flow-side stream via # `node.stream(...)`/`stream_text()`). Replay the response's parts as events so that stream # produces content. With a handler, events were already delivered inside the task, so the # flow-side stream stays empty to avoid delivering them twice. yield CompletedStreamedResponse( response, model_request_parameters=model_request_parameters, replay_events=self._get_event_stream_handler() is None, )