Repository: pydantic/pydantic-ai
File: pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py (Python)

Code before:

    async def before_response(self) -> AsyncIterator[BaseChunk]:
        if self._step_started:
            yield FinishStepChunk()

        self._step_started = True
        yield StartStepChunk()

    async def after_stream(self) -> AsyncIterator[BaseChunk]:
        yield FinishStepChunk()

        yield FinishChunk(finish_reason=self._finish_reason)
        yield DoneChunk()

    async def handle_run_result(self, event: AgentRunResultEvent) -> AsyncIterator[BaseChunk]:
        messages = event.result.all_messages()
        if messages and isinstance(messages[-1], ModelResponse):
            pydantic_reason = messages[-1].finish_reason
            if pydantic_reason:
                self._finish_reason = _FINISH_REASON_MAP.get(pydantic_reason)

        # Emit tool approval requests for deferred approvals
        output = event.result.output
        if isinstance(output, DeferredToolRequests):
            for tool_call in output.approvals:
                yield ToolApprovalRequestChunk(
                    approval_id=str(uuid4()),
                    tool_call_id=tool_call.tool_call_id,
                )

    async def on_error(self, error: Exception) -> AsyncIterator[BaseChunk]:
        yield ErrorChunk(error_text=str(error))

    async def handle_text_start(self, part: TextPart, follows_text: bool = False) -> AsyncIterator[BaseChunk]:
        if follows_text:
            message_id = self.message_id
        else:
            message_id = self.new_message_id()
            yield TextStartChunk(id=message_id)

        if part.content:
            yield TextDeltaChunk(id=message_id, delta=part.content)

    async def handle_text_delta(self, delta: TextPartDelta) -> AsyncIterator[BaseChunk]:
        if delta.content_delta:  # pragma: no branch
            yield TextDeltaChunk(id=self.message_id, delta=delta.content_delta)

    async def handle_text_end(self, part: TextPart, followed_by_text: bool = False) -> AsyncIterator[BaseChunk]:
        if not followed_by_text:
            yield TextEndChunk(id=self.message_id)


Review comment (suggestion):
Do you know how Vercel AI will send back the approval/denial? We'd also have to map that to `DeferredToolResults` for the entire flow to work automatically. Let's do that in a separate PR (if you're up for it!) and limit this one again to the types + finish reason.

Code after:
            yield FinishStepChunk()

        self._step_started = True
        yield StartStepChunk()

    async def after_stream(self) -> AsyncIterator[BaseChunk]:
        yield FinishStepChunk()

        yield FinishChunk(finish_reason=self._finish_reason)
        yield DoneChunk()

    async def handle_run_result(self, event: AgentRunResultEvent) -> AsyncIterator[BaseChunk]:
        pydantic_reason = event.result.response.finish_reason
        if pydantic_reason:
            self._finish_reason = _FINISH_REASON_MAP.get(pydantic_reason)
        return
        yield

    async def on_error(self, error: Exception) -> AsyncIterator[BaseChunk]:
        self._finish_reason = 'error'
        yield ErrorChunk(error_text=str(error))

    async def handle_text_start(self, part: TextPart, follows_text: bool = False) -> AsyncIterator[BaseChunk]:
        if follows_text:
            message_id = self.message_id
        else:
            message_id = self.new_message_id()
            yield TextStartChunk(id=message_id)

        if part.content:
            yield TextDeltaChunk(id=message_id, delta=part.content)

    async def handle_text_delta(self, delta: TextPartDelta) -> AsyncIterator[BaseChunk]:
        if delta.content_delta:  # pragma: no branch
            yield TextDeltaChunk(id=self.message_id, delta=delta.content_delta)

    async def handle_text_end(self, part: TextPart, followed_by_text: bool = False) -> AsyncIterator[BaseChunk]:
        if not followed_by_text:
            yield TextEndChunk(id=self.message_id)

    async def handle_thinking_start(
        self, part: ThinkingPart, follows_thinking: bool = False
    ) -> AsyncIterator[BaseChunk]:
        message_id = self.new_message_id()
        yield ReasoningStartChunk(id=message_id)
        if part.content:
            yield ReasoningDeltaChunk(id=message_id, delta=part.content)

    async def handle_thinking_delta(self, delta: ThinkingPartDelta) -> AsyncIterator[BaseChunk]:
        if delta.content_delta:  # pragma: no branch
            yield ReasoningDeltaChunk(id=self.message_id, delta=delta.content_delta)