diff --git a/README.md b/README.md index 74a9e4d..c7aaa6e 100644 --- a/README.md +++ b/README.md @@ -185,7 +185,7 @@ ollama run llama3.2

(back to top)

### Websocket AI chat service -A general-purpose, websocket-based chat service lives at `/v1/chat/ws`, built with FastAPI 0.139.0 +A general-purpose, websocket-based chat service lives at `/v1/chat/ws`, built with FastAPI 0.159.1 and Pydantic models for every message shape exchanged over the wire. It follows the [Pydantic AI chat-app](https://github.com/pydantic/pydantic-ai/blob/main/docs/web.md) interaction pattern - a session id, role-labeled messages, and streamed assistant chunks - so a UI can be wired @@ -203,10 +203,15 @@ up later without changing the backend contract. chat flow. **Lifespan wiring**: `app/main.py`'s `lifespan` builds the configured `ChatAgent` via -`build_chat_agent(settings.chat)` and a fresh `ChatSessionManager` on startup (`app.chat_agent` -/ `app.state.chat_sessions`), and calls `chat_agent.aclose()` on shutdown to release any held +`build_chat_agent(settings.chat)` and a fresh `ChatSessionManager` on startup. The agent is stored +on `app.chat_agent`, and the session manager on `app.state.chat_sessions`; the websocket route +uses that manager to track active sessions. On shutdown, `chat_agent.aclose()` releases any held connections. +The WebSocket integration tests use HTTPX2's async `ASGIWebSocketTransport` +(`httpx2[ws]`) to exercise the application without starting a server. They manage +the app lifespan explicitly because the transport does not start it. + **Configuration** (`ChatConfig` in `app/config.py`, all overridable via environment variables): | Env var | Default | Description | diff --git a/app/api/chat.py b/app/api/chat.py index a33247e..7d80d2c 100644 --- a/app/api/chat.py +++ b/app/api/chat.py @@ -100,7 +100,7 @@ async def chat_websocket(websocket: WebSocket) -> None: await websocket.accept() agent: ChatAgent = websocket.app.chat_agent - session_manager: ChatSessionManager = ChatSessionManager() + session_manager: ChatSessionManager = websocket.app.state.chat_sessions session = await session_manager.create() await websocket.send_json(Connected(session_id=session.id).model_dump(mode="json")) diff --git a/app/main.py b/app/main.py index c5f6068..1817e11 100644 --- a/app/main.py +++ b/app/main.py @@ -22,6 +22,7 @@ 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 from app.services.scheduler import scheduler_lifespan templates = Jinja2Templates(directory=Path(__file__).parent.parent / "templates") @@ -32,11 +33,9 @@ async def lifespan(app: FastAPI): app.logger = get_logger() app.redis = await get_redis() global_settings.postgres_url.unicode_string() - # Chat service: initialize the pluggable model-client adapter. The session - # manager is now a singleton and will be auto-instantiated on first access. - # See app/services/chat_agent.py to swap the local stub for a real model - # client (OpenAI, Ollama, etc.). + # Initialize the chat agent and per-application session registry. app.chat_agent = build_chat_agent(global_settings.chat) + app.state.chat_sessions = ChatSessionManager() try: # app.postgres_pool = await asyncpg.create_pool( # dsn=postgres_dsn, diff --git a/app/services/chat_session.py b/app/services/chat_session.py index 4de6ab6..92d87f4 100644 --- a/app/services/chat_session.py +++ b/app/services/chat_session.py @@ -13,7 +13,6 @@ from uuid import UUID, uuid4 import attrs from app.schemas.chat import ChatMessage -from app.utils.singleton import SingletonMetaNoArgs @attrs.define(slots=True) @@ -28,11 +27,11 @@ class ChatSession: return list(self.messages) -class ChatSessionManager(metaclass=SingletonMetaNoArgs): +class ChatSessionManager: """Tracks active chat sessions keyed by their opaque session id. - Implemented as a singleton to ensure exactly one instance per application, - maintaining a consistent registry of all active websocket chat sessions. + Each application lifespan owns its own manager, maintaining a consistent + registry of active websocket chat sessions for that application. """ def __init__(self) -> None: diff --git a/pyproject.toml b/pyproject.toml index 2910ee0..c76e5b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,7 @@ test = [ "inline-snapshot==0.36.1", "dirty-equals==0.11", "polyfactory==3.5.0", + "httpx2[ws]==2.16.0", ] [tool.pyrefly.errors] diff --git a/tests/api/test_chat_ws.py b/tests/api/test_chat_ws.py index 8867df6..629d30d 100644 --- a/tests/api/test_chat_ws.py +++ b/tests/api/test_chat_ws.py @@ -1,81 +1,106 @@ """Tests for the websocket chat service using the local stub agent. -Uses Starlette's synchronous ``TestClient`` websocket support (rather than -the ``httpx.AsyncClient`` fixture used elsewhere) since it triggers the -FastAPI lifespan on ``__enter__``/``__exit__``, wiring up -``app.chat_agent`` / ``app.state.chat_sessions`` exactly like a real -server run. +Uses HTTPX2's async ASGI WebSocket transport. The app lifespan is managed +explicitly because the transport does not trigger lifespan events itself. """ +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + import pytest -from fastapi.testclient import TestClient +from httpx2 import AsyncClient +from httpx2.websockets import ASGIWebSocketTransport from app.main import app pytestmark = pytest.mark.anyio -async def test_chat_ws_full_conversation_flow() -> None: - with ( - TestClient(app) as client, - client.websocket_connect("/v1/chat/ws") as websocket, - ): +@pytest.fixture() +async def chat_client( + monkeypatch: pytest.MonkeyPatch, +) -> AsyncIterator[AsyncClient]: + class RedisStub: + async def aclose(self) -> None: + pass + + async def get_redis_stub() -> RedisStub: + return RedisStub() + + @asynccontextmanager + async def scheduler_lifespan_stub() -> AsyncIterator[None]: + yield + + monkeypatch.setattr("app.main.get_redis", get_redis_stub) + monkeypatch.setattr("app.main.scheduler_lifespan", scheduler_lifespan_stub) + + async with app.router.lifespan_context(app): + async with AsyncClient( + transport=ASGIWebSocketTransport(app), + base_url="http://testserver", + ) as client: + yield client + + +async def test_chat_ws_full_conversation_flow(chat_client: AsyncClient) -> None: + assert await app.state.chat_sessions.count() == 0 + + async with chat_client.websocket("ws://testserver/v1/chat/ws") as websocket: # 1. connect - connected = websocket.receive_json() + connected = await websocket.receive_json() assert connected["type"] == "connected" session_id = connected["session_id"] + assert await app.state.chat_sessions.count() == 1 # 2. start conversation with a system prompt - websocket.send_json( + await websocket.send_json( {"type": "start", "system_prompt": "You are a helpful assistant."} ) - started = websocket.receive_json() - print(f"{started=}") + started = await websocket.receive_json() assert started == {"type": "started", "session_id": session_id} # 3. send a user message and aggregate the streamed assistant chunks - websocket.send_json({"type": "user_message", "content": "hello"}) + await websocket.send_json({"type": "user_message", "content": "hello"}) - event = websocket.receive_json() + event = await websocket.receive_json() aggregated = "" while event["type"] == "assistant_chunk": assert event["index"] >= 0 aggregated += event["content"] - event = websocket.receive_json() + event = await websocket.receive_json() assert event["type"] == "assistant_message" assert event["message"]["role"] == "assistant" assert event["message"]["content"] == aggregated - print(f"{event=}") - print(f"{aggregated=}") assert "hello" in aggregated.lower() # 4. request full history (system + user + assistant) - websocket.send_json({"type": "history_request"}) - history = websocket.receive_json() + await websocket.send_json({"type": "history_request"}) + history = await websocket.receive_json() assert history["type"] == "history" roles = [m["role"] for m in history["messages"]] assert roles == ["system", "user", "assistant"] # 5. gracefully end the conversation - websocket.send_json({"type": "end"}) + await websocket.send_json({"type": "end"}) + + assert await app.state.chat_sessions.count() == 0 -async def test_chat_ws_invalid_event_returns_error() -> None: - with ( - TestClient(app) as client, - client.websocket_connect("/v1/chat/ws") as websocket, - ): - websocket.receive_json() # connected +async def test_chat_ws_invalid_event_returns_error(chat_client: AsyncClient) -> None: + async with chat_client.websocket("ws://testserver/v1/chat/ws") as websocket: + await websocket.receive_json() # connected - websocket.send_json({"type": "not_a_real_event"}) - error = websocket.receive_json() + await websocket.send_json({"type": "not_a_real_event"}) + error = await websocket.receive_json() assert error["type"] == "error" - websocket.send_json({"type": "end"}) + await websocket.send_json({"type": "end"}) -async def test_chat_ws_ask_question_and_print_model_answer() -> None: +async def test_chat_ws_ask_question_and_print_model_answer( + chat_client: AsyncClient, +) -> None: """Ask a real question over the websocket chat and print the model's answer. Runs against whatever ``ChatAgent`` is wired up via ``app.state.chat_agent`` @@ -85,30 +110,27 @@ async def test_chat_ws_ask_question_and_print_model_answer() -> None: """ question = "What is the capital of France?" - with ( - TestClient(app) as client, - client.websocket_connect("/v1/chat/ws") as websocket, - ): + async with chat_client.websocket("ws://testserver/v1/chat/ws") as websocket: # 1. connect - connected = websocket.receive_json() + connected = await websocket.receive_json() assert connected["type"] == "connected" session_id = connected["session_id"] # 2. start conversation - websocket.send_json( + await websocket.send_json( {"type": "start", "system_prompt": "You are a helpful assistant."} ) - started = websocket.receive_json() + started = await websocket.receive_json() assert started == {"type": "started", "session_id": session_id} # 3. ask the question and aggregate the streamed answer - websocket.send_json({"type": "user_message", "content": question}) + await websocket.send_json({"type": "user_message", "content": question}) - event = websocket.receive_json() + event = await websocket.receive_json() answer = "" while event["type"] == "assistant_chunk": answer += event["content"] - event = websocket.receive_json() + event = await websocket.receive_json() assert event["type"] == "assistant_message" answer = event["message"]["content"] @@ -120,5 +142,4 @@ async def test_chat_ws_ask_question_and_print_model_answer() -> None: assert answer.strip() != "" # 4. gracefully end the conversation - websocket.send_json({"type": "end"}) - + await websocket.send_json({"type": "end"}) diff --git a/uv.lock b/uv.lock index 7ef411d..41ffe86 100644 --- a/uv.lock +++ b/uv.lock @@ -359,6 +359,7 @@ dev = [ ] test = [ { name = "dirty-equals" }, + { name = "httpx2", extra = ["ws"] }, { name = "inline-snapshot" }, { name = "polyfactory" }, { name = "pytest" }, @@ -404,6 +405,7 @@ dev = [ ] test = [ { name = "dirty-equals", specifier = "==0.11" }, + { name = "httpx2", extras = ["ws"], specifier = "==2.16.0" }, { name = "inline-snapshot", specifier = "==0.36.1" }, { name = "polyfactory", specifier = "==3.5.0" }, { name = "pytest", specifier = "==9.1.1" }, @@ -513,6 +515,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, ] +[[package]] +name = "httpcore2" +version = "2.16.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "h11" }, + { name = "truststore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c0/61/afa9d6b7452f9eca73f73d45324eb0bb52f098fd5e6049bf6b8263be9b4e/httpcore2-2.16.0.tar.gz", hash = "sha256:c72cd0413563b17e1cd8354677dff887466f8ec2551679759e93126534428643", size = 69886, upload-time = "2026-10-21T07:09:45.31Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/af/c2/ce417c397afa5181f0e3c6ff01a4d01639414e1adef08730d97571297151/httpcore2-2.16.0-py3-none-any.whl", hash = "sha256:5c4358319f2602d984ad93182c9fee5160f22e9da7babdda870a5e7c4e2dc34f", size = 85070, upload-time = "2026-10-21T07:09:42.802Z" }, +] + [[package]] name = "httptools" version = "0.9.0" @@ -562,6 +577,36 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, ] +[[package]] +name = "httpx2" +version = "2.16.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio", marker = "sys_platform != 'emscripten'" }, + { name = "httpcore2", marker = "sys_platform != 'emscripten'" }, + { name = "httpx2-jsfetch", marker = "sys_platform == 'emscripten'" }, + { name = "idna" }, + { name = "truststore", marker = "sys_platform != 'emscripten'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/73/5ca2cedf5e4e50423d1c258877865a6b62793633b561c58499f8aeb33b7e/httpx2-2.16.0.tar.gz", hash = "sha256:0d54c98128e8f2ff7fbbde0fca8c359e5f032edeaaaac443cdb7f31f685bcd7f", size = 103116, upload-time = "2026-10-21T07:09:46.348Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7b/3d/2e4bc4fd1ef9770871669adb348d612a6e4786b303a23f9935b81c48dd60/httpx2-2.16.0-py3-none-any.whl", hash = "sha256:fd2c5d7909e8d86fe92d53020c310ca0e4e431b838c207c5754b9bbf0308d49a", size = 97880, upload-time = "2026-10-21T07:09:44.01Z" }, +] + +[package.optional-dependencies] +ws = [ + { name = "wsproto" }, +] + +[[package]] +name = "httpx2-jsfetch" +version = "1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/c4/0e5636363151a2a1795e0a77617168b9ca438e1748ec05fc9b5687f93d64/httpx2_jsfetch-1.0.tar.gz", hash = "sha256:70a0e3eabfef7cce5ad9c629f7d01ca05e418f586646f4ddf14782e4c1454c60", size = 6872, upload-time = "2026-08-07T00:13:07.492Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9b/43/832f631d32e4f1211caa2ba368317739fe71f0b8530e4c9d15dc454bac2a/httpx2_jsfetch-1.0-py3-none-any.whl", hash = "sha256:cb916b707601e69a07721aabc8f3f6659be3a6893bc1ff5c6f9e02241df2da32", size = 6382, upload-time = "2026-08-07T00:13:06.567Z" }, +] + [[package]] name = "idna" version = "3.20" @@ -1287,6 +1332,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ad/66/0d785f0bc5e4315a96c989bb476d0fc07ea4f85132550c7b156ca2035d52/traitlets-5.16.1-py3-none-any.whl", hash = "sha256:f775618166caa0396c8e337099240f2bd3e5e917d203b2e6fbe21a58d3cb1f6b", size = 86211, upload-time = "2026-08-03T08:32:34.48Z" }, ] +[[package]] +name = "truststore" +version = "0.10.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ee/9f/c5201d42a484c061e528825fc8e2d565f5abd50a4ced6fb7d29c4ec99b2b/truststore-0.10.5.tar.gz", hash = "sha256:30d36967ccaded5cbb38d602c433f53600036c79d502f4533a49b60a03bbefcd", size = 28091, upload-time = "2026-10-12T22:27:31.808Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/e9/3a7820be2bb0fe53b6bc9c3be26d3d1158004e4c3ab953aa6840b955b1e9/truststore-0.10.5-py3-none-any.whl", hash = "sha256:9aaaedaefaf06d8b206278cf8b5012bc897f485a874503501e12d776df78951c", size = 19017, upload-time = "2026-10-12T22:27:30.377Z" }, +] + [[package]] name = "tryceratops" version = "2.4.1" @@ -1541,3 +1595,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/31/06/c969504aa9887010bac12484df2f03aa0e7ca5ada615a989da04e8f33046/whenever-0.10.5-cp314-cp314t-win_amd64.whl", hash = "sha256:06c32f8931bbc52b6b3f513e831aeb936e231f038b2a66cd11d4a00c131b82e6", size = 544821, upload-time = "2026-08-07T15:02:39.579Z" }, { url = "https://files.pythonhosted.org/packages/f2/84/c937c7ec9de2aa59fbcdd60b52a8da60069636d2252c7461e43d389c04a4/whenever-0.10.5-py3-none-any.whl", hash = "sha256:193091633fc1dece4c86e7303cbd4039320b4f45c69dc5ae69b8f711bca82603", size = 129302, upload-time = "2026-08-07T15:03:29.285Z" }, ] + +[[package]] +name = "wsproto" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "h11" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c7/79/12135bdf8b9c9367b8701c2c19a14c913c120b882d50b014ca0d38083c2c/wsproto-1.3.2.tar.gz", hash = "sha256:b86885dcf294e15204919950f666e06ffc6c7c114ca900b060d6e16293528294", size = 50116, upload-time = "2025-11-20T18:18:01.871Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/f5/10b68b7b1544245097b2a1b8238f66f2fc6dcaeb24ba5d917f52bd2eed4f/wsproto-1.3.2-py3-none-any.whl", hash = "sha256:61eea322cdf56e8cc904bd3ad7573359a242ba65688716b0710a5eb12beab584", size = 24405, upload-time = "2025-11-20T18:18:00.454Z" }, +]