Get chat app working with new messages format (#251)

Co-authored-by: Samuel Colvin <[email protected]>

Author: sydney-runkle

pydantic_ai_examples/chat_app.html

@@ -53,7 +53,7 @@ <h1>Chat App</h1>
       </div>
     </form>
     <div id="error" class="d-none text-danger">
-      Error occurred, check the console for more information.
+      Error occurred, check the browser developer console for more information.
     </div>
   </main>
 </body>

pydantic_ai_examples/chat_app.py

@@ -8,27 +8,31 @@
 from __future__ import annotations as _annotations
 
 import asyncio
+import json
 import sqlite3
 from collections.abc import AsyncIterator
 from concurrent.futures.thread import ThreadPoolExecutor
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
+from datetime import datetime, timezone
 from functools import partial
 from pathlib import Path
-from typing import Annotated, Any, Callable, TypeVar
+from typing import Annotated, Any, Callable, Literal, TypeVar
 
 import fastapi
 import logfire
 from fastapi import Depends, Request
 from fastapi.responses import HTMLResponse, Response, StreamingResponse
 from pydantic import Field, TypeAdapter
-from typing_extensions import LiteralString, ParamSpec
+from typing_extensions import LiteralString, ParamSpec, TypedDict
 
 from pydantic_ai import Agent
+from pydantic_ai.exceptions import UnexpectedModelBehavior
 from pydantic_ai.messages import (
     Message,
     MessagesTypeAdapter,
     ModelResponse,
+    TextPart,
     UserPrompt,
 )
 
@@ -68,19 +72,54 @@ async def get_db(request: Request) -> Database:
 async def get_chat(database: Database = Depends(get_db)) -> Response:
     msgs = await database.get_messages()
     return Response(
-        b'\n'.join(MessageTypeAdapter.dump_json(m) for m in msgs),
+        b'\n'.join(json.dumps(to_chat_message(m)).encode('utf-8') for m in msgs),
         media_type='text/plain',
     )
 
 
+class ChatMessage(TypedDict):
+    """Format of messages sent to the browser."""
+
+    role: Literal['user', 'model']
+    timestamp: str
+    content: str
+
+
+def to_chat_message(m: Message) -> ChatMessage:
+    if isinstance(m, UserPrompt):
+        return {
+            'role': 'user',
+            'timestamp': m.timestamp.isoformat(),
+            'content': m.content,
+        }
+    elif isinstance(m, ModelResponse):
+        first_part = m.parts[0]
+        if isinstance(first_part, TextPart):
+            return {
+                'role': 'model',
+                'timestamp': m.timestamp.isoformat(),
+                'content': first_part.content,
+            }
+    raise UnexpectedModelBehavior(f'Unexpected message type for chat app: {m}')
+
+
 @app.post('/chat/')
 async def post_chat(
     prompt: Annotated[str, fastapi.Form()], database: Database = Depends(get_db)
 ) -> StreamingResponse:
     async def stream_messages():
         """Streams new line delimited JSON `Message`s to the client."""
         # stream the user prompt so that can be displayed straight away
-        yield MessageTypeAdapter.dump_json(UserPrompt(content=prompt)) + b'\n'
+        yield (
+            json.dumps(
+                {
+                    'role': 'user',
+                    'timestamp': datetime.now(tz=timezone.utc).isoformat(),
+                    'content': prompt,
+                }
+            ).encode('utf-8')
+            + b'\n'
+        )
         # get the chat history so far to pass as context to the agent
         messages = await database.get_messages()
         # run the agent with the user prompt and the chat history
@@ -89,7 +128,7 @@ async def stream_messages():
                 # text here is a `str` and the frontend wants
                 # JSON encoded ModelResponse, so we create one
                 m = ModelResponse.from_text(content=text, timestamp=result.timestamp())
-                yield MessageTypeAdapter.dump_json(m) + b'\n'
+                yield json.dumps(to_chat_message(m)).encode('utf-8') + b'\n'
 
         # add new messages (e.g. the user prompt and the agent response in this case) to the database
         await database.add_messages(result.new_messages_json())

pydantic_ai_examples/chat_app.ts

@@ -36,7 +36,7 @@ async function onFetchResponse(response: Response): Promise<void> {
 // The format of messages, this matches pydantic-ai both for brevity and understanding
 // in production, you might not want to keep this format all the way to the frontend
 interface Message {
-  message_kind: string
+  role: string
   content: string
   timestamp: string
 }
@@ -50,14 +50,14 @@ function addMessages(responseText: string) {
   const messages: Message[] = lines.filter(line => line.length > 1).map(j => JSON.parse(j))
   for (const message of messages) {
     // we use the timestamp as a crude element id
-    const {timestamp, message_kind, content} = message
+    const {timestamp, role, content} = message
     const id = `msg-${timestamp}`
     let msgDiv = document.getElementById(id)
     if (!msgDiv) {
       msgDiv = document.createElement('div')
       msgDiv.id = id
-      msgDiv.title = `${message_kind} at ${timestamp}`
-      msgDiv.classList.add('border-top', 'pt-2', message_kind)
+      msgDiv.title = `${role} at ${timestamp}`
+      msgDiv.classList.add('border-top', 'pt-2', role)
       convElement.appendChild(msgDiv)
     }
     msgDiv.innerHTML = marked.parse(content)