Set Content-Type header on StreamingResponse

Author: DouweM

pydantic_ai_slim/pydantic_ai/ui/__init__.py

@@ -7,11 +7,12 @@
 from __future__ import annotations
 
 from .adapter import BaseAdapter, OnCompleteFunc, StateDeps, StateHandler
-from .event_stream import BaseEventStream
+from .event_stream import SSE_CONTENT_TYPE, BaseEventStream
 
 __all__ = [
     'BaseAdapter',
     'BaseEventStream',
+    'SSE_CONTENT_TYPE',
     'StateDeps',
     'StateHandler',
     'OnCompleteFunc',

pydantic_ai_slim/pydantic_ai/ui/adapter.py

@@ -128,10 +128,18 @@ def load_messages(cls, messages: Sequence[MessageT]) -> list[ModelMessage]:
         """Load messages from the request and return the loaded messages."""
         raise NotImplementedError
 
-    @property
     @abstractmethod
-    def event_stream(self) -> BaseEventStream[RunRequestT, EventT, AgentDepsT, OutputDataT]:
-        """Create an event stream for the adapter."""
+    def build_event_stream(
+        self, accept: str | None = None
+    ) -> BaseEventStream[RunRequestT, EventT, AgentDepsT, OutputDataT]:
+        """Create an event stream for the adapter.
+
+        Args:
+            accept: The accept header value.
+
+        Returns:
+            The event stream.
+        """
         raise NotImplementedError
 
     @cached_property
@@ -167,9 +175,9 @@ def encode_stream(self, stream: AsyncIterator[EventT], accept: str | None = None
 
         Args:
             stream: The stream of events to encode.
-            accept: The accept header value for encoding format.
+            accept: The accept header value.
         """
-        return self.event_stream.encode_stream(stream, accept)
+        return self.build_event_stream(accept).encode_stream(stream)
 
     async def process_stream(
         self,
@@ -182,7 +190,7 @@ async def process_stream(
             stream: The stream of events to process.
             on_complete: Optional callback function called when the agent run completes successfully.
         """
-        async for event in self.event_stream.handle_stream(stream, on_complete=on_complete):
+        async for event in self.build_event_stream().handle_stream(stream, on_complete=on_complete):
             yield event
 
     async def run_stream(
@@ -266,7 +274,7 @@ async def stream_response(self, stream: AsyncIterator[EventT], accept: str | Non
 
         Args:
             stream: The stream of events to encode.
-            accept: The accept header value for encoding format.
+            accept: The accept header value.
         """
         try:
             from starlette.responses import StreamingResponse
@@ -276,12 +284,11 @@ async def stream_response(self, stream: AsyncIterator[EventT], accept: str | Non
                 'you can use the `ui` optional group — `pip install "pydantic-ai-slim[ui]"`'
             ) from e
 
+        event_stream = self.build_event_stream(accept)
         return StreamingResponse(
-            self.encode_stream(
-                stream,
-                accept=accept,
-            ),
+            event_stream.encode_stream(stream),
             headers=self.response_headers,
+            media_type=event_stream.content_type,
         )
 
     @classmethod

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_adapter.py

@@ -101,10 +101,18 @@ async def validate_request(cls, request: Request) -> RunAgentInput:
         """Validate the request and return the validated request."""
         return RunAgentInput.model_validate(await request.json())
 
-    @property
-    def event_stream(self) -> BaseEventStream[RunAgentInput, BaseEvent, AgentDepsT, OutputDataT]:
-        """Create an event stream for the adapter."""
-        return AGUIEventStream(self.request)
+    def build_event_stream(
+        self, accept: str | None = None
+    ) -> BaseEventStream[RunAgentInput, BaseEvent, AgentDepsT, OutputDataT]:
+        """Create an event stream for the adapter.
+
+        Args:
+            accept: The accept header value.
+
+        Returns:
+            The event stream.
+        """
+        return AGUIEventStream(self.request, accept=accept)
 
     @cached_property
     def toolset(self) -> AbstractToolset[AgentDepsT] | None:

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_event_stream.py

@@ -26,7 +26,7 @@
 )
 from ...output import OutputDataT
 from ...tools import AgentDepsT
-from .. import BaseEventStream
+from .. import SSE_CONTENT_TYPE, BaseEventStream
 
 try:
     from ag_ui.core import (
@@ -64,10 +64,6 @@
     'RunFinishedEvent',
 ]
 
-SSE_CONTENT_TYPE: Final[str] = 'text/event-stream'
-"""Content type header value for Server-Sent Events (SSE)."""
-
-
 BUILTIN_TOOL_CALL_ID_PREFIX: Final[str] = 'pyd_ai_builtin'
 
 
@@ -79,18 +75,29 @@ class AGUIEventStream(BaseEventStream[RunAgentInput, BaseEvent, AgentDepsT, Outp
     _builtin_tool_call_ids: dict[str, str] = field(default_factory=dict)
     _error: bool = False
 
-    def encode_event(self, event: BaseEvent, accept: str | None = None) -> str:
+    @property
+    def _event_encoder(self) -> EventEncoder:
+        return EventEncoder(accept=self.accept or SSE_CONTENT_TYPE)
+
+    @property
+    def content_type(self) -> str:
+        """Get the content type for the event stream, compatible with the accept header value.
+
+        Args:
+            accept: The accept header value.
+        """
+        return self._event_encoder.get_content_type()
+
+    def encode_event(self, event: BaseEvent) -> str:
         """Encode an AG-UI event as SSE.
 
         Args:
             event: The AG-UI event to encode.
-            accept: The accept header value for encoding format.
 
         Returns:
             The SSE-formatted string.
         """
-        encoder = EventEncoder(accept=accept or SSE_CONTENT_TYPE)
-        return encoder.encode(event)
+        return self._event_encoder.encode(event)
 
     async def before_stream(self) -> AsyncIterator[BaseEvent]:
         """Yield events before agent streaming starts."""

pydantic_ai_slim/pydantic_ai/ui/event_stream.py

@@ -44,6 +44,9 @@
     'BaseEventStream',
 ]
 
+SSE_CONTENT_TYPE = 'text/event-stream'
+"""Content type header value for Server-Sent Events (SSE)."""
+
 EventT = TypeVar('EventT')
 """Type variable for protocol-specific event types."""
 
@@ -65,12 +68,15 @@ class BaseEventStream(ABC, Generic[RunRequestT, EventT, AgentDepsT, OutputDataT]
     """TODO (DouwM): Docstring."""
 
     request: RunRequestT
-    result: AgentRunResult[OutputDataT] | None = None
+
+    accept: str | None = None
+    """TODO (DouweM): Docstring"""
 
     message_id: str = field(default_factory=lambda: str(uuid4()))
 
     _turn: Literal['request', 'response'] | None = None
 
+    _result: AgentRunResult[OutputDataT] | None = None
     _final_result_event: FinalResultEvent | None = None
 
     def new_message_id(self) -> str:
@@ -82,25 +88,35 @@ def new_message_id(self) -> str:
         self.message_id = str(uuid4())
         return self.message_id
 
+    @property
+    def content_type(self) -> str:
+        """Get the content type for the event stream, compatible with the accept header value.
+
+        By default, this returns the SSE content type (`text/event-stream`).
+        If a subclass supports other types as well, it should consider `self.accept` in `encode_event` and return the resulting content type here.
+
+        Args:
+            accept: The accept header value.
+        """
+        return SSE_CONTENT_TYPE
+
     @abstractmethod
-    def encode_event(self, event: EventT, accept: str | None = None) -> str:
+    def encode_event(self, event: EventT) -> str:
         """Encode an event as a string.
 
         Args:
             event: The event to encode.
-            accept: The accept header value for encoding format.
         """
         raise NotImplementedError
 
-    async def encode_stream(self, stream: AsyncIterator[EventT], accept: str | None = None) -> AsyncIterator[str]:
+    async def encode_stream(self, stream: AsyncIterator[EventT]) -> AsyncIterator[str]:
         """Encode a stream of events as SSE strings.
 
         Args:
             stream: The stream of events to encode.
-            accept: The accept header value for encoding format.
         """
         async for event in stream:
-            yield self.encode_event(event, accept)
+            yield self.encode_event(event)
 
     async def handle_stream(  # noqa: C901
         self, stream: AsyncIterator[SourceEvent], on_complete: OnCompleteFunc[EventT] | None = None
@@ -147,19 +163,20 @@ async def handle_stream(  # noqa: C901
                         async for e in self.handle_function_tool_result(output_tool_result_event):
                             yield e
 
-                    self.result = cast(AgentRunResult[OutputDataT], event.result)
+                    result = cast(AgentRunResult[OutputDataT], event.result)
+                    self._result = result
 
                     async for e in self._turn_to(None):
                         yield e
 
                     if on_complete is not None:
                         if inspect.isasyncgenfunction(on_complete):
-                            async for e in on_complete(self.result):
+                            async for e in on_complete(result):
                                 yield e
                         elif _utils.is_async_callable(on_complete):
-                            await on_complete(self.result)
+                            await on_complete(result)
                         else:
-                            await _utils.run_in_executor(on_complete, self.result)
+                            await _utils.run_in_executor(on_complete, result)
                 elif isinstance(event, FinalResultEvent):
                     self._final_result_event = event
 

pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py

@@ -72,9 +72,10 @@ async def validate_request(cls, request: Request) -> RequestData:
         """Validate a Vercel AI request."""
         return request_data_ta.validate_json(await request.body())
 
-    @property
-    def event_stream(self) -> BaseEventStream[RequestData, BaseChunk, AgentDepsT, OutputDataT]:
-        return VercelAIEventStream(self.request)
+    def build_event_stream(
+        self, accept: str | None = None
+    ) -> BaseEventStream[RequestData, BaseChunk, AgentDepsT, OutputDataT]:
+        return VercelAIEventStream(self.request, accept=accept)
 
     @property
     def response_headers(self) -> Mapping[str, str] | None:

pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py

@@ -61,7 +61,7 @@ class VercelAIEventStream(BaseEventStream[RequestData, BaseChunk, AgentDepsT, Ou
 
     _step_started: bool = False
 
-    def encode_event(self, event: BaseChunk, accept: str | None = None) -> str:
+    def encode_event(self, event: BaseChunk) -> str:
         if isinstance(event, DoneChunk):
             return 'data: [DONE]\n\n'
         return f'data: {event.model_dump_json(by_alias=True, exclude_none=True)}\n\n'

tests/test_ui.py

@@ -92,9 +92,8 @@ async def validate_request(cls, request: Request) -> UIRequest:
     def load_messages(cls, messages: Sequence[ModelMessage]) -> list[ModelMessage]:
         return list(messages)
 
-    @property
-    def event_stream(self) -> UIEventStream[AgentDepsT, OutputDataT]:
-        return UIEventStream[AgentDepsT, OutputDataT](self.request)
+    def build_event_stream(self, accept: str | None = None) -> UIEventStream[AgentDepsT, OutputDataT]:
+        return UIEventStream[AgentDepsT, OutputDataT](self.request, accept=accept)
 
     @cached_property
     def messages(self) -> list[ModelMessage]:
@@ -115,7 +114,7 @@ def response_headers(self) -> dict[str, str]:
 
 @dataclass(kw_only=True)
 class UIEventStream(BaseEventStream[UIRequest, str, AgentDepsT, OutputDataT]):
-    def encode_event(self, event: str, accept: str | None = None) -> str:
+    def encode_event(self, event: str) -> str:
         return event
 
     async def handle_event(self, event: SourceEvent) -> AsyncIterator[str]:
@@ -629,7 +628,11 @@ async def send(data: MutableMapping[str, Any]) -> None:
 
     assert chunks == snapshot(
         [
-            {'type': 'http.response.start', 'status': 200, 'headers': [(b'x-test', b'test')]},
+            {
+                'type': 'http.response.start',
+                'status': 200,
+                'headers': [(b'x-test', b'test'), (b'content-type', b'text/event-stream; charset=utf-8')],
+            },
             {'type': 'http.response.body', 'body': b'<stream>', 'more_body': True},
             {'type': 'http.response.body', 'body': b'<response>', 'more_body': True},
             {'type': 'http.response.body', 'body': b'<text follows_text=False>', 'more_body': True},