Support Vercel AI Data Stream Protocol

.gitignore

@@ -10,7 +10,7 @@ env*/
 /TODO.md
 /postgres-data/
 .DS_Store
-examples/pydantic_ai_examples/.chat_app_messages.sqlite
+.chat_app_messages.sqlite
 .cache/
 .vscode/
 /question_graph_history.json

examples/pydantic_ai_examples/chat_app.py

@@ -7,216 +7,96 @@
 
 from __future__ import annotations as _annotations
 
-import asyncio
-import json
-import sqlite3
-from collections.abc import AsyncIterator, Callable
-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, Literal, TypeVar
 
 import fastapi
 import logfire
-from fastapi import Depends, Request
-from fastapi.responses import FileResponse, Response, StreamingResponse
-from typing_extensions import LiteralString, ParamSpec, TypedDict
-
-from pydantic_ai import (
-    Agent,
-    ModelMessage,
-    ModelMessagesTypeAdapter,
-    ModelRequest,
-    ModelResponse,
-    TextPart,
-    UnexpectedModelBehavior,
-    UserPromptPart,
-)
+from fastapi import Depends, Request, Response
+
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.vercel_ai.starlette import StarletteChat
+
+from .sqlite_database import Database
 
 # 'if-token-present' means nothing will be sent (and the example will work) if you don't have logfire configured
 logfire.configure(send_to_logfire='if-token-present')
 logfire.instrument_pydantic_ai()
 
-agent = Agent('openai:gpt-4o')
 THIS_DIR = Path(__file__).parent
+sql_schema = """
+create table if not exists memory(
+    id integer primary key,
+    user_id integer not null,
+    value text not null,
+    unique(user_id, value)
+);"""
 
 
 @asynccontextmanager
 async def lifespan(_app: fastapi.FastAPI):
-    async with Database.connect() as db:
+    async with Database.connect(sql_schema) as db:
         yield {'db': db}
 
 
-app = fastapi.FastAPI(lifespan=lifespan)
-logfire.instrument_fastapi(app)
-
+@dataclass
+class Deps:
+    conn: Database
+    user_id: int
 
-@app.get('/')
-async def index() -> FileResponse:
-    return FileResponse((THIS_DIR / 'chat_app.html'), media_type='text/html')
 
+chat_agent = Agent(
+    'openai:gpt-4.1',
+    deps_type=Deps,
+    instructions="""
+You are a helpful assistant.
 
-@app.get('/chat_app.ts')
-async def main_ts() -> FileResponse:
-    """Get the raw typescript code, it's compiled in the browser, forgive me."""
-    return FileResponse((THIS_DIR / 'chat_app.ts'), media_type='text/plain')
+Always reply with markdown. ALWAYS use code fences for code examples and lines of code.
+""",
+)
 
 
-async def get_db(request: Request) -> Database:
-    return request.state.db
+@chat_agent.tool
+async def record_memory(ctx: RunContext[Deps], value: str) -> str:
+    """Use this tool to store information in memory."""
+    await ctx.deps.conn.execute(
+        'insert into memory(user_id, value) values(?, ?) on conflict do nothing',
+        ctx.deps.user_id,
+        value,
+        commit=True,
+    )
+    return 'Value added to memory.'
 
 
-@app.get('/chat/')
-async def get_chat(database: Database = Depends(get_db)) -> Response:
-    msgs = await database.get_messages()
-    return Response(
-        b'\n'.join(json.dumps(to_chat_message(m)).encode('utf-8') for m in msgs),
-        media_type='text/plain',
+@chat_agent.tool
+async def retrieve_memories(ctx: RunContext[Deps], memory_contains: str) -> str:
+    """Get all memories about the user."""
+    rows = await ctx.deps.conn.fetchall(
+        'select value from memory where user_id = ? and value like ?',
+        ctx.deps.user_id,
+        f'%{memory_contains}%',
     )
+    return '\n'.join([row[0] for row in rows])
 
 
-class ChatMessage(TypedDict):
-    """Format of messages sent to the browser."""
-
-    role: Literal['user', 'model']
-    timestamp: str
-    content: str
-
-
-def to_chat_message(m: ModelMessage) -> ChatMessage:
-    first_part = m.parts[0]
-    if isinstance(m, ModelRequest):
-        if isinstance(first_part, UserPromptPart):
-            assert isinstance(first_part.content, str)
-            return {
-                'role': 'user',
-                'timestamp': first_part.timestamp.isoformat(),
-                'content': first_part.content,
-            }
-    elif isinstance(m, ModelResponse):
-        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 (
-            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
-        async with agent.run_stream(prompt, message_history=messages) as result:
-            async for text in result.stream_output(debounce_by=0.01):
-                # text here is a `str` and the frontend wants
-                # JSON encoded ModelResponse, so we create one
-                m = ModelResponse(parts=[TextPart(text)], timestamp=result.timestamp())
-                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())
-
-    return StreamingResponse(stream_messages(), media_type='text/plain')
-
-
-P = ParamSpec('P')
-R = TypeVar('R')
+starlette_chat = StarletteChat(chat_agent)
+app = fastapi.FastAPI(lifespan=lifespan)
+logfire.instrument_fastapi(app)
 
 
-@dataclass
-class Database:
-    """Rudimentary database to store chat messages in SQLite.
-
-    The SQLite standard library package is synchronous, so we
-    use a thread pool executor to run queries asynchronously.
-    """
-
-    con: sqlite3.Connection
-    _loop: asyncio.AbstractEventLoop
-    _executor: ThreadPoolExecutor
-
-    @classmethod
-    @asynccontextmanager
-    async def connect(
-        cls, file: Path = THIS_DIR / '.chat_app_messages.sqlite'
-    ) -> AsyncIterator[Database]:
-        with logfire.span('connect to DB'):
-            loop = asyncio.get_event_loop()
-            executor = ThreadPoolExecutor(max_workers=1)
-            con = await loop.run_in_executor(executor, cls._connect, file)
-            slf = cls(con, loop, executor)
-        try:
-            yield slf
-        finally:
-            await slf._asyncify(con.close)
-
-    @staticmethod
-    def _connect(file: Path) -> sqlite3.Connection:
-        con = sqlite3.connect(str(file))
-        con = logfire.instrument_sqlite3(con)
-        cur = con.cursor()
-        cur.execute(
-            'CREATE TABLE IF NOT EXISTS messages (id INT PRIMARY KEY, message_list TEXT);'
-        )
-        con.commit()
-        return con
-
-    async def add_messages(self, messages: bytes):
-        await self._asyncify(
-            self._execute,
-            'INSERT INTO messages (message_list) VALUES (?);',
-            messages,
-            commit=True,
-        )
-        await self._asyncify(self.con.commit)
-
-    async def get_messages(self) -> list[ModelMessage]:
-        c = await self._asyncify(
-            self._execute, 'SELECT message_list FROM messages order by id'
-        )
-        rows = await self._asyncify(c.fetchall)
-        messages: list[ModelMessage] = []
-        for row in rows:
-            messages.extend(ModelMessagesTypeAdapter.validate_json(row[0]))
-        return messages
-
-    def _execute(
-        self, sql: LiteralString, *args: Any, commit: bool = False
-    ) -> sqlite3.Cursor:
-        cur = self.con.cursor()
-        cur.execute(sql, args)
-        if commit:
-            self.con.commit()
-        return cur
-
-    async def _asyncify(
-        self, func: Callable[P, R], *args: P.args, **kwargs: P.kwargs
-    ) -> R:
-        return await self._loop.run_in_executor(  # type: ignore
-            self._executor,
-            partial(func, **kwargs),
-            *args,  # type: ignore
-        )
+async def get_db(request: Request) -> Database:
+    return request.state.db
+
+
+@app.options('/api/chat')
+def options_chat():
+    pass
+
+
+@app.post('/api/chat')
+async def get_chat(request: Request, database: Database = Depends(get_db)) -> Response:
+    return await starlette_chat.dispatch_request(request, deps=Deps(database, 123))
 
 
 if __name__ == '__main__':

examples/pydantic_ai_examples/sqlite_database.py

@@ -0,0 +1,81 @@
+from __future__ import annotations as _annotations
+
+import asyncio
+import sqlite3
+from collections.abc import AsyncIterator, Callable
+from concurrent.futures.thread import ThreadPoolExecutor
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+from typing import Any, LiteralString, ParamSpec, TypeVar
+
+import logfire
+
+P = ParamSpec('P')
+R = TypeVar('R')
+
+
+@dataclass
+class Database:
+    """Rudimentary database to store chat messages in SQLite.
+
+    The SQLite standard library package is synchronous, so we
+    use a thread pool executor to run queries asynchronously.
+    """
+
+    con: sqlite3.Connection
+    _loop: asyncio.AbstractEventLoop
+    _executor: ThreadPoolExecutor
+
+    @classmethod
+    @asynccontextmanager
+    async def connect(
+        cls, schema_sql: str, file: Path = Path('.chat_app_messages.sqlite')
+    ) -> AsyncIterator[Database]:
+        with logfire.span('connect to DB'):
+            loop = asyncio.get_event_loop()
+            executor = ThreadPoolExecutor(max_workers=1)
+            con = await loop.run_in_executor(executor, cls._connect, schema_sql, file)
+            slf = cls(con, loop, executor)
+        try:
+            yield slf
+        finally:
+            await slf._asyncify(con.close)
+
+    @staticmethod
+    def _connect(schema_sql: str, file: Path) -> sqlite3.Connection:
+        con = sqlite3.connect(str(file))
+        con = logfire.instrument_sqlite3(con)
+        cur = con.cursor()
+        cur.execute(schema_sql)
+        con.commit()
+        return con
+
+    async def execute(self, sql: LiteralString, *args: Any, commit: bool = False):
+        await self._asyncify(self._execute, sql, *args, commit=True)
+        if commit:
+            await self._asyncify(self.con.commit)
+
+    async def fetchall(self, sql: LiteralString, *args: Any) -> list[tuple[str, ...]]:
+        c = await self._asyncify(self._execute, sql, *args)
+        rows = await self._asyncify(c.fetchall)
+        return [tuple(row) for row in rows]
+
+    def _execute(
+        self, sql: LiteralString, *args: Any, commit: bool = False
+    ) -> sqlite3.Cursor:
+        cur = self.con.cursor()
+        cur.execute(sql, args)
+        if commit:
+            self.con.commit()
+        return cur
+
+    async def _asyncify(
+        self, func: Callable[P, R], *args: P.args, **kwargs: P.kwargs
+    ) -> R:
+        return await self._loop.run_in_executor(  # type: ignore
+            self._executor,
+            partial(func, **kwargs),
+            *args,  # type: ignore
+        )