Repository: pydantic/pydantic-ai
File: pydantic_ai_slim/pydantic_ai/ag_ui.py (Python)

Code before:

        async with agent.iter(
            user_prompt=None,
            output_type=[output_type or agent.output_type, DeferredToolCalls],
            message_history=messages,
            model=model,
            deps=deps,
            model_settings=model_settings,
            usage_limits=usage_limits,
            usage=usage,
            infer_name=infer_name,
            toolsets=toolsets,
        ) as run:
            completed_run = run
            async for event in _agent_stream(run):
                yield encoder.encode(event)
    except _RunError as e:
        yield encoder.encode(
            RunErrorEvent(message=e.message, code=e.code),
        )
    except Exception as e:
        yield encoder.encode(
            RunErrorEvent(message=str(e)),
        )
        raise e
    else:
        yield encoder.encode(
            RunFinishedEvent(
                thread_id=run_input.thread_id,
                run_id=run_input.run_id,
            ),
        )
    finally:
        # Call the callback if provided and we have a completed run
        if on_complete is not None and completed_run is not None:
            import inspect

            if inspect.iscoroutinefunction(on_complete):
                await on_complete(completed_run)
            else:
                on_complete(completed_run)


async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEvent]:
    """Run the agent streaming responses using AG-UI protocol events.

    Args:
        run: The agent run to process.

    Yields:
        AG-UI Server-Sent Events (SSE).

Review comment (bug):
If we're also going to call this on failure, I think we should pass the exception to the callback handler.

I'm also fine with just calling this after the `async for event` is complete.

Code after:
        ) as run:
            async for event in _agent_stream(run):
                yield encoder.encode(event)

        if on_complete is not None and run.result is not None:
            if _utils.is_async_callable(on_complete):
                await on_complete(run.result)
            else:
                await _utils.run_in_executor(on_complete, run.result)
    except _RunError as e:
        yield encoder.encode(
            RunErrorEvent(message=e.message, code=e.code),
        )
    except Exception as e:
        yield encoder.encode(
            RunErrorEvent(message=str(e)),
        )
        raise e
    else:
        yield encoder.encode(
            RunFinishedEvent(
                thread_id=run_input.thread_id,
                run_id=run_input.run_id,
            ),
        )


async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEvent]:
    """Run the agent streaming responses using AG-UI protocol events.

    Args:
        run: The agent run to process.

    Yields:
        AG-UI Server-Sent Events (SSE).
    """
    async for node in run:
        stream_ctx = _RequestStreamContext()
        if isinstance(node, ModelRequestNode):
            async with node.stream(run.ctx) as request_stream:
                async for agent_event in request_stream:
                    async for msg in _handle_model_request_event(stream_ctx, agent_event):
                        yield msg

                if stream_ctx.part_end:  # pragma: no branch
                    yield stream_ctx.part_end
                    stream_ctx.part_end = None
                if stream_ctx.thinking:
                    yield ThinkingEndEvent(
                        type=EventType.THINKING_END,
                    )