diff --git a/app/config.py b/app/config.py index c4a6992..0662cee 100644 --- a/app/config.py +++ b/app/config.py @@ -41,6 +41,21 @@ class Settings(BaseSettings): # exhausts GPU/CPU. Excess requests get 429. chat_max_concurrency: int = 4 + # Bounds on a single request (Codex M3). The concurrency cap limits how many + # generations run at once, but says nothing about how large or how long any + # one of them is -- four callers could hold every slot for the full Ollama + # timeout with a prompt the size of a book. + # + # A question is a question: 4000 characters is longer than anyone types. + chat_max_message_chars: int = 4000 + # Turns of prior conversation kept. Each one is re-sent to the model, so an + # unbounded history is an unbounded prompt, paid for on every request. + chat_max_history_turns: int = 20 + chat_max_history_chars: int = 16000 + # Hard ceiling on one streamed response, independent of the model's own + # timeout. A generation that will not stop still ends. + chat_deadline_seconds: int = 180 + model_config = {"env_prefix": "FORAIL_ASSISTANT_"} diff --git a/app/main.py b/app/main.py index abacf92..d2929d0 100644 --- a/app/main.py +++ b/app/main.py @@ -4,6 +4,7 @@ import hmac import json import logging +import time from fastapi import FastAPI, Header, HTTPException from fastapi.middleware.cors import CORSMiddleware @@ -72,6 +73,46 @@ class ChatRequest(BaseModel): history: list[dict] | None = None +def _bounded_request(req: "ChatRequest") -> tuple[str, list[dict]]: + """ + The message and history this request is allowed to spend, or 413. + + The concurrency cap limits how many generations run at once and says nothing + about how large any one of them is: four callers could hold every slot for + the full Ollama timeout with a prompt the size of a book. History matters + more than the message, because every turn is re-sent to the model and paid + for again on the next request. + """ + message = (req.message or "").strip() + if not message: + raise HTTPException(status_code=400, detail="message must not be empty") + if len(message) > settings.chat_max_message_chars: + raise HTTPException( + status_code=413, + detail=f"message must be at most {settings.chat_max_message_chars} characters", + ) + + history = req.history or [] + if not isinstance(history, list): + raise HTTPException(status_code=400, detail="history must be a list") + + # Trimmed rather than rejected: dropping the oldest turns degrades the answer + # a little, while a 413 in the middle of a conversation ends it. + history = history[-settings.chat_max_history_turns:] + budget = settings.chat_max_history_chars + kept: list[dict] = [] + for turn in reversed(history): + if not isinstance(turn, dict): + continue + cost = len(str(turn.get("content", ""))) + if cost > budget: + break + budget -= cost + kept.append(turn) + kept.reverse() + return message, kept + + class HealthResponse(BaseModel): status: str version: str @@ -129,18 +170,30 @@ async def chat(req: ChatRequest, authorization: str | None = Header(default=None if _chat_semaphore.locked(): raise HTTPException(status_code=429, detail="Assistant busy, retry shortly") + message, history = _bounded_request(req) + page_context = "" if req.context and req.context.get("page"): - page_context = req.context["page"] + page_context = str(req.context["page"])[:200] async def event_generator(): async with _chat_semaphore: + deadline = time.monotonic() + settings.chat_deadline_seconds try: async for token in stream_chat( - message=req.message, + message=message, page_context=page_context, - history=req.history, + history=history, ): + # A generation that will not stop still has to end: the + # slot it holds is one of only chat_max_concurrency. + if time.monotonic() > deadline: + logger.warning( + "Chat generation exceeded %ss deadline; cutting the stream", + settings.chat_deadline_seconds, + ) + yield {"data": json.dumps({"error": "response timed out", "done": True})} + return yield {"data": json.dumps({"token": token})} yield {"data": json.dumps({"done": True})} except Exception: diff --git a/docs/configuration.md b/docs/configuration.md index 9edf6bd..4dd39b0 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -23,6 +23,17 @@ All configuration is via environment variables with the `FORAIL_ASSISTANT_` pref | `FORAIL_ASSISTANT_CHROMA_PORT` | `8000` | ChromaDB port | | `FORAIL_ASSISTANT_CHROMA_COLLECTION` | `forail_docs` | Collection name for indexed documents | +### Request Limits + +| Variable | Default | Description | +|----------|---------|-------------| +| `FORAIL_ASSISTANT_CHAT_TOKEN` | `""` | Bearer token required on `/api/v1/chat`. **Empty means the endpoint is open** — set it whenever the service is reachable by anything you do not control | +| `FORAIL_ASSISTANT_CHAT_MAX_CONCURRENCY` | `4` | Concurrent generations; excess requests get 429 | +| `FORAIL_ASSISTANT_CHAT_MAX_MESSAGE_CHARS` | `4000` | Longest accepted question; over it returns 413 | +| `FORAIL_ASSISTANT_CHAT_MAX_HISTORY_TURNS` | `20` | Prior turns kept. Trimmed, not rejected — every turn is re-sent to the model on each request | +| `FORAIL_ASSISTANT_CHAT_MAX_HISTORY_CHARS` | `16000` | Total history size kept, oldest turns dropped first | +| `FORAIL_ASSISTANT_CHAT_DEADLINE_SECONDS` | `180` | Hard ceiling on one streamed response. A generation that will not stop still ends, so it cannot hold a concurrency slot indefinitely | + ### RAG Settings | Variable | Default | Description | diff --git a/tests/test_api.py b/tests/test_api.py index 2b60176..aae88a3 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -156,3 +156,113 @@ def test_openapi_schema(self, client): schema = resp.json() assert "/api/v1/health" in schema["paths"] assert "/api/v1/chat" in schema["paths"] + + +class TestChatRequestBounds: + """ + Codex M3: the concurrency cap limits how many generations run at once and + says nothing about how large or how long any one of them is. Four callers + could hold every slot for the full Ollama timeout with a prompt the size of + a book. + """ + + def _stream(self): + async def mock_stream(*args, **kwargs): + yield "ok" + return mock_stream + + def test_oversized_message_is_refused(self, client): + from app.config import settings + + resp = client.post( + "/api/v1/chat", + json={"message": "x" * (settings.chat_max_message_chars + 1)}, + ) + assert resp.status_code == 413 + + def test_message_at_the_limit_is_accepted(self, client): + from app.config import settings + + with patch("app.main.stream_chat", side_effect=self._stream()): + resp = client.post( + "/api/v1/chat", + json={"message": "x" * settings.chat_max_message_chars}, + ) + assert resp.status_code == 200 + + def test_blank_message_is_refused(self, client): + resp = client.post("/api/v1/chat", json={"message": " "}) + assert resp.status_code == 400 + + def test_history_is_trimmed_to_the_most_recent_turns(self, client): + # Trimmed rather than rejected: dropping the oldest turns costs a little + # context, while a 413 mid-conversation ends it. + from app.config import settings + + captured = {} + + async def mock_stream(*args, **kwargs): + captured.update(kwargs) + yield "ok" + + history = [{"role": "user", "content": f"turn {i}"} for i in range(200)] + with patch("app.main.stream_chat", side_effect=mock_stream): + resp = client.post("/api/v1/chat", json={"message": "hi", "history": history}) + + assert resp.status_code == 200 + assert len(captured["history"]) <= settings.chat_max_history_turns + # The turns kept are the recent ones, not the first ones. + assert captured["history"][-1]["content"] == "turn 199" + + def test_history_is_trimmed_by_total_size(self, client): + from app.config import settings + + captured = {} + + async def mock_stream(*args, **kwargs): + captured.update(kwargs) + yield "ok" + + history = [{"role": "user", "content": "x" * 5000} for _ in range(10)] + with patch("app.main.stream_chat", side_effect=mock_stream): + client.post("/api/v1/chat", json={"message": "hi", "history": history}) + + total = sum(len(t["content"]) for t in captured["history"]) + assert total <= settings.chat_max_history_chars + + def test_page_context_is_truncated(self, client): + captured = {} + + async def mock_stream(*args, **kwargs): + captured.update(kwargs) + yield "ok" + + with patch("app.main.stream_chat", side_effect=mock_stream): + client.post( + "/api/v1/chat", + json={"message": "hi", "context": {"page": "/x" * 5000}}, + ) + assert len(captured["page_context"]) <= 200 + + def test_a_generation_that_will_not_stop_is_cut(self, client): + # The slot it holds is one of only chat_max_concurrency, so an endless + # stream is a denial of service against the other three. A deadline + # already in the past is the same code path as one that runs out. + from app.config import settings + + emitted = 0 + + async def endless(*args, **kwargs): + nonlocal emitted + for _ in range(100): + emitted += 1 + yield "token" + + with patch("app.main.stream_chat", side_effect=endless), \ + patch.object(settings, "chat_deadline_seconds", -1): + resp = client.post("/api/v1/chat", json={"message": "hi"}) + + assert resp.status_code == 200 + assert "timed out" in resp.text + # Cut, not drained: the generator does not run to completion. + assert emitted < 100