PersonalPublic/fastapi-sqlalchemy-asyncpg
1
0
Fork 0
mirror of https://github.com/grillazz/fastapi-sqlalchemy-asyncpg.git synced 2026-11-01 01:24:45 +03:00
Code Issues Packages Projects Releases Wiki Activity

feat: refactor ChatSessionManager to be application-specific and integrate httpx2 for websocket transport

Browse Source
This commit is contained in:
grillazz
2026-10-22 13:39:46 +02:00
parent d1661f6036
commit 96b383e391
7 changed files with 148 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+8 -3
View File
@@ -185,7 +185,7 @@ ollama run llama3.2
<p align="right">(<a href="#readme-top">back to top</a>)</p>
### 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 |
app/api/chat.py
+1 -1
View File
@@ -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"))
app/main.py
+3 -4
View File
@@ -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,
app/services/chat_session.py
+3 -4
View File
@@ -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:
pyproject.toml
+1
View File
@@ -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]
tests/api/test_chat_ws.py
+66 -45
View File
@@ -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"})
uv.lock
Generated
+66
View File
@@ -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" },
]