Update app.py

app.py

@@ -1,223 +1,219 @@
 from __future__ import annotations as _annotations
 
+import asyncio
 import json
-import os
+import sqlite3
+from collections.abc import AsyncIterator
+from concurrent.futures.thread import ThreadPoolExecutor
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from typing import Any
+from datetime import datetime, timezone
+from functools import partial
+from pathlib import Path
+from typing import Annotated, Any, Callable, 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
+from pydantic_ai.exceptions import UnexpectedModelBehavior
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelMessagesTypeAdapter,
+    ModelRequest,
+    ModelResponse,
+    TextPart,
+    UserPromptPart,
+)
 
-import gradio as gr
-from dotenv import load_dotenv
-from httpx import AsyncClient
-from pydantic_ai import Agent, ModelRetry, RunContext
-from pydantic_ai.messages import ModelStructuredResponse, ModelTextResponse, ToolReturn
+# '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')
 
-load_dotenv()
+agent = Agent('openai:gpt-4o')
+THIS_DIR = Path(__file__).parent
 
 
-@dataclass
-class Deps:
-    client: AsyncClient
-    weather_api_key: str | None
-    geo_api_key: str | None
+@asynccontextmanager
+async def lifespan(_app: fastapi.FastAPI):
+    async with Database.connect() as db:
+        yield {'db': db}
 
 
-weather_agent = Agent(
-    "openai:gpt-4o",
-    system_prompt="You are an expert packer. A user will ask you for help packing for a trip given a destination. Use your weather tools to provide a concise and effective packing list. Also ask follow up questions if neccessary.",
-    deps_type=Deps,
-    retries=2,
-)
+app = fastapi.FastAPI(lifespan=lifespan)
+logfire.instrument_fastapi(app)
 
 
-@weather_agent.tool
-async def get_lat_lng(
-    ctx: RunContext[Deps], location_description: str
-) -> dict[str, float]:
-    """Get the latitude and longitude of a location.
+@app.get('/')
+async def index() -> FileResponse:
+    return FileResponse((THIS_DIR / 'chat_app.html'), media_type='text/html')
 
-    Args:
-        ctx: The context.
-        location_description: A description of a location.
-    """
-    if ctx.deps.geo_api_key is None:
-        # if no API key is provided, return a dummy response (London)
-        return {"lat": 51.1, "lng": -0.1}
-
-    params = {
-        "q": location_description,
-        "api_key": ctx.deps.geo_api_key,
-    }
-    r = await ctx.deps.client.get("https://geocode.maps.co/search", params=params)
-    r.raise_for_status()
-    data = r.json()
-
-    if data:
-        return {"lat": data[0]["lat"], "lng": data[0]["lon"]}
-    else:
-        raise ModelRetry("Could not find the location")
-
-
-@weather_agent.tool
-async def get_weather(ctx: RunContext[Deps], lat: float, lng: float) -> dict[str, Any]:
-    """Get the weather at a location.
-
-    Args:
-        ctx: The context.
-        lat: Latitude of the location.
-        lng: Longitude of the location.
-    """
-    if ctx.deps.weather_api_key is None:
-        # if no API key is provided, return a dummy response
-        return {"temperature": "21 °C", "description": "Sunny"}
-
-    params = {
-        "apikey": ctx.deps.weather_api_key,
-        "location": f"{lat},{lng}",
-        "units": "metric",
-    }
-    r = await ctx.deps.client.get(
-        "https://api.tomorrow.io/v4/weather/realtime", params=params
-    )
-    r.raise_for_status()
-    data = r.json()
-
-    values = data["data"]["values"]
-    # https://docs.tomorrow.io/reference/data-layers-weather-codes
-    code_lookup = {
-        1000: "Clear, Sunny",
-        1100: "Mostly Clear",
-        1101: "Partly Cloudy",
-        1102: "Mostly Cloudy",
-        1001: "Cloudy",
-        2000: "Fog",
-        2100: "Light Fog",
-        4000: "Drizzle",
-        4001: "Rain",
-        4200: "Light Rain",
-        4201: "Heavy Rain",
-        5000: "Snow",
-        5001: "Flurries",
-        5100: "Light Snow",
-        5101: "Heavy Snow",
-        6000: "Freezing Drizzle",
-        6001: "Freezing Rain",
-        6200: "Light Freezing Rain",
-        6201: "Heavy Freezing Rain",
-        7000: "Ice Pellets",
-        7101: "Heavy Ice Pellets",
-        7102: "Light Ice Pellets",
-        8000: "Thunderstorm",
-    }
-    return {
-        "temperature": f'{values["temperatureApparent"]:0.0f}°C',
-        "description": code_lookup.get(values["weatherCode"], "Unknown"),
-    }
-
-
-TOOL_TO_DISPLAY_NAME = {"get_lat_lng": "Geocoding API", "get_weather": "Weather API"}
-
-client = AsyncClient()
-weather_api_key = os.getenv("WEATHER_API_KEY")
-# create a free API key at https://geocode.maps.co/
-geo_api_key = os.getenv("GEO_API_KEY")
-deps = Deps(client=client, weather_api_key=weather_api_key, geo_api_key=geo_api_key)
-
-
-async def stream_from_agent(prompt: str, chatbot: list[dict], past_messages: list):
-    chatbot.append({"role": "user", "content": prompt})
-    yield gr.Textbox(interactive=False, value=""), chatbot, gr.skip()
-    async with weather_agent.run_stream(
-        prompt, deps=deps, message_history=past_messages
-    ) as result:
-        for message in result.new_messages():
-            past_messages.append(message)
-            if isinstance(message, ModelStructuredResponse):
-                for call in message.calls:
-                    gr_message = {
-                        "role": "assistant",
-                        "content": "",
-                        "metadata": {
-                            "title": f"### 🛠️ Using {TOOL_TO_DISPLAY_NAME[call.tool_name]}",
-                            "id": call.tool_id,
-                        },
-                    }
-                    chatbot.append(gr_message)
-            if isinstance(message, ToolReturn):
-                for gr_message in chatbot:
-                    if gr_message.get("metadata", {}).get("id", "") == message.tool_id:
-                        gr_message["content"] = f"Output: {json.dumps(message.content)}"
-            yield gr.skip(), chatbot, gr.skip()
-        chatbot.append({"role": "assistant", "content": ""})
-        async for message in result.stream_text():
-            chatbot[-1]["content"] = message
-            yield gr.skip(), chatbot, gr.skip()
-        data = await result.get_data()
-        past_messages.append(ModelTextResponse(content=data))
-        yield gr.Textbox(interactive=True), gr.skip(), past_messages
-
-
-async def handle_retry(chatbot, past_messages: list, retry_data: gr.RetryData):
-    new_history = chatbot[: retry_data.index]
-    previous_prompt = chatbot[retry_data.index]["content"]
-    past_messages = past_messages[: retry_data.index]
-    async for update in stream_from_agent(previous_prompt, new_history, past_messages):
-        yield update
-
-
-def undo(chatbot, past_messages: list, undo_data: gr.UndoData):
-    new_history = chatbot[: undo_data.index]
-    past_messages = past_messages[: undo_data.index]
-    return chatbot[undo_data.index]["content"], new_history, past_messages
-
-
-def select_data(message: gr.SelectData) -> str:
-    return message.value["text"]
-
-
-with gr.Blocks() as demo:
-    gr.HTML(
-        """
-<div style="display: flex; justify-content: center; align-items: center; gap: 2rem; padding: 1rem; width: 100%">
-    <img src="https://ai.pydantic.dev/img/logo-white.svg" style="max-width: 200px; height: auto">
-    <div>
-        <h1 style="margin: 0 0 1rem 0">Vacation Packing Assistant</h1>
-        <h3 style="margin: 0 0 0.5rem 0">
-            This assistant will help you pack for your vacation. Enter your destination and it will provide you with a concise packing list based on the weather forecast.
-        </h3>
-        <h3 style="margin: 0">
-            Feel free to ask for help with any other questions you have about your trip!
-        </h3>
-    </div>
-</div>
-"""
-    )
-    past_messages = gr.State([])
-    chatbot = gr.Chatbot(
-        label="Packing Assistant",
-        type="messages",
-        avatar_images=(None, "https://ai.pydantic.dev/img/logo-white.svg"),
-        examples=[
-            {"text": "I am going to Paris for the holidays, what should I pack?"},
-            {"text": "I am going to Tokyo this week."},
-        ],
+
+@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.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',
     )
-    with gr.Row():
-        prompt = gr.Textbox(
-            lines=1,
-            show_label=False,
-            placeholder="I am planning a trip to Miami, what should I pack?",
+
+
+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):
+            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'
         )
-    generation = prompt.submit(
-        stream_from_agent,
-        inputs=[prompt, chatbot, past_messages],
-        outputs=[prompt, chatbot, past_messages],
-    )
-    chatbot.example_select(select_data, None, [prompt])
-    chatbot.retry(
-        handle_retry, [chatbot, past_messages], [prompt, chatbot, past_messages]
-    )
-    chatbot.undo(undo, [chatbot, past_messages], [prompt, chatbot, past_messages])
+        # 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(debounce_by=0.01):
+                # 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 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__':
+    #import uvicorn
 
-if __name__ == "__main__":
+    #uvicorn.run(
+    #    'pydantic_ai_examples.chat_app:app', reload=True, reload_dirs=[str(THIS_DIR)]
+    #)
     demo.launch()