From 0ce4dec232982db0985a11f58bae0e33460824e7 Mon Sep 17 00:00:00 2001 From: Duane Millar Barlow Date: Thu, 12 Nov 2026 17:36:14 -0500 Subject: [PATCH] Bound chat turns in time, log their usage, and name spend limits A stalled provider could hold a stream open until the proxy gave up, so a turn now ends with a readable error after CHAT_TURN_TIMEOUT seconds (180 by default). Each turn writes one chat_turn log line with its outcome, tokens, and cost, never the question or answer. Provider spend limits that arrive as 403 or 429 instead of 402 now read as the usage limit being reached. --- .env.example | 3 +- backend/chat/agent.py | 104 +++++++++++++++++++++++++++------ backend/chat/config.py | 4 ++ backend/chat/router.py | 10 +++- backend/tests/test_chat.py | 69 ++++++++++++++++++++++ docs/architecture/data-chat.md | 7 ++- justfile | 2 +- 7 files changed, 176 insertions(+), 23 deletions(-) diff --git a/.env.example b/.env.example index 5cacacd3..5f5d4b78 100644 --- a/.env.example +++ b/.env.example @@ -97,12 +97,13 @@ CENSUS_API_KEY= # `just chat-eval-ui` sets this; never set it in production. # CHAT_EVALS_DIR= -# Per-question bounds (model requests, tool calls, output tokens), rows of each +# Per-question bounds (model requests, tool calls, output tokens, seconds), rows of each # tool result shown to the model, history limits, questions per IP per hour # (per worker process), and rows returned for one chart. # CHAT_MAX_MODEL_REQUESTS=20 # CHAT_MAX_TOOL_CALLS=30 # CHAT_MAX_OUTPUT_TOKENS=40000 +# CHAT_TURN_TIMEOUT=180 # CHAT_PREVIEW_ROWS=20 # CHAT_MAX_HISTORY_MESSAGES=20 # CHAT_MAX_MESSAGE_CHARS=4000 diff --git a/backend/chat/agent.py b/backend/chat/agent.py index 70690790..d1200228 100644 --- a/backend/chat/agent.py +++ b/backend/chat/agent.py @@ -2,7 +2,10 @@ from __future__ import annotations +import asyncio +import json import logging +import time from collections.abc import AsyncIterator from typing import Any @@ -41,6 +44,10 @@ from data_tools.service import DataToolService logger = logging.getLogger(__name__) +# One line per turn with counts and cost, never the question or the answer: +# the only record of real usage until there is a spend ledger. The router +# gives it a handler, since the API configures no INFO logging of its own. +usage_log = logging.getLogger("chat.usage") # export_data returns inline CSV, which the browser gets by reference instead. AGENT_DATA_TOOLS = tuple(name for name in TOOL_MODELS if name != "export_data") @@ -260,11 +267,26 @@ def build_agent(settings: ChatSettings, service: DataToolService) -> Agent: ) -def http_error_message(status: int) -> str: - """What to tell the user when the model provider answers with an error.""" +# How providers say a spend limit is used up when they don't answer 402: +# OpenRouter can send 403 with metadata.limit_source "openrouter_key_limit", +# and OpenAI sends 429 with code "project_spend_limit_exceeded" (or +# "organization_..."). Without these, the first reads as "not configured" and +# the second as "busy, try again". +SPEND_LIMIT_MARKERS = ("openrouter_key_limit", "spend_limit_exceeded") + + +def spend_limit_reached(status: int, body: object = None) -> bool: if status == 402: - # OpenRouter's reply once the key's credit limit is spent: the spend - # backstop until we have a ledger of our own. + return True + text = json.dumps(body, default=str) if body is not None else "" + return any(marker in text for marker in SPEND_LIMIT_MARKERS) + + +def http_error_message(status: int, body: object = None) -> str: + """What to tell the user when the model provider answers with an error.""" + if spend_limit_reached(status, body): + # The provider key's credit limit is spent: the spend backstop until + # we have a ledger of our own. return "Data Chat has reached its usage limit for now. Try again later." if status in (401, 403, 404): return "Chat isn't configured on this server yet." @@ -298,6 +320,26 @@ def to_history(messages: list[dict[str, str]]) -> list[ModelMessage]: return history +def log_turn(model: str, outcome: str, usage: dict[str, Any], seconds: float) -> None: + usage_log.info( + "chat_turn outcome=%s model=%s seconds=%.1f requests=%s tool_calls=%s " + "input_tokens=%s output_tokens=%s cost_usd=%s", + outcome, + model, + seconds, + *( + usage.get(key) + for key in ( + "requests", + "tool_calls", + "input_tokens", + "output_tokens", + "cost_usd", + ) + ), + ) + + async def stream_answer( agent: Agent, settings: ChatSettings, @@ -310,10 +352,19 @@ async def stream_answer( tool_calls_limit=settings.max_tool_calls, output_tokens_limit=settings.max_output_tokens, ) + started = time.monotonic() + # "stopped" unless the turn finishes or fails: the reader went away. + outcome = "stopped" + usage: dict[str, Any] = {} try: - async with agent.run_stream_events( - prompt, message_history=history, usage_limits=limits - ) as events: + # A stalled provider would otherwise hold the stream open until the + # proxy gives up; evals saw chats hang for minutes. + async with ( + asyncio.timeout(settings.turn_timeout), + agent.run_stream_events( + prompt, message_history=history, usage_limits=limits + ) as events, + ): async for event in events: if isinstance(event, PartStartEvent) and isinstance( event.part, TextPart @@ -346,42 +397,59 @@ async def stream_answer( if chart: yield {"type": "chart", "id": part.tool_call_id, "chart": chart} elif isinstance(event, AgentRunResultEvent): - usage = event.result.usage - yield { - "type": "done", - "usage": { - "requests": usage.requests, - "tool_calls": usage.tool_calls, - "input_tokens": usage.input_tokens, - "output_tokens": usage.output_tokens, - "cost_usd": turn_cost(event.result.new_messages()), - }, + run_usage = event.result.usage + usage = { + "requests": run_usage.requests, + "tool_calls": run_usage.tool_calls, + "input_tokens": run_usage.input_tokens, + "output_tokens": run_usage.output_tokens, + "cost_usd": turn_cost(event.result.new_messages()), } + outcome = "ok" + yield {"type": "done", "usage": usage} except UsageLimitExceeded: + outcome = "step_limit" yield { "type": "error", "message": "This question needed more steps than one answer allows. " "Try asking something narrower.", } + except TimeoutError: + outcome = "timeout" + yield { + "type": "error", + "message": "That took too long to answer. Try again, or ask " + "something narrower.", + } except UserError as exc: # Usually a missing provider API key or an unknown model name. + outcome = "not_configured" logger.error("Chat model configuration error: %s", exc) yield {"type": "error", "message": "Chat isn't configured on this server yet."} except ModelHTTPError as exc: # The status says what to do: limit reached, busy, misconfigured, or # down. The body can hold provider details, so it isn't logged. + limit = spend_limit_reached(exc.status_code, exc.body) + outcome = "spend_limit" if limit else f"http_{exc.status_code}" logger.error( "Chat model request failed: HTTP %s from %s", exc.status_code, exc.model_name, ) - yield {"type": "error", "message": http_error_message(exc.status_code)} + yield { + "type": "error", + "message": http_error_message(exc.status_code, exc.body), + } except ModelAPIError as exc: + outcome = "provider_error" logger.error("Chat model request failed (%s)", type(exc).__name__) yield { "type": "error", "message": "The language model service is unavailable. Try again shortly.", } except Exception as exc: # noqa: BLE001 -- the stream must end with a readable event + outcome = "error" logger.error("Chat turn failed (%s)", type(exc).__name__) yield {"type": "error", "message": "Something went wrong answering that."} + finally: + log_turn(settings.model, outcome, usage, time.monotonic() - started) diff --git a/backend/chat/config.py b/backend/chat/config.py index 04c908cc..39914b37 100644 --- a/backend/chat/config.py +++ b/backend/chat/config.py @@ -37,6 +37,9 @@ class ChatSettings: max_model_requests: int = 20 max_tool_calls: int = 30 max_output_tokens: int = 40_000 + # Seconds one answer may take before it ends with an error, so a stalled + # provider can't hold a stream open. + turn_timeout: float = 180.0 # Rows of each data tool result the model sees; the browser gets full rows # by reference through /api/chat/data. preview_rows: int = 20 @@ -67,6 +70,7 @@ def from_env(cls, environ: Mapping[str, str] | None = None) -> ChatSettings: max_model_requests=int(env.get("CHAT_MAX_MODEL_REQUESTS", "20")), max_tool_calls=int(env.get("CHAT_MAX_TOOL_CALLS", "30")), max_output_tokens=int(env.get("CHAT_MAX_OUTPUT_TOKENS", "40000")), + turn_timeout=float(env.get("CHAT_TURN_TIMEOUT", "180")), preview_rows=int(env.get("CHAT_PREVIEW_ROWS", "20")), max_history_messages=int(env.get("CHAT_MAX_HISTORY_MESSAGES", "20")), max_message_chars=int(env.get("CHAT_MAX_MESSAGE_CHARS", "4000")), diff --git a/backend/chat/router.py b/backend/chat/router.py index 8374f403..e9e20770 100644 --- a/backend/chat/router.py +++ b/backend/chat/router.py @@ -16,6 +16,7 @@ from __future__ import annotations import json +import logging import threading import time from collections import defaultdict, deque @@ -27,7 +28,7 @@ from pydantic import BaseModel, Field, model_validator from chat import access -from chat.agent import build_agent, stream_answer, to_history, tool_error +from chat.agent import build_agent, stream_answer, to_history, tool_error, usage_log from chat.charts import DataRef, resolve_rows from chat.config import ChatSettings from data_tools.service import DataToolService @@ -119,6 +120,13 @@ def create_router( limiter = RateLimiter(settings.rate_limit) unlock_limiter = RateLimiter(UNLOCK_ATTEMPTS) router = APIRouter(prefix="/chat", tags=["chat"]) + # Uvicorn only shows its own loggers at INFO, so send usage lines to + # stderr ourselves (once, however many routers are built). + if not usage_log.handlers: + handler = logging.StreamHandler() + handler.setFormatter(logging.Formatter("%(asctime)s %(message)s")) + usage_log.addHandler(handler) + usage_log.setLevel(logging.INFO) def unlocked(request: Request) -> bool: code = settings.access_code diff --git a/backend/tests/test_chat.py b/backend/tests/test_chat.py index 19765daf..8c915228 100644 --- a/backend/tests/test_chat.py +++ b/backend/tests/test_chat.py @@ -2,6 +2,7 @@ import asyncio import json +import logging import duckdb import pytest @@ -417,6 +418,74 @@ async def run(): assert phrase in events[-1]["message"] +@pytest.mark.parametrize( + ("status", "body"), + [ + # OpenRouter's key credit limit, when it answers 403 instead of 402. + ( + 403, + { + "message": "Key limit", + "metadata": {"limit_source": "openrouter_key_limit"}, + }, + ), + # OpenAI's project or organization hard limit. + (429, {"code": "project_spend_limit_exceeded", "message": "Limit reached"}), + (429, {"error": {"code": "organization_spend_limit_exceeded"}}), + ], +) +def test_spend_limit_errors_say_the_usage_limit_was_reached(warehouse, status, body): + async def fails(_messages, _info): + raise ModelHTTPError(status_code=status, model_name="some/model", body=body) + yield "" + + settings = ChatSettings(model="demo", warehouse_path=warehouse) + agent = build_agent(settings, DataToolService(warehouse)) + + async def run(): + with agent.override(model=FunctionModel(stream_function=fails)): + return [e async for e in stream_answer(agent, settings, "hi", [])] + + assert "usage limit" in asyncio.run(run())[-1]["message"] + + +def test_a_stalled_model_ends_the_turn_with_a_readable_error(warehouse): + async def stalls(_messages, _info): + await asyncio.sleep(30) + yield "too late" + + settings = ChatSettings(model="demo", warehouse_path=warehouse, turn_timeout=0.05) + agent = build_agent(settings, DataToolService(warehouse)) + + async def run(): + with agent.override(model=FunctionModel(stream_function=stalls)): + return [e async for e in stream_answer(agent, settings, "hi", [])] + + events = asyncio.run(run()) + assert [e["type"] for e in events] == ["error"] + assert "too long" in events[0]["message"] + + +def test_each_turn_logs_usage_but_not_the_question(warehouse, caplog): + settings = ChatSettings(model="demo", warehouse_path=warehouse) + agent = build_agent(settings, DataToolService(warehouse)) + + async def run(): + return [ + e + async for e in stream_answer( + agent, settings, "A question about Winooski", [] + ) + ] + + with caplog.at_level(logging.INFO, logger="chat.usage"): + asyncio.run(run()) + (line,) = [r.getMessage() for r in caplog.records if r.name == "chat.usage"] + assert line.startswith("chat_turn outcome=ok model=demo ") + assert "requests=3 tool_calls=2" in line + assert "Winooski" not in line + + def test_openrouter_models_report_cost_and_skip_data_collecting_providers(warehouse): settings = ChatSettings( model="anthropic/claude-haiku-5.5", warehouse_path=warehouse diff --git a/docs/architecture/data-chat.md b/docs/architecture/data-chat.md index 680a77f7..4e3fb591 100644 --- a/docs/architecture/data-chat.md +++ b/docs/architecture/data-chat.md @@ -33,7 +33,8 @@ Decisions so far: | The data tools, in-process (with previews trimmed for the model) | Server-side chat history and conversation logging | | `render_chart` tool using data references, and `POST /api/chat/data` to resolve them | PDF report blocks | | `/chat` page; chats saved in localStorage like the working report | Session workspace handles that persist across turns | -| Per-turn bounds: output length, tool-loop iterations, maximum history size | | +| Per-turn bounds: output length, tool-loop iterations, seconds, maximum history size | | +| One log line per turn (`chat_turn ...`: outcome, model, tokens, cost; never the text) | | | In-memory per-IP rate limit (same approach as the MCP's; per worker, resets on restart) | | | A dedicated OpenRouter key with a credit limit, which hard-stops requests once spent | | @@ -244,7 +245,9 @@ depth, from the inside out: - **OpenRouter:** an optional credit limit per API key. Once it's spent, that key's requests fail even if the account has credit. The docs give 402 ("insufficient credits") in one place and 403 with `error.metadata.limit_source: "openrouter_key_limit"` in another. A limit can reset on a schedule or never; `GET /api/v1/key` shows what's left. - **OpenAI** (only for a direct `openai:` model): a project or organization hard limit, which must be switched on ("Enforce a hard limit"); spend alerts alone don't block anything. Hitting it returns 429 with `project_spend_limit_exceeded` (or `organization_spend_limit_exceeded`). - Enforcement isn't instant on OpenAI, so spend can slightly exceed the limit; set it a little under what we can afford. - - **Known gap:** the chat only recognizes 402 as "usage limit reached". An OpenRouter 403 key-limit error would read as "isn't configured", and an OpenAI 429 spend-limit error as "busy, try again". Reading `limit_source` or the error code from the response body would fix both. + - The chat reads all three as "usage limit reached": a 402, an OpenRouter + 403 naming `openrouter_key_limit`, and an OpenAI 429 whose code ends in + `spend_limit_exceeded` (`spend_limit_reached` in `chat/agent.py`). 6. **Reconcile** (optional, later): a nightly job compares the ledger with the provider's usage/cost API. When the cap is reached, `GET /api/chat/status` returns `{ "available": false, "resets_at": "2026-11-01" }` and the UI shows: _"Data Chat has reached its monthly usage limit and will be back on November 1."_ diff --git a/justfile b/justfile index 228d8c22..98242c87 100644 --- a/justfile +++ b/justfile @@ -12,7 +12,7 @@ mcp_port := env("MCP_PORT", "6768") # Pass named variables without placing bearer-token values in podman arguments. mcp_container_env := "-e MCP_ENABLED -e MCP_AUTH_MODE -e MCP_BEARER_TOKENS -e MCP_ALLOWED_HOSTS -e MCP_ALLOWED_ORIGINS -e MCP_MAX_ROWS -e MCP_MAX_BYTES -e MCP_QUERY_TIMEOUT -e MCP_MAX_CONCURRENCY -e MCP_RATE_LIMIT -e MCP_MAX_REQUEST_BYTES -e MCP_REQUEST_BODY_TIMEOUT" # Data chat (docs/architecture/data-chat.md); the provider key is passed by name. -chat_container_env := "-e CHAT_ENABLED -e CHAT_ACCESS_CODE -e CHAT_ACCESS_DAYS -e CHAT_MODEL -e OPENROUTER_API_KEY -e OPENAI_API_KEY -e CHAT_MAX_MODEL_REQUESTS -e CHAT_MAX_TOOL_CALLS -e CHAT_MAX_OUTPUT_TOKENS -e CHAT_PREVIEW_ROWS -e CHAT_MAX_HISTORY_MESSAGES -e CHAT_MAX_MESSAGE_CHARS -e CHAT_MAX_HISTORY_CHARS -e CHAT_RATE_LIMIT -e CHAT_CHART_MAX_ROWS" +chat_container_env := "-e CHAT_ENABLED -e CHAT_ACCESS_CODE -e CHAT_ACCESS_DAYS -e CHAT_MODEL -e OPENROUTER_API_KEY -e OPENAI_API_KEY -e CHAT_MAX_MODEL_REQUESTS -e CHAT_MAX_TOOL_CALLS -e CHAT_MAX_OUTPUT_TOKENS -e CHAT_TURN_TIMEOUT -e CHAT_PREVIEW_ROWS -e CHAT_MAX_HISTORY_MESSAGES -e CHAT_MAX_MESSAGE_CHARS -e CHAT_MAX_HISTORY_CHARS -e CHAT_RATE_LIMIT -e CHAT_CHART_MAX_ROWS" # Host-specific podman flags. Empty by default (Docker, macOS, rootful podman). # On the VM with rootless podman, override: just podman_flags="--userns=keep-id:uid=1000,gid=1000" podman_flags := env("PODMAN_FLAGS", "")