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

.env

@@ -22,4 +22,6 @@ EMAIL_HOST=
 EMAIL_HOST_USER=
 EMAIL_HOST_PASSWORD=
 
+CHAT_BACKEND=ollama
+
 

app/config.py

@@ -1,4 +1,5 @@
 import os
+from typing import Literal
 
 from pydantic import BaseModel, PostgresDsn, RedisDsn, computed_field
 from pydantic_core import MultiHostUrl
@@ -13,6 +14,25 @@ class SMTPConfig(BaseModel):
     template_path: str = os.getenv("EMAIL_TEMPLATE_PATH", "templates")
 
 
+class ChatConfig(BaseModel):
+    """Configuration for the websocket chat service's model-client adapter.
+
+    ``backend`` selects which :class:`~app.services.chat_agent.ChatAgent`
+    implementation is built by
+    :func:`~app.services.chat_agent.build_chat_agent` during app startup:
+
+    - ``"stub"``: a local, dependency-free echo agent (default, no network
+      calls, ideal for local dev/tests).
+    - ``"ollama"``: streams completions from an OpenAI-compatible endpoint
+      (e.g. a local Ollama server), reusing the same adapter interface.
+    """
+
+    backend: Literal["stub", "ollama"] = os.getenv("CHAT_BACKEND", "stub")
+    base_url: str = os.getenv("CHAT_BASE_URL", "http://localhost:11434/v1")
+    model: str = os.getenv("CHAT_MODEL", "llama3.2")
+    stream_delay_seconds: float = float(os.getenv("CHAT_STREAM_DELAY", "0.02"))
+
+
 class Settings(BaseSettings):
     model_config = SettingsConfigDict(
         env_file=".env", env_ignore_empty=True, extra="ignore"
@@ -21,6 +41,7 @@ class Settings(BaseSettings):
     jwt_expire: int = os.getenv("JWT_EXPIRE")
 
     smtp: SMTPConfig = SMTPConfig()
+    chat: ChatConfig = ChatConfig()
 
     REDIS_HOST: str
     REDIS_PORT: int

app/main.py

@@ -1,4 +1,3 @@
-from starlette.templating import _TemplateResponse
 from contextlib import asynccontextmanager
 from pathlib import Path
 
@@ -9,7 +8,9 @@ from fastapi.templating import Jinja2Templates
 from rotoger import get_logger
 from starlette.middleware import Middleware
 from starlette.middleware.gzip import GZipMiddleware
+from starlette.templating import _TemplateResponse
 
+from app.api.chat import router as chat_ws_router
 from app.api.health import router as health_router
 from app.api.ml import router as ml_router
 from app.api.nonsense import router as nonsense_router
@@ -21,6 +22,8 @@ from app.exception_handlers import register_exception_handlers
 from app.middleware.profiler import ProfilingMiddleware
 from app.redis import get_redis
 from app.services.auth import AuthBearer
+from app.services.chat_agent import build_chat_agent
+from app.services.chat_session import ChatSessionManager
 
 templates = Jinja2Templates(directory=Path(__file__).parent.parent / "templates")
 
@@ -30,21 +33,27 @@ async def lifespan(app: FastAPI):
     app.logger = get_logger()
     app.redis = await get_redis()
     postgres_dsn = global_settings.postgres_url.unicode_string()
+    # Chat service: initialize the pluggable model-client adapter and the
+    # in-memory session manager. See app/services/chat_agent.py to swap the
+    # local stub for a real model client (OpenAI, Ollama, etc.).
+    app.chat_agent = build_chat_agent(global_settings.chat)
+    app.chat_sessions = ChatSessionManager()
     try:
-        app.postgres_pool = await asyncpg.create_pool(
-            dsn=postgres_dsn,
-            min_size=5,
-            max_size=20,
-        )
-        await app.logger.ainfo(
-            "Postgres pool created", idle_size=app.postgres_pool.get_idle_size()
-        )
+        # app.postgres_pool = await asyncpg.create_pool(
+        #     dsn=postgres_dsn,
+        #     min_size=5,
+        #     max_size=20,
+        # )
+        # await app.logger.ainfo(
+        #     "Postgres pool created", idle_size=app.postgres_pool.get_idle_size()
+        # )
         yield
     except Exception as e:
         await app.logger.aerror("Error during app startup", error=repr(e))
         raise
     finally:
         await app.redis.close()
+        await app.chat_agent.aclose()
         # await app.postgres_pool.close()
 
 
@@ -66,6 +75,7 @@ def create_app() -> FastAPI:
     app.include_router(shakespeare_router)
     app.include_router(user_router)
     app.include_router(ml_router, prefix="/v1/ml", tags=["ML"])
+    app.include_router(chat_ws_router, prefix="/v1/chat", tags=["Chat"])
     app.include_router(
         health_router, prefix="/v1/public/health", tags=["Health, Public"]
     )

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
+)

app/services/chat_agent.py

@@ -0,0 +1,116 @@
+"""Pluggable async model-client adapter for the websocket chat service.
+
+``ChatAgent`` is the small interface every model connector must satisfy:
+given the full message history (system/user/assistant, matching Pydantic AI
+message-history semantics), yield the assistant's reply as a stream of text
+chunks. Swapping the local stub for a real model (OpenAI, a local Ollama
+server, etc.) only requires implementing this protocol and pointing
+``build_chat_agent`` at it - no changes to the websocket endpoint or session
+handling are needed.
+"""
+
+import asyncio
+import random
+from collections.abc import AsyncIterator
+from typing import Protocol, runtime_checkable
+
+import httpx
+import orjson
+
+from app.config import ChatConfig
+from app.schemas.chat import ChatMessage
+
+
+@runtime_checkable
+class ChatAgent(Protocol):
+    """Adapter interface implemented by every model connector."""
+
+    async def stream_reply(self, messages: list[ChatMessage]) -> AsyncIterator[str]:
+        """Yield the assistant reply for ``messages`` chunk by chunk.
+
+        ``messages`` is the full conversation history (oldest first),
+        following Pydantic AI's role-labeled message-history convention.
+        """
+        ...  # pragma: no cover - protocol stub, never called directly
+
+    async def aclose(self) -> None:
+        """Release any held resources (connections, clients, ...)."""
+
+
+class LocalEchoAgent:
+    """Dependency-free stub agent used for local development and tests.
+
+    It requires no API keys or network access: it "thinks" briefly, then
+    streams back a canned/echo response word by word, emulating the token
+    streaming behaviour of a real LLM backend closely enough to exercise the
+    full websocket flow end-to-end.
+    """
+
+    def __init__(self, stream_delay_seconds: float = 0.02) -> None:
+        self.stream_delay_seconds = stream_delay_seconds
+
+    def _compose_reply(self, messages: list[ChatMessage]) -> str:
+        last_user = next(
+            (m.content for m in reversed(messages) if m.role == "user"), ""
+        )
+        if not last_user:
+            return "Hello! I'm a local stub agent. Send me a message to get started."
+        greetings = ("hi", "hello", "hey")
+        if last_user.strip().lower() in greetings:
+            return "Hello there! How can I help you today?"
+        return f"You said: {last_user!r}. This is a local echo response (stub agent)."
+
+    async def stream_reply(self, messages: list[ChatMessage]) -> AsyncIterator[str]:
+        reply = self._compose_reply(messages)
+        for word in reply.split(" "):
+            await asyncio.sleep(self.stream_delay_seconds + random.uniform(0, 0.01))
+            yield word + " "
+
+    async def aclose(self) -> None:
+        return None
+
+
+class OllamaChatAgent:
+    """Streams chat completions from an OpenAI-compatible endpoint.
+
+    Works out of the box with a local Ollama server (``ollama serve``) but
+    any OpenAI-compatible ``/chat/completions`` endpoint works too. This is
+    a ready-to-swap-in replacement for :class:`LocalEchoAgent` once a real
+    model should be used.
+    """
+
+    def __init__(self, base_url: str, model: str) -> None:
+        self.model = model
+        self._client = httpx.AsyncClient(base_url=base_url, timeout=60.0)
+
+    async def stream_reply(self, messages: list[ChatMessage]) -> AsyncIterator[str]:
+        payload = {
+            "model": self.model,
+            "messages": [{"role": m.role, "content": m.content} for m in messages],
+            "stream": True,
+        }
+        async with self._client.stream(
+            "POST", "/chat/completions", json=payload
+        ) as response:
+            async for line in response.aiter_lines():
+                if not line.startswith("data: ") or line == "data: [DONE]":
+                    continue
+                try:
+                    data = orjson.loads(line[6:])
+                    content = (
+                        data.get("choices", [{}])[0].get("delta", {}).get("content", "")
+                    )
+                except Exception:
+                    content = ""
+                if content:
+                    yield content
+
+    async def aclose(self) -> None:
+        await self._client.aclose()
+
+
+def build_chat_agent(config: ChatConfig) -> ChatAgent:
+    """Factory selecting the concrete :class:`ChatAgent` from ``config``."""
+    if config.backend == "ollama":
+        return OllamaChatAgent(base_url=config.base_url, model=config.model)
+    return LocalEchoAgent(stream_delay_seconds=config.stream_delay_seconds)

app/services/chat_session.py

@@ -0,0 +1,52 @@
+"""In-memory conversation/session management for the websocket chat service.
+
+Sessions are intentionally kept simple (a dict guarded by an ``asyncio.Lock``)
+since each websocket connection owns exactly one session for its lifetime.
+Swapping this for a Redis-backed store later (for multi-worker deployments)
+only requires changing this module; the websocket endpoint only depends on
+the small public API below.
+"""
+
+import asyncio
+from dataclasses import dataclass, field
+from uuid import UUID, uuid4
+
+from app.schemas.chat import ChatMessage
+
+
+@dataclass(slots=True)
+class ChatSession:
+    id: UUID = field(default_factory=uuid4)
+    messages: list[ChatMessage] = field(default_factory=list)
+
+    def add(self, message: ChatMessage) -> None:
+        self.messages.append(message)
+
+    def history(self) -> list[ChatMessage]:
+        return list(self.messages)
+
+
+class ChatSessionManager:
+    """Tracks active chat sessions keyed by their opaque session id."""
+
+    def __init__(self) -> None:
+        self._sessions: dict[UUID, ChatSession] = {}
+        self._lock = asyncio.Lock()
+
+    async def create(self) -> ChatSession:
+        session = ChatSession()
+        async with self._lock:
+            self._sessions[session.id] = session
+        return session
+
+    async def get(self, session_id: UUID) -> ChatSession | None:
+        async with self._lock:
+            return self._sessions.get(session_id)
+
+    async def remove(self, session_id: UUID) -> None:
+        async with self._lock:
+            self._sessions.pop(session_id, None)
+
+    async def count(self) -> int:
+        async with self._lock:
+            return len(self._sessions)

app/ws.py

@@ -0,0 +1,11 @@
+"""Thin re-export module matching the suggested ``app/ws.py`` layout.
+
+The actual websocket router lives in :mod:`app.api.chat` (consistent with
+this project's convention of keeping all routers under ``app/api``). This
+module simply re-exports it so the chat service can also be wired up as
+``from app.ws import router``.
+"""
+
+from app.api.chat import router
+
+__all__ = ["router"]

pyproject.toml

@@ -5,7 +5,7 @@ description = "A modern FastAPI application with SQLAlchemy 2.0 and AsyncPG for
 readme = "README.md"
 requires-python = "==3.14.4"
 dependencies = [
-    "fastapi[all]==0.136.3",
+    "fastapi[all]==0.139.0",
     "pydantic==2.13.4",
     "pydantic-settings==2.14.1",
     "sqlalchemy==2.0.50",
@@ -31,6 +31,7 @@ dependencies = [
     "apscheduler[redis,sqlalchemy]>=4.0.0a6",
     "rotoger==0.3.0",
     "pyinstrument>=5.1.2",
+    "websockets>=15.0.1",
 ]
 
 [tool.uv]

uv.lock

@@ -308,7 +308,7 @@ wheels = [
 
 [[package]]
 name = "fastapi"
-version = "0.136.3"
+version = "0.139.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "annotated-doc" },
@@ -317,9 +317,9 @@ dependencies = [
     { name = "typing-extensions" },
     { name = "typing-inspection" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/81/2d/ff8d91d7b564d464629a0fd50a4489c97fcb836ac230bf3a7269232a9b1f/fastapi-0.136.3.tar.gz", hash = "sha256:e487fae93ad408e6f47641ee4dfe389864fd7bec92e547ea8498fc13f43e83ab", size = 396410, upload-time = "2026-05-23T18:53:15.192Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/d3/af/a5f50ccfa659ec1802cb4ca842c23f06d906a8cc9aef6016a2caeea3d4ed/fastapi-0.139.0.tar.gz", hash = "sha256:99ab7b2d92223c76d6cf10757ab3f89d45b38267fc20b2a136cf02f6beac3145", size = 423016, upload-time = "2026-07-01T16:35:33.436Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/e0/82/45359b62a067409bd929ae8a56b8ed13e5a8c8a61194b3c236920999ab83/fastapi-0.136.3-py3-none-any.whl", hash = "sha256:3d2a69bdf04b7e9f3afa292c3bc7a98816bbfafa10bc9b45f3f3700d2f761620", size = 117481, upload-time = "2026-05-23T18:53:16.924Z" },
+    { url = "https://files.pythonhosted.org/packages/9e/7c/8e3c6ad324ea5cb36604fc3f968554887891c316d9dfde57761611d907ad/fastapi-0.139.0-py3-none-any.whl", hash = "sha256:cf15e1e9e667ddb0ad63811e60bd11390d1aac838ca4a7a23f421807b2308189", size = 130339, upload-time = "2026-07-01T16:35:32.19Z" },
 ]
 
 [package.optional-dependencies]
@@ -405,6 +405,7 @@ dependencies = [
     { name = "sqlalchemy" },
     { name = "uvicorn", extra = ["standard"] },
     { name = "uvloop" },
+    { name = "websockets" },
 ]
 
 [package.dev-dependencies]
@@ -423,7 +424,7 @@ requires-dist = [
     { name = "asyncpg", specifier = "==0.31.0" },
     { name = "bcrypt", specifier = "==5.0.0" },
     { name = "dirty-equals", specifier = "==0.11" },
-    { name = "fastapi", extras = ["all"], specifier = "==0.136.3" },
+    { name = "fastapi", extras = ["all"], specifier = "==0.139.0" },
     { name = "fastexcel", specifier = "==0.20.2" },
     { name = "granian", specifier = "==2.7.6" },
     { name = "httptools", specifier = "==0.8.0" },
@@ -444,6 +445,7 @@ requires-dist = [
     { name = "sqlalchemy", specifier = "==2.0.50" },
     { name = "uvicorn", extras = ["standard"], specifier = "==0.38.0" },
     { name = "uvloop", specifier = "==0.22.1" },
+    { name = "websockets", specifier = ">=15.0.1" },
 ]
 
 [package.metadata.requires-dev]