From 2cf22f89cc35b5931848e19efea6575d850dbe1c Mon Sep 17 00:00:00 2001 From: grillazz Date: Tue, 14 Jul 2026 11:12:45 +0200 Subject: [PATCH] wip: implement websocket chat service with Pydantic models and session management --- .env | 2 + app/config.py | 21 ++++++ app/main.py | 28 +++++--- app/schemas/chat.py | 124 +++++++++++++++++++++++++++++++++++ app/services/chat_agent.py | 116 ++++++++++++++++++++++++++++++++ app/services/chat_session.py | 52 +++++++++++++++ app/ws.py | 11 ++++ pyproject.toml | 3 +- uv.lock | 10 +-- 9 files changed, 353 insertions(+), 14 deletions(-) create mode 100644 app/schemas/chat.py create mode 100644 app/services/chat_agent.py create mode 100644 app/services/chat_session.py create mode 100644 app/ws.py diff --git a/.env b/.env index 842a23b..e2d9cab 100644 --- a/.env +++ b/.env @@ -22,4 +22,6 @@ EMAIL_HOST= EMAIL_HOST_USER= EMAIL_HOST_PASSWORD= +CHAT_BACKEND=ollama + diff --git a/app/config.py b/app/config.py index dbe4877..b25d9b2 100644 --- a/app/config.py +++ b/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 diff --git a/app/main.py b/app/main.py index a484155..faa9e5b 100644 --- a/app/main.py +++ b/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"] ) diff --git a/app/schemas/chat.py b/app/schemas/chat.py new file mode 100644 index 0000000..8f844f5 --- /dev/null +++ b/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 +) diff --git a/app/services/chat_agent.py b/app/services/chat_agent.py new file mode 100644 index 0000000..93ebcae --- /dev/null +++ b/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) diff --git a/app/services/chat_session.py b/app/services/chat_session.py new file mode 100644 index 0000000..90c7579 --- /dev/null +++ b/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) diff --git a/app/ws.py b/app/ws.py new file mode 100644 index 0000000..396a690 --- /dev/null +++ b/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"] diff --git a/pyproject.toml b/pyproject.toml index bb6e282..6735c79 100644 --- a/pyproject.toml +++ b/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] diff --git a/uv.lock b/uv.lock index 6e1a25d..ce793a3 100644 --- a/uv.lock +++ b/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]