refactoring

Author: DouweM

examples/pydantic_ai_examples/chat_app.py

@@ -16,7 +16,7 @@
 from fastapi import Depends, Request, Response
 
 from pydantic_ai import Agent, RunContext
-from pydantic_ai.vercel_ai_elements.starlette import StarletteChat
+from pydantic_ai.vercel_ai.starlette import StarletteChat
 
 from .sqlite_database import Database
 

pydantic_ai_slim/pydantic_ai/vercel_ai_elements/__init__.py renamed to pydantic_ai_slim/pydantic_ai/vercel_ai/__init__.py

@@ -249,15 +249,16 @@ class UIMessage(CamelBaseModel):
     """
 
 
-class SubmitMessage(CamelBaseModel):
+class SubmitMessage(CamelBaseModel, extra='allow'):
     """Submit a message to the agent."""
 
     trigger: Literal['submit-message']
     id: str
     messages: list[UIMessage]
 
-    model: str
-    web_search: bool
+    # TODO (DouweM): Update, make variable? I like `inference_params` from OpenAI ChatKit.
+    # model: str
+    # web_search: bool
 
 
 class RegenerateMessage(CamelBaseModel):
@@ -269,5 +270,5 @@ class RegenerateMessage(CamelBaseModel):
     message_id: str
 
 
-RequestData = SubmitMessage | RegenerateMessage
-request_data_schema: TypeAdapter[RequestData] = TypeAdapter(Annotated[RequestData, Discriminator('trigger')])
+RequestData = Annotated[SubmitMessage | RegenerateMessage, Discriminator('trigger')]
+request_data_ta: TypeAdapter[RequestData] = TypeAdapter(RequestData)

pydantic_ai_slim/pydantic_ai/vercel_ai_elements/_utils.py renamed to pydantic_ai_slim/pydantic_ai/vercel_ai/_utils.py

@@ -8,34 +8,12 @@
 from pydantic_core import to_json
 
 from .. import messages
-from ..agent import Agent
-from ..run import AgentRunResultEvent
-from ..tools import AgentDepsT
 from . import response_types as _t
 
-__all__ = 'sse_stream', 'VERCEL_AI_ELEMENTS_HEADERS', 'EventStreamer'
-# no idea if this is important, but vercel sends it, therefore so am I
-VERCEL_AI_ELEMENTS_HEADERS = {'x-vercel-ai-ui-message-stream': 'v1'}
+__all__ = 'VERCEL_AI_DSP_HEADERS', 'EventStreamer'
 
-
-async def sse_stream(agent: Agent[AgentDepsT], user_prompt: str, deps: Any) -> AsyncIterator[str]:
-    """Stream events from an agent run as Vercel AI Elements events.
-
-    Args:
-        agent: The agent to run.
-        user_prompt: The user prompt to run the agent with.
-        deps: The dependencies to pass to the agent.
-
-    Yields:
-        An async iterator text lines to stream over SSE.
-    """
-    event_streamer = EventStreamer()
-    async for event in agent.run_stream_events(user_prompt, deps=deps):
-        if not isinstance(event, AgentRunResultEvent):
-            async for chunk in event_streamer.event_to_chunks(event):
-                yield chunk.sse()
-    async for chunk in event_streamer.finish():
-        yield chunk.sse()
+# See https://ai-sdk.dev/docs/ai-sdk-ui/stream-protocol#data-stream-protocol
+VERCEL_AI_DSP_HEADERS = {'x-vercel-ai-ui-message-stream': 'v1'}
 
 
 @dataclass
@@ -136,6 +114,9 @@ def sse(self) -> str:
     def __str__(self) -> str:
         return 'DoneChunk<marker for the end of sse stream message>'
 
+    def __eq__(self, other: Any) -> bool:
+        return isinstance(other, DoneChunk)
+
 
 def _json_dumps(obj: Any) -> str:
     return to_json(obj).decode('utf-8')

pydantic_ai_slim/pydantic_ai/vercel_ai_elements/request_types.py renamed to pydantic_ai_slim/pydantic_ai/vercel_ai/request_types.py

@@ -0,0 +1,85 @@
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from typing import Generic
+
+from pydantic import ValidationError
+
+from ..agent import Agent
+from ..run import AgentRunResultEvent
+from ..tools import AgentDepsT
+from .request_types import RequestData, TextUIPart, request_data_ta
+from .response_stream import VERCEL_AI_DSP_HEADERS, DoneChunk, EventStreamer
+from .response_types import AbstractSSEChunk
+
+try:
+    from sse_starlette.sse import EventSourceResponse
+    from starlette.requests import Request
+    from starlette.responses import JSONResponse, Response
+except ImportError as e:
+    raise ImportError('To use Vercel AI Elements, please install starlette and sse_starlette') from e
+
+
+@dataclass
+class StarletteChat(Generic[AgentDepsT]):
+    """Starlette support for Pydantic AI's Vercel AI Elements integration.
+
+    This can be used with either FastAPI or Starlette apps.
+    """
+
+    agent: Agent[AgentDepsT]
+
+    async def dispatch_request(self, request: Request, deps: AgentDepsT) -> Response:
+        """Handle a request and return a streamed SSE response.
+
+        Args:
+            request: The incoming Starlette/FastAPI request.
+            deps: The dependencies for the agent.
+
+        Returns:
+            A streamed SSE response.
+        """
+        try:
+            data = request_data_ta.validate_json(await request.json())
+
+            async def run_sse() -> AsyncIterator[str]:
+                async for chunk in self.run(data, deps=deps):
+                    yield chunk.sse()
+
+            return EventSourceResponse(run_sse(), headers=VERCEL_AI_DSP_HEADERS)
+        except ValidationError as e:
+            return JSONResponse({'errors': e.errors()}, status_code=422)
+        except Exception as e:
+            return JSONResponse({'errors': str(e)}, status_code=500)
+
+    async def run(self, data: RequestData, deps: AgentDepsT = None) -> AsyncIterator[AbstractSSEChunk | DoneChunk]:
+        """Stream events from an agent run as Vercel AI Elements events.
+
+        Args:
+            data: The data to run the agent with.
+            deps: The dependencies to pass to the agent.
+
+        Yields:
+            An async iterator text lines to stream over SSE.
+        """
+        # TODO (DouweM): Use .model and .builtin_tools
+
+        # TODO: Use entire message history
+
+        if not data.messages:
+            raise ValueError('no messages provided')
+
+        message = data.messages[-1]
+        prompt: list[str] = []
+        for part in message.parts:
+            if isinstance(part, TextUIPart):
+                prompt.append(part.text)
+            else:
+                raise ValueError(f'Only text parts are supported yet, got {part}')
+
+        event_streamer = EventStreamer()
+        async for event in self.agent.run_stream_events('\n'.join(prompt), deps=deps):
+            if not isinstance(event, AgentRunResultEvent):
+                async for chunk in event_streamer.event_to_chunks(event):
+                    yield chunk
+        async for chunk in event_streamer.finish():
+            yield chunk