https://prefect.io logo
Join Slack
Powered by
<@ULVA73B9P> PydanticAI prefectagent seems to fail...
# ask-marvin
r
@Marvin PydanticAI prefectagent seems to fail when running outside of a flow deployment (aka just calling the prefec agent directly) when the tools have a non none timeout. It always throws a "timeout error" despite there being no real timeout (throws the timeout instantaneously). I can fix this by removing the timeout arg on the tools. can you investigate
m
✅ thought for 238.3 seconds
I dug into this, and your observation looks real: 
PrefectAgent
appears to mis-handle tool timeouts when invoked outside a Prefect flow context. What I found: * 
PrefectAgent
and related classes are coming from 
pydantic_ai.durable_exec.prefect
* The API does support 
tool_task_config
/ per-tool 
TaskConfig
* I did not find an existing Prefect GitHub issue or discussion matching this exact behavior * The likely failure mode is that timeout enforcement depends on Prefect task execution context, and when you call the agent “bare” instead of from within a 
@flow
, the timeout path is still activated but the surrounding runtime/context isn’t fully established So the symptom you described: * tool has 
timeout != None
* agent run is outside a flow/deployment * timeout error is raised immediately * removing timeout makes it work …is consistent with a context-sensitive bug/limitation in the integration. What I can say confidently: * Running the agent inside a Prefect flow is the intended usage pattern for this integration * Your workaround of removing 
timeout
is valid * This does not look like an actual tool-duration timeout if it fires instantaneously Recommended workaround:
Copy code
from prefect import flow

@flow
async def run_agent():
    result = await agent.run("...")
    return result
Then call the flow locally or deploy it. If you need to run it outside a flow for local/dev usage, the safest options are: * omit Prefect tool timeouts * enforce timeout inside the tool implementation yourself * or wrap the call in a lightweight 
@flow
even for local execution A couple relevant references: * Prefect docs home: Prefect docs * Example showing Prefect + PydanticAI usage in flow-oriented style: AI data analyst with PydanticAI If you want, I can help with one of these next steps: 1. draft a minimal reproducible example for a GitHub issue 2. help confirm whether this is in 
pydantic-ai
vs Prefect integration code 3. suggest a pattern for adding non-Prefect timeouts when running the agent outside flows
r
@Marvin if we call the flow, can we get access to the event stream handler?
m
✅ thought for 157.3 seconds
n
hmm lemme take a look at this
m
Yes — if you call the agent from inside a 
@flow
, you can use 
event_stream_handler
. A few key details: * 
PrefectAgent
accepts 
event_stream_handler
directly * Inside a flow, Prefect detects flow context and wraps event handling so it can run within Prefect’s task model * So you can still consume/log the stream while using 
agent.run(...)
inside the flow Example shape:
Copy code
python
from prefect import flow, get_run_logger
from pydantic_ai import Agent
from pydantic_ai.durable_exec.prefect import PrefectAgent

async def handle_events(ctx, stream):
    logger = get_run_logger()
    async for event in stream:
        <http://logger.info|logger.info>(f"event: {type(event).__name__}")

base_agent = Agent("openai:gpt-4o")
agent = PrefectAgent(
    base_agent,
    event_stream_handler=handle_events,
)

@flow
async def run_agent():
    result = await agent.run("hello")
    return result.output
A few caveats: * Inside flows, use 
agent.run(...)
with 
event_stream_handler
* Don’t expect 
run_stream()
/ 
run_stream_events()
to be the right pattern inside a Prefect flow * The handler is best for side effects: logging, writing events somewhere, emitting metrics, buffering to a queue, etc. * If you want to return streamed data to your caller, you’ll usually need to collect it yourself in some external structure/store as the handler runs Relevant docs/example: * AI data analyst with PydanticAI If you want, I can also show: * how to capture the streamed events into a list and return them from the flow, or * how to wire the handler to a websocket / SSE endpoint for UI streaming.
n
hm im not seeing a timeout bug here, but if you have a code example behaving unexpectedly, please do share. and yes i think marvin made the correct recommendation for event stream handler
r
Basically i was running the PrefectAgent in a fastapi endpoint directly via agent.run, but passing timeouts seemed to result in instantaneous timeout errors
@Marvin can you give me an example of how to wire the PrefectAgent handler and flow to a websocket / SSE endpoint for UI streaming.
m
✅ thought for 32.6 seconds
Yes — the basic pattern is: 1. create a per-request queue / pub-sub channel 2. have 
event_stream_handler
push agent events into that channel 3. run the 
@flow
in the background 4. have your WebSocket or SSE endpoint read from the channel and forward events to the UI Below is a minimal FastAPI example using an in-memory 
asyncio.Queue
. Important caveat: this works best when the flow is being run in-process with your API server. If you deploy the flow to a remote worker, your API process won’t share memory with the flow run, so you’d want Redis, Postgres, Kafka, etc. instead of an in-memory queue. WebSocket example
Copy code
python
import asyncio
import json
import uuid
from collections.abc import AsyncIterable
from typing import Any

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from prefect import flow, get_run_logger
from pydantic_ai import Agent, RunContext
from pydantic_ai.durable_exec.prefect import PrefectAgent
from pydantic_ai.messages import AgentStreamEvent

app = FastAPI()

streams: dict[str, asyncio.Queue] = {}


def serialize_event(event: AgentStreamEvent) -> dict[str, Any]:
    if hasattr(event, "model_dump"):
        payload = event.model_dump(mode="json")
    else:
        payload = {"repr": repr(event)}

    return {
        "type": type(event).__name__,
        "payload": payload,
    }


def build_agent(stream_id: str) -> PrefectAgent:
    async def handler(
        ctx: RunContext[Any],
        stream: AsyncIterable[AgentStreamEvent],
    ) -> None:
        logger = get_run_logger()
        queue = streams[stream_id]

        async for event in stream:
            message = serialize_event(event)
            <http://logger.info|logger.info>(f"Streaming event {message['type']}")
            await queue.put({
                "kind": "event",
                "data": message,
            })

    base_agent = Agent(
        "openai:gpt-4o",
        system_prompt="You are a helpful assistant.",
    )

    return PrefectAgent(
        base_agent,
        event_stream_handler=handler,
    )


@flow
async def run_agent_flow(prompt: str, stream_id: str) -> str:
    queue = streams[stream_id]
    agent = build_agent(stream_id)

    try:
        result = await agent.run(prompt)

        await queue.put({
            "kind": "result",
            "data": {
                "output": result.output,
            },
        })
        return result.output

    except Exception as exc:
        await queue.put({
            "kind": "error",
            "data": {
                "message": str(exc),
            },
        })
        raise

    finally:
        await queue.put({
            "kind": "done",
            "data": {},
        })


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()

    stream_id = str(uuid.uuid4())
    queue: asyncio.Queue = asyncio.Queue()
    streams[stream_id] = queue

    flow_task = None

    try:
        init_msg = await websocket.receive_json()
        prompt = init_msg["prompt"]

        await websocket.send_json({
            "kind": "session",
            "data": {
                "stream_id": stream_id,
            },
        })

        flow_task = asyncio.create_task(run_agent_flow(prompt, stream_id))

        while True:
            item = await queue.get()
            await websocket.send_json(item)

            if item["kind"] == "done":
                break

        await flow_task

    except WebSocketDisconnect:
        if flow_task:
            flow_task.cancel()

    except Exception as exc:
        await websocket.send_json({
            "kind": "error",
            "data": {
                "message": str(exc),
            },
        })

    finally:
        streams.pop(stream_id, None)
        await websocket.close()
Frontend messages look like:
Copy code
python
{
  "kind": "event",
  "data": {
    "type": "PartDeltaEvent",
    "payload": {...}
  }
}
and then finally: ```python { "kind": "result", "data": {
"output": "final answer here" } }
Copy code
*How the UI uses it*
* connect to `/ws`
* send `{ "prompt": "..." }`
* render `event` messages incrementally
* stop when `done` arrives

*SSE example*

If you prefer Server-Sent Events:

```python
import asyncio
import json
import uuid
from collections.abc import AsyncIterable, AsyncGenerator
from typing import Any

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from prefect import flow, get_run_logger
from pydantic_ai import Agent, RunContext
from pydantic_ai.durable_exec.prefect import PrefectAgent
from pydantic_ai.messages import AgentStreamEvent

app = FastAPI()

streams: dict[str, asyncio.Queue] = {}


def serialize_event(event: AgentStreamEvent) -> dict[str, Any]:
    if hasattr(event, "model_dump"):
        payload = event.model_dump(mode="json")
    else:
        payload = {"repr": repr(event)}

    return {
        "type": type(event).__name__,
        "payload": payload,
    }


def build_agent(stream_id: str) -> PrefectAgent:
    async def handler(
        ctx: RunContext[Any],
        stream: AsyncIterable[AgentStreamEvent],
    ) -> None:
        logger = get_run_logger()
        queue = streams[stream_id]

        async for event in stream:
            message = {
                "kind": "event",
                "data": serialize_event(event),
            }
            <http://logger.info|logger.info>(f"Streaming event {message['data']['type']}")
            await queue.put(message)

    base_agent = Agent(
        "openai:gpt-4o",
        system_prompt="You are a helpful assistant.",
    )

    return PrefectAgent(
        base_agent,
        event_stream_handler=handler,
    )


@flow
async def run_agent_flow(prompt: str, stream_id: str) -> str:
    queue = streams[stream_id]
    agent = build_agent(stream_id)

    try:
        result = await agent.run(prompt)

        await queue.put({
            "kind": "result",
            "data": {
                "output": result.output,
            },
        })
        return result.output

    except Exception as exc:
        await queue.put({
            "kind": "error",
            "data": {
                "message": str(exc),
            },
        })
        raise

    finally:
        await queue.put({
            "kind": "done",
            "data": {},
        })


@app.get("/sse")
async def sse(prompt: str):
    stream_id = str(uuid.uuid4())
    queue: asyncio.Queue = asyncio.Queue()
    streams[stream_id] = queue

    async def event_generator() -> AsyncGenerator[str, None]:
        flow_task = asyncio.create_task(run_agent_flow(prompt, stream_id))

        try:
            yield f"data: {json.dumps({'kind': 'session', 'data': {'stream_id': stream_id}})}\n\n"

            while True:
                item = await queue.get()
                yield f"data: {json.dumps(item)}\n\n"

                if item["kind"] == "done":
                    break

            await flow_task

        except asyncio.CancelledError:
            flow_task.cancel()
            raise

        finally:
            streams.pop(stream_id, None)

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
    )
Production note If your flow runs on a Prefect worker / separate process, replace 
asyncio.Queue
with shared infrastructure: * Redis pub/sub or Redis streams * Postgres table polled by the API * Kafka / NATS * a durable event bus Then your handler becomes:
Copy code
python
async def handler(ctx, stream):
    async for event in stream:
        await redis.xadd(
            f"agent-stream:{stream_id}",
            {"data": json.dumps(serialize_event(event))}
        )
and your WebSocket/SSE endpoint reads from Redis instead of local memory. Recommended message envelope I’d suggest standardizing on something like:
Copy code
python
{
  "kind": "event" | "result" | "error" | "done" | "session",
  "data": {...},
  "run_id": "...",
  "timestamp": "..."
}
That makes the frontend much easier to write. One more practical caveat The Prefect event handler runs inline with event processing, so keep it fast. If your websocket/SSE publishing path is slow, use a queue quickly in the handler and let another coroutine handle delivery. If you want, I can give you a more production-ready example next with: * FastAPI + Redis * reconnect/resume support * browser-side React example for consuming the stream
111 Views
Open in Slack
PreviousNext
Powered by