diff --git a/README.md b/README.md index a321257..c7aaa6e 100644 --- a/README.md +++ b/README.md @@ -130,11 +130,13 @@ In `pl.read_excel()` “calamine” engine can be used for reading all major typ
### Worker aware async scheduler -The project uses the APScheduler library to schedule tasks in the background. -The APScheduler library is a powerful and flexible in-process task scheduler with Cron-like capabilities. -It allows you to schedule jobs to run at specific times or intervals, and it supports multiple job stores, triggers, and executors. -The library is designed to be easy to use and highly configurable, making it suitable for a wide range of use cases. -It was added to project in version 4.0.0a5 with Redis as event broker and SQLAlchemy as data store. +The app starts an `AsyncScheduler` as part of FastAPI's lifespan. Each worker shares +the persistent SQLAlchemy data store in the `scheduler` schema and uses Redis as the +event broker, allowing APScheduler to coordinate scheduling across workers. The +`tick-sql-25` schedule runs `app.services.scheduler.tick` every 25 seconds and is +replaced on startup so its definition stays current. The project uses APScheduler +4.0.0a6, a pre-release; see the +[ASGI/FastAPI integration example](https://github.com/agronholm/apscheduler/blob/4.0.0a6/examples/web/asgi_fastapi.py). @@ -183,7 +185,7 @@ ollama run llama3.2 ### 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 @@ -201,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 e3ab516..1817e11 100644 --- a/app/main.py +++ b/app/main.py @@ -22,6 +22,8 @@ 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") @@ -31,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, @@ -45,7 +45,8 @@ async def lifespan(app: FastAPI): # await app.logger.ainfo( # "Postgres pool created", idle_size=app.postgres_pool.get_idle_size() # ) - yield + async with scheduler_lifespan(): + yield except Exception as e: await app.logger.aerror("Error during app startup", error=repr(e)) raise @@ -96,16 +97,4 @@ def create_app() -> FastAPI: app = create_app() -# --- Unused/experimental code and TODOs --- -# from apscheduler import AsyncScheduler -# from apscheduler.datastores.sqlalchemy import SQLAlchemyDataStore -# from apscheduler.eventbrokers.redis import RedisEventBroker -# from app.database import engine -# from app.services.scheduler import SchedulerMiddleware -# _scheduler_data_store = SQLAlchemyDataStore(engine, schema="scheduler") -# _scheduler_event_broker = RedisEventBroker(client_or_url=global_settings.redis_url.unicode_string()) -# _scheduler_himself = AsyncScheduler(_scheduler_data_store, _scheduler_event_broker) -# app.add_middleware(SchedulerMiddleware, scheduler=_scheduler_himself) # TODO: every non-GET method should reset cache -# TODO: scheduler tasks needing DB should access connection pool via request -# TODO: https://stackoverflow.com/questions/16053364/make-sure-only-one-worker-launches-the-apscheduler-event-in-a-pyramid-web-app-ru diff --git a/app/services/__init__.py b/app/services/__init__.py index e69de29..a0dac8d 100644 --- a/app/services/__init__.py +++ b/app/services/__init__.py @@ -0,0 +1 @@ +"""Application service implementations.""" diff --git a/app/services/auth.py b/app/services/auth.py index 0a4f973..6499879 100644 --- a/app/services/auth.py +++ b/app/services/auth.py @@ -1,3 +1,5 @@ +"""Bearer-token authentication and Redis-backed token storage.""" + import time import jwt @@ -11,24 +13,27 @@ from app.models.user import User logger = get_logger() -async def get_from_redis(request: Request, key: str): +async def get_from_redis(request: Request, key: str) -> str | None: + """Return the Redis value stored under ``key``.""" return await request.app.redis.get(key) -async def set_to_redis(request: Request, key: str, value: str, ex: int): +async def set_to_redis(request: Request, key: str, value: str, ex: int) -> bool: + """Store ``value`` in Redis for ``ex`` seconds.""" return await request.app.redis.set(key, value, ex=ex) async def verify_jwt(request: Request, token: str) -> bool: + """Check whether Redis recognizes the bearer token.""" payload = await get_from_redis(request, token) return bool(payload) class AuthBearer(HTTPBearer): - def __init__(self, auto_error: bool = True) -> None: - super().__init__(auto_error=auto_error) + """Validate bearer credentials against the active Redis token store.""" - async def __call__(self, request: Request): + async def __call__(self, request: Request) -> str: + """Return a valid token or raise an authorization error.""" credentials: HTTPAuthorizationCredentials = await super().__call__(request) if not credentials: raise HTTPException(status_code=403, detail="Invalid authorization code.") @@ -44,8 +49,8 @@ class AuthBearer(HTTPBearer): return credentials.credentials -async def create_access_token(user: User, request: Request): - # sourcery skip: avoid-builtin-shadow +async def create_access_token(user: User, request: Request) -> str | None: + """Sign a user token and store it in Redis until its configured expiry.""" payload = { "email": user.email, "expiry": time.time() + global_settings.jwt_expire, @@ -55,8 +60,8 @@ async def create_access_token(user: User, request: Request): payload, str(user.password), algorithm=global_settings.jwt_algorithm ) - _bool = await set_to_redis( + stored = await set_to_redis( request, token, str(payload), ex=global_settings.jwt_expire ) - if _bool: + if stored: return token diff --git a/app/services/chat_agent.py b/app/services/chat_agent.py index 7ef5725..6a237b6 100644 --- a/app/services/chat_agent.py +++ b/app/services/chat_agent.py @@ -16,7 +16,7 @@ from typing import Protocol, runtime_checkable from urllib.parse import urlparse import attrs -import httpx +import httpx2 import orjson from app.config import ChatConfig @@ -49,9 +49,11 @@ class LocalEchoAgent: """ def __init__(self, stream_delay_seconds: float = 0.02) -> None: + """Initialize the delay between streamed words.""" self.stream_delay_seconds = stream_delay_seconds def _compose_reply(self, messages: list[ChatMessage]) -> str: + """Compose a canned response to the latest user message.""" last_user = next( (m.content for m in reversed(messages) if m.role == "user"), "" ) @@ -63,26 +65,20 @@ class LocalEchoAgent: 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]: + """Yield the canned reply one word at a time.""" 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 + """Close resources; the local agent holds none.""" -def _validate_base_url(_instance: OllamaChatAgent, _attribute: attrs.Attribute, value: str) -> None: - """Validate that base_url is a valid HTTP(S) URL. - - Args: - instance: The OllamaChatAgent instance being initialized. - attribute: The attrs attribute descriptor for base_url. - value: The URL string to validate. - - Raises: - ValueError: If base_url is not a valid HTTP(S) URL. - """ +def _validate_base_url( + _instance: OllamaChatAgent, _attribute: attrs.Attribute, value: str +) -> None: + """Require an absolute HTTP(S) endpoint URL.""" parsed = urlparse(value) if parsed.scheme not in ("http", "https"): msg = f"base_url must be HTTP(S), got scheme '{parsed.scheme}' from '{value}'" @@ -92,37 +88,15 @@ def _validate_base_url(_instance: OllamaChatAgent, _attribute: attrs.Attribute, raise ValueError(msg) -def _validate_timeout(_instance: OllamaChatAgent, _attribute: attrs.Attribute, value: float) -> None: - """Validate that timeout is positive. - - Args: - instance: The OllamaChatAgent instance being initialized. - attribute: The attrs attribute descriptor for timeout. - value: The timeout value in seconds. - - Raises: - ValueError: If timeout is not positive. - """ +def _validate_timeout( + _instance: OllamaChatAgent, _attribute: attrs.Attribute, value: float +) -> None: + """Require a positive request timeout.""" if value <= 0: msg = f"timeout must be positive, got {value}" raise ValueError(msg) -def _create_httpx_client(instance: OllamaChatAgent) -> httpx.AsyncClient: - """Factory function to create the httpx.AsyncClient with validated config. - - This function is called from __attrs_post_init__ after all field validators - have run, ensuring the client is created with validated configuration. - - Args: - instance: The OllamaChatAgent instance being initialized. - - Returns: - An initialized httpx.AsyncClient configured with base_url and timeout. - """ - return httpx.AsyncClient(base_url=instance.base_url, timeout=instance.timeout) - - @attrs.define(slots=True, eq=False, hash=False) class OllamaChatAgent: """Streams chat completions from an OpenAI-compatible endpoint. @@ -132,9 +106,6 @@ class OllamaChatAgent: a ready-to-swap-in replacement for :class:`LocalEchoAgent` once a real model should be used. - Attrs Configuration: - - slots=True: Memory-efficient attribute storage (~40-50% reduction) - - eq=False, hash=False: Instances are not comparable (contain async resources) """ model: str = attrs.field( @@ -160,20 +131,15 @@ class OllamaChatAgent: "constraints": "Must be positive", }, ) - _client: httpx.AsyncClient = attrs.field( + _client: httpx2.AsyncClient = attrs.field( init=False, repr=False, metadata={"description": "Internal HTTP client for API communication"}, ) def __attrs_post_init__(self) -> None: - """Initialize the HTTP client after field validation. - - This hook is called by attrs after __init__ completes and all field - validators have run. It's used to initialize the internal _client - field which depends on validated configuration. - """ - self._client = _create_httpx_client(self) + """Create the HTTPX2 client after attrs validates the fields.""" + self._client = httpx2.AsyncClient(base_url=self.base_url, timeout=self.timeout) async def stream_reply(self, messages: list[ChatMessage]) -> AsyncIterator[str]: """Stream chat completion responses from the configured model. @@ -195,13 +161,11 @@ class OllamaChatAgent: 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 = "" + data = orjson.loads(line.removeprefix("data: ")) + choices = data.get("choices", []) + if not choices: + continue + content = choices[0].get("delta", {}).get("content", "") if content: yield content diff --git a/app/services/chat_session.py b/app/services/chat_session.py index 4de6ab6..8b8b3bd 100644 --- a/app/services/chat_session.py +++ b/app/services/chat_session.py @@ -13,46 +13,54 @@ from uuid import UUID, uuid4 import attrs from app.schemas.chat import ChatMessage -from app.utils.singleton import SingletonMetaNoArgs @attrs.define(slots=True) class ChatSession: + """Conversation history associated with one websocket connection.""" + id: UUID = attrs.field(factory=uuid4) messages: list[ChatMessage] = attrs.field(factory=list) def add(self, message: ChatMessage) -> None: + """Append a message to this conversation.""" self.messages.append(message) def history(self) -> list[ChatMessage]: + """Return a copy of the conversation history.""" 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: + """Create an empty, lock-protected session registry.""" self._sessions: dict[UUID, ChatSession] = {} self._lock = asyncio.Lock() async def create(self) -> ChatSession: + """Create and register a new conversation.""" session = ChatSession() async with self._lock: self._sessions[session.id] = session return session async def get(self, session_id: UUID) -> ChatSession | None: + """Return the matching session, if it is active.""" async with self._lock: return self._sessions.get(session_id) async def remove(self, session_id: UUID) -> None: + """Remove a session from the active registry.""" async with self._lock: self._sessions.pop(session_id, None) async def count(self) -> int: + """Return the number of active sessions.""" async with self._lock: return len(self._sessions) diff --git a/app/services/llm.py b/app/services/llm.py index ceaa455..cc1f4a3 100644 --- a/app/services/llm.py +++ b/app/services/llm.py @@ -1,25 +1,28 @@ -from collections.abc import AsyncGenerator +"""Streaming client for OpenAI-compatible chat completion endpoints.""" -import httpx +from collections.abc import AsyncIterator + +import httpx2 import orjson class StreamLLMService: + """Send prompts to an LLM endpoint and yield newline-delimited JSON events.""" + def __init__(self, base_url: str = "http://localhost:11434/v1") -> None: + """Configure the endpoint and default model.""" self.base_url = base_url self.model = "llama3.2" - async def stream_chat(self, prompt: str) -> AsyncGenerator[bytes]: - """Stream chat completion responses from LLM.""" - # Send the user a message first + async def stream_chat(self, prompt: str) -> AsyncIterator[bytes]: + """Yield the prompt followed by streamed model response chunks.""" user_msg = { "role": "user", "content": prompt, } yield orjson.dumps(user_msg) + b"\n" - # Open client as context manager and stream responses - async with httpx.AsyncClient(base_url=self.base_url) as client: + async with httpx2.AsyncClient(base_url=self.base_url) as client: async with client.stream( "POST", "/chat/completions", @@ -32,21 +35,16 @@ class StreamLLMService: ) as response: async for line in response.aiter_lines(): if line.startswith("data: ") and line != "data: [DONE]": - try: - json_line = line[6:] # Remove "data: " prefix - data = orjson.loads(json_line) - content = ( - data.get("choices", [{}])[0] - .get("delta", {}) - .get("content", "") - ) - if content: - model_msg = {"role": "model", "content": content} - yield orjson.dumps(model_msg) + b"\n" - except Exception: - pass + data = orjson.loads(line.removeprefix("data: ")) + choices = data.get("choices", []) + if not choices: + continue + content = choices[0].get("delta", {}).get("content", "") + if content: + model_msg = {"role": "model", "content": content} + yield orjson.dumps(model_msg) + b"\n" -# FastAPI dependency def get_llm_service(base_url: str | None = None) -> StreamLLMService: + """Build the service using the supplied URL or its local default.""" return StreamLLMService(base_url=base_url or "http://localhost:11434/v1") diff --git a/app/services/scheduler.py b/app/services/scheduler.py index 6352fb6..be2ff5c 100644 --- a/app/services/scheduler.py +++ b/app/services/scheduler.py @@ -1,18 +1,24 @@ +"""Background database health-check scheduling.""" + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from datetime import datetime -from apscheduler import AsyncScheduler +from apscheduler import AsyncScheduler, ConflictPolicy +from apscheduler.datastores.sqlalchemy import SQLAlchemyDataStore +from apscheduler.eventbrokers.redis import RedisEventBroker from apscheduler.triggers.interval import IntervalTrigger -from attrs import define from rotoger import get_logger from sqlalchemy import text -from starlette.types import ASGIApp, Receive, Scope, Send -from app.database import AsyncSessionFactory +from app.config import settings as global_settings +from app.database import AsyncSessionFactory, engine logger = get_logger() -async def tick(): +async def tick() -> bool: + """Run the scheduler's lightweight database heartbeat.""" async with AsyncSessionFactory() as session: stmt = text("select 1;") await logger.ainfo(f">>>> Be or not to be...{datetime.now()}") @@ -21,26 +27,20 @@ async def tick(): return True -@define -class SchedulerMiddleware: - app: ASGIApp - scheduler: AsyncScheduler +@asynccontextmanager +async def scheduler_lifespan() -> AsyncIterator[None]: + """Start the scheduler and stop it when the app lifespan ends.""" + scheduler = AsyncScheduler( + SQLAlchemyDataStore(engine, schema="scheduler"), + RedisEventBroker(client_or_url=global_settings.redis_url.unicode_string()), + ) - async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: - """ - Handles the incoming request and schedules tasks if the scope type is 'lifespan'. - - Args: - scope (Scope): The ASGI scope dictionary containing request information. - receive (Receive): The ASGI receive callable. - send (Send): The ASGI send callable. - """ - if scope["type"] == "lifespan": - async with self.scheduler: - await self.scheduler.add_schedule( - tick, IntervalTrigger(seconds=25), id="tick-sql-25" - ) - await self.scheduler.start_in_background() - await self.app(scope, receive, send) - else: - await self.app(scope, receive, send) + async with scheduler: + await scheduler.add_schedule( + tick, + IntervalTrigger(seconds=25), + id="tick-sql-25", + conflict_policy=ConflictPolicy.replace, + ) + await scheduler.start_in_background() + yield diff --git a/app/services/smtp.py b/app/services/smtp.py index 3d09932..16734f7 100644 --- a/app/services/smtp.py +++ b/app/services/smtp.py @@ -1,3 +1,5 @@ +"""Email delivery through a configured SMTP server.""" + import smtplib from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText diff --git a/pyproject.toml b/pyproject.toml index 9b88b57..966be12 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,7 @@ dependencies = [ "uvicorn[standard]==0.54.0", "asyncpg==0.32.0", "alembic==1.20.1", - "httpx==0.28.1", + "httpx2==2.16.0", "uvloop==0.23.0", "httptools==0.9.0", "rich==15.0.0", @@ -29,6 +29,7 @@ dependencies = [ "websockets==17.3", "jinja2==3.1.6", "whenever==0.10.5", + "sniffio==1.3.1", ] [dependency-groups] @@ -46,6 +47,8 @@ test = [ "inline-snapshot==0.36.1", "dirty-equals==0.11", "polyfactory==3.5.0", + "httpx==0.28.1", + "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/tests/test_ollama_chat_agent_attrs.py b/tests/test_ollama_chat_agent_attrs.py index 19250f2..3866677 100644 --- a/tests/test_ollama_chat_agent_attrs.py +++ b/tests/test_ollama_chat_agent_attrs.py @@ -1,5 +1,6 @@ """Test suite for refactored OllamaChatAgent with attrs features.""" -import httpx + +import httpx2 import pytest from app.config import ChatConfig @@ -54,14 +55,14 @@ class TestOllamaChatAgentAttrsFeatures: assert isinstance(agent.timeout, float) assert agent.timeout == 30.0 - def test_factory_initialization_of_client(self): - """Test that _client is created via factory function.""" + def test_client_initialization_after_field_validation(self): + """Test that the async client uses the validated configuration.""" agent = OllamaChatAgent(model="llama3.2", base_url="http://localhost:11434/v1") assert hasattr(agent, "_client") - assert isinstance(agent._client, httpx.AsyncClient) - # httpx normalizes URLs by adding a trailing slash + assert isinstance(agent._client, httpx2.AsyncClient) + # httpx2 normalizes URLs by adding a trailing slash assert str(agent._client.base_url) == "http://localhost:11434/v1/" - assert agent._client.timeout == httpx.Timeout(60.0) + assert agent._client.timeout == httpx2.Timeout(60.0) def test_slots_enabled_no_dict(self): """Test that slots=True prevents __dict__ attribute.""" @@ -72,12 +73,8 @@ class TestOllamaChatAgentAttrsFeatures: def test_equality_disabled(self): """Test that eq=False means instances are not equal even with same values.""" - agent1 = OllamaChatAgent( - model="llama3.2", base_url="http://localhost:11434/v1" - ) - agent2 = OllamaChatAgent( - model="llama3.2", base_url="http://localhost:11434/v1" - ) + agent1 = OllamaChatAgent(model="llama3.2", base_url="http://localhost:11434/v1") + agent2 = OllamaChatAgent(model="llama3.2", base_url="http://localhost:11434/v1") # With eq=False, only identity comparison works assert agent1 != agent2 assert agent1 == agent1 @@ -112,9 +109,7 @@ class TestOllamaChatAgentAttrsFeatures: def test_https_urls_accepted(self): """Test that HTTPS URLs are properly accepted.""" - agent = OllamaChatAgent( - model="llama3.2", base_url="https://api.example.com/v1" - ) + agent = OllamaChatAgent(model="llama3.2", base_url="https://api.example.com/v1") assert agent.base_url == "https://api.example.com/v1" def test_metadata_annotations_present(self): @@ -133,4 +128,3 @@ class TestOllamaChatAgentAttrsFeatures: if __name__ == "__main__": pytest.main([__file__, "-v"]) - diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py new file mode 100644 index 0000000..a21db8b --- /dev/null +++ b/tests/test_scheduler.py @@ -0,0 +1,62 @@ +import pytest +from apscheduler import ConflictPolicy +from apscheduler.triggers.interval import IntervalTrigger + +import app.services.scheduler as scheduler_service + + +@pytest.mark.anyio +async def test_scheduler_lifespan_starts_and_stops_scheduler(monkeypatch) -> None: + events = [] + data_store = object() + event_broker = object() + + class SchedulerStub: + async def __aenter__(self): + events.append("enter") + return self + + async def __aexit__(self, *_args: object) -> None: + events.append("exit") + + async def add_schedule( + self, func, trigger, *, id: str, conflict_policy: ConflictPolicy + ) -> None: + events.append(("schedule", func, trigger, id, conflict_policy)) + + async def start_in_background(self) -> None: + events.append("start") + + scheduler = SchedulerStub() + + def create_data_store(actual_engine, *, schema): + assert actual_engine is scheduler_service.engine + assert schema == "scheduler" + return data_store + + def create_event_broker(*, client_or_url): + assert client_or_url == ( + scheduler_service.global_settings.redis_url.unicode_string() + ) + return event_broker + + def create_scheduler(actual_data_store, actual_event_broker): + assert actual_data_store is data_store + assert actual_event_broker is event_broker + return scheduler + + monkeypatch.setattr(scheduler_service, "SQLAlchemyDataStore", create_data_store) + monkeypatch.setattr(scheduler_service, "RedisEventBroker", create_event_broker) + monkeypatch.setattr(scheduler_service, "AsyncScheduler", create_scheduler) + + async with scheduler_service.scheduler_lifespan(): + assert events[0] == "enter" + event_type, func, trigger, schedule_id, conflict_policy = events[1] + assert event_type == "schedule" + assert func is scheduler_service.tick + assert isinstance(trigger, IntervalTrigger) + assert schedule_id == "tick-sql-25" + assert conflict_policy is ConflictPolicy.replace + assert events[2:] == ["start"] + + assert events[-1] == "exit" diff --git a/uv.lock b/uv.lock index e44ae8f..ce7e279 100644 --- a/uv.lock +++ b/uv.lock @@ -330,7 +330,7 @@ dependencies = [ { name = "fastexcel" }, { name = "granian" }, { name = "httptools" }, - { name = "httpx" }, + { name = "httpx2" }, { name = "jinja2" }, { name = "polars" }, { name = "pydantic", extra = ["email"] }, @@ -341,6 +341,7 @@ dependencies = [ { name = "redis" }, { name = "rich" }, { name = "rotoger" }, + { name = "sniffio" }, { name = "sqlalchemy" }, { name = "uvicorn", extra = ["standard"] }, { name = "uvloop" }, @@ -358,6 +359,8 @@ dev = [ ] test = [ { name = "dirty-equals" }, + { name = "httpx" }, + { name = "httpx2", extra = ["ws"] }, { name = "inline-snapshot" }, { name = "polyfactory" }, { name = "pytest" }, @@ -374,7 +377,7 @@ requires-dist = [ { name = "fastexcel", specifier = "==0.21.0" }, { name = "granian", specifier = "==2.8.5" }, { name = "httptools", specifier = "==0.9.0" }, - { name = "httpx", specifier = "==0.28.1" }, + { name = "httpx2", specifier = "==2.16.0" }, { name = "jinja2", specifier = "==3.1.6" }, { name = "polars", specifier = "==2.0.0" }, { name = "pydantic", extras = ["email"], specifier = "==2.14.1" }, @@ -385,6 +388,7 @@ requires-dist = [ { name = "redis", specifier = "==8.1.0" }, { name = "rich", specifier = "==15.0.0" }, { name = "rotoger", specifier = "==0.3.0" }, + { name = "sniffio", specifier = "==1.3.1" }, { name = "sqlalchemy", specifier = "==2.1.6" }, { name = "uvicorn", extras = ["standard"], specifier = "==0.54.0" }, { name = "uvloop", specifier = "==0.23.0" }, @@ -402,6 +406,8 @@ dev = [ ] test = [ { name = "dirty-equals", specifier = "==0.11" }, + { name = "httpx", specifier = "==0.28.1" }, + { 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" }, @@ -511,6 +517,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" @@ -560,6 +579,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" @@ -1177,6 +1226,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] +[[package]] +name = "sniffio" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a2/87/a6771e1546d97e7e041b6ae58d80074f81b7d5121207425c964ddf5cfdbd/sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc", size = 20372, upload-time = "2024-02-25T23:20:04.057Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235, upload-time = "2024-02-25T23:20:01.196Z" }, +] + [[package]] name = "sqlalchemy" version = "2.1.6" @@ -1276,6 +1334,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" @@ -1530,3 +1597,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" }, +]