Reset chat app example

Author: DouweM

examples/pydantic_ai_examples/chat_app.py

@@ -7,97 +7,216 @@
 
 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, Response
-
-from pydantic_ai import Agent, RunContext
-from pydantic_ai.ui.vercel_ai import VercelAIAdapter
-
-from .sqlite_database import Database
+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,
+)
 
 # '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(sql_schema) as db:
+    async with Database.connect() as db:
         yield {'db': db}
 
 
-@dataclass
-class Deps:
-    conn: Database
-    user_id: int
-
-
-chat_agent = Agent(
-    'openai:gpt-4.1',
-    deps_type=Deps,
-    instructions="""
-You are a helpful assistant.
-
-Always reply with markdown. ALWAYS use code fences for code examples and lines of code.
-""",
-)
-
-
-@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 = fastapi.FastAPI(lifespan=lifespan)
+logfire.instrument_fastapi(app)
 
 
-@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])
+@app.get('/')
+async def index() -> FileResponse:
+    return FileResponse((THIS_DIR / 'chat_app.html'), media_type='text/html')
 
 
-app = fastapi.FastAPI(lifespan=lifespan)
-logfire.instrument_fastapi(app)
+@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')
 
 
 async def get_db(request: Request) -> Database:
     return request.state.db
 
 
-@app.options('/api/chat')
-def options_chat():
-    pass
+@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',
+    )
 
 
-@app.post('/api/chat')
-async def get_chat(request: Request, database: Database = Depends(get_db)) -> Response:
-    return await VercelAIAdapter[Deps].dispatch_request(
-        request, agent=chat_agent, deps=Deps(database, 123)
-    )
+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')
+
+
+@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
+        )
 
 
 if __name__ == '__main__':