wip: implement websocket chat service with Pydantic models and session management

app/schemas/chat.py

@@ -0,0 +1,124 @@
+"""Pydantic models for the websocket chat service.
+
+These models define the wire format exchanged over the ``/v1/chat/ws``
+websocket endpoint. The shapes intentionally mirror the conventions used by
+the Pydantic AI chat-app example: messages carry a ``role`` and ``content``,
+conversations are keyed by an opaque ``session_id``, and assistant replies
+can be streamed as a sequence of chunks before a final aggregated message.
+
+Python 3.14 evaluates annotations lazily by default (PEP 649/749), so plain
+modern typing (``str | None``, builtin generics, etc.) is used throughout
+without needing ``from __future__ import annotations``.
+"""
+
+from typing import Annotated, Literal
+from uuid import UUID, uuid4
+
+from pydantic import BaseModel, Field
+
+Role = Literal["system", "user", "assistant"]
+
+
+class ChatMessage(BaseModel):
+    """A single role-labeled message, matching Pydantic AI message shapes."""
+
+    role: Role
+    content: str
+
+
+# --------------------------------------------------------------------------
+# Client -> Server events
+# --------------------------------------------------------------------------
+
+
+class StartConversation(BaseModel):
+    """Sent by the client to (re)start a conversation, optionally seeding it
+    with a system prompt."""
+
+    type: Literal["start"] = "start"
+    system_prompt: str | None = None
+
+
+class SendUserMessage(BaseModel):
+    """Sent by the client with a new user message to append to history and
+    forward to the agent."""
+
+    type: Literal["user_message"] = "user_message"
+    content: str
+
+
+class RequestHistory(BaseModel):
+    """Sent by the client to request the full message history for the
+    current session."""
+
+    type: Literal["history_request"] = "history_request"
+
+
+class EndConversation(BaseModel):
+    """Sent by the client to gracefully close the conversation/websocket."""
+
+    type: Literal["end"] = "end"
+
+
+ClientEvent = Annotated[
+    StartConversation | SendUserMessage | RequestHistory | EndConversation,
+    Field(discriminator="type"),
+]
+
+
+# --------------------------------------------------------------------------
+# Server -> Client events
+# --------------------------------------------------------------------------
+
+
+class Connected(BaseModel):
+    """First event sent right after the websocket handshake completes."""
+
+    type: Literal["connected"] = "connected"
+    session_id: UUID = Field(default_factory=uuid4)
+
+
+class ConversationStarted(BaseModel):
+    """Acknowledges a ``start`` event."""
+
+    type: Literal["started"] = "started"
+    session_id: UUID
+
+
+class AssistantChunk(BaseModel):
+    """A single streamed token/chunk of the assistant's reply."""
+
+    type: Literal["assistant_chunk"] = "assistant_chunk"
+    index: int
+    content: str
+
+
+class AssistantMessage(BaseModel):
+    """The final, aggregated assistant message once streaming completes."""
+
+    type: Literal["assistant_message"] = "assistant_message"
+    message: ChatMessage
+
+
+class HistoryResponse(BaseModel):
+    """Response to a ``history_request`` event."""
+
+    type: Literal["history"] = "history"
+    messages: list[ChatMessage]
+
+
+class ChatError(BaseModel):
+    """Emitted whenever something goes wrong processing a client event."""
+
+    type: Literal["error"] = "error"
+    message: str
+
+
+ServerEvent = (
+    Connected
+    | ConversationStarted
+    | AssistantChunk
+    | AssistantMessage
+    | HistoryResponse
+    | ChatError
+)