From 9e2caeac48aaf62b5abdc686b4f03d591fbb422c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jaroslav=20Bene=C5=A1?= Date: Sat, 1 Aug 2026 00:36:12 +0200 Subject: [PATCH] Ask the endpoint what a streamed reply cost A streamed completion carries no token counts unless you ask for them, and `stream_options: {include_usage: true}` is how. Not every server implements it, and an unknown key is a 400 from some -- the same hazard as sending a tools array to an endpoint without support. So it is asked for once per base URL per process, and an endpoint that refuses is remembered and retried without it. The retry is safe because the status is checked before a single line is read: nothing has been yielded, so there is nothing to duplicate. chunk_usage() reads the resulting chunk. It needed no change to the loop above it: a usage chunk carries `choices: []`, which is exactly the shape delta_text, delta_reasoning, delta_tool_calls and finish_reason have always returned early on. All-zero counts are treated as absent, because some servers attach zeros to every chunk and the real numbers only at the end. services/tokens.py is the fallback for endpoints that never report: four characters to a token, counting the tools array because thirteen schemas is a meaningful slice of a short window, and counting nothing for an image because its cost depends on tiling and an invented number would be worse than the omission. Crude on purpose -- a real tokeniser means one per model family, for a figure that is displayed beside a tilde. Nothing uses any of this yet. Co-Authored-By: Claude Opus 5 (1M context) --- src/lembas/services/llm/openai_client.py | 76 ++++++++++++ src/lembas/services/tokens.py | 75 ++++++++++++ tests/test_metrics.py | 140 ++++++++++++++++++++++- 3 files changed, 290 insertions(+), 1 deletion(-) create mode 100644 src/lembas/services/tokens.py diff --git a/src/lembas/services/llm/openai_client.py b/src/lembas/services/llm/openai_client.py index f78878e..70c7c4b 100644 --- a/src/lembas/services/llm/openai_client.py +++ b/src/lembas/services/llm/openai_client.py @@ -190,6 +190,12 @@ def context_from(entry: dict[str, Any]) -> int: return 0 +# Endpoints that rejected `stream_options`, so it is asked for once per base URL +# per process and then never again. Not persisted: it is a property of whatever +# is running there now, and a restart is the right time to find out afresh. +_NO_STREAM_OPTIONS: set[str] = set() + + async def stream_chat( endpoint: Endpoint, payload: dict[str, Any], @@ -198,8 +204,41 @@ async def stream_chat( Yields the raw upstream chunks; interpreting them is the caller's job. The terminating "[DONE]" sentinel is consumed here and not yielded. + + `stream_options` asks for the final usage chunk, which is the only way to + learn what a streamed reply actually cost. Not every server implements it, + and an unknown key is a 400 from some of them -- the same hazard as sending + a `tools` array to an endpoint without support. So it is asked for once, + and an endpoint that refuses is remembered and never asked again. Retrying + is safe because the status is checked before a single line is read: nothing + has been yielded, so there is nothing to duplicate. """ + wants_usage = endpoint.base_url not in _NO_STREAM_OPTIONS + + try: + async for chunk in _stream_once(endpoint, payload, usage=wants_usage): + yield chunk + except LLMError as exc: + if not wants_usage or exc.status_code not in (400, 422): + raise + _NO_STREAM_OPTIONS.add(endpoint.base_url) + log.info( + "%s rejected stream_options; token counts will be estimated there", + endpoint.base_url, + ) + async for chunk in _stream_once(endpoint, payload, usage=False): + yield chunk + + +async def _stream_once( + endpoint: Endpoint, + payload: dict[str, Any], + *, + usage: bool, +) -> AsyncIterator[dict[str, Any]]: body = {**payload, "stream": True} + if usage: + body["stream_options"] = {"include_usage": True} try: async with ( @@ -318,6 +357,43 @@ def finish_reason(chunk: dict[str, Any]) -> str: return "" +def chunk_usage(chunk: dict[str, Any]) -> dict[str, int] | None: + """Token counts from a usage chunk, or None if this is not one. + + A usage chunk carries `choices: []`, which is exactly the shape delta_text, + delta_reasoning, delta_tool_calls and finish_reason all return early on -- + they have always tolerated it, so nothing else needs to change to let one + through. + + Fields are read defensively because "the endpoint returned something odd" + must never be the reason a reply fails; a bad shape simply means no counts. + """ + try: + raw = chunk.get("usage") + except AttributeError: + return None + if not isinstance(raw, dict): + return None + + counts: dict[str, int] = {} + for key in ("prompt_tokens", "completion_tokens", "total_tokens"): + value = raw.get(key) + if isinstance(value, bool) or not isinstance(value, (int, float)): + continue + if value >= 0: + counts[key] = int(value) + + # Some servers send a usage object of zeros on every chunk and the real + # numbers only at the end. All-zero is indistinguishable from that, and + # treating it as an answer would freeze the count at nothing. + if not counts or not any(counts.values()): + return None + counts.setdefault( + "total_tokens", counts.get("prompt_tokens", 0) + counts.get("completion_tokens", 0) + ) + return counts + + def delta_text(chunk: dict[str, Any]) -> str: """Pull the text out of one streamed chunk, tolerating provider variation.""" try: diff --git a/src/lembas/services/tokens.py b/src/lembas/services/tokens.py new file mode 100644 index 0000000..6e60342 --- /dev/null +++ b/src/lembas/services/tokens.py @@ -0,0 +1,75 @@ +"""A rough token count, for when the endpoint does not give a real one. + +Deliberately crude. Counting tokens properly means the model's own tokeniser, +which means shipping one per model family and a dependency that has to be kept +in step with them -- for a number that is displayed beside a `~` and used to +decide when to summarise. + +Four characters per token is the usual English approximation. It is optimistic +on code and badly wrong on CJK, which is why anything derived from it is marked +as an estimate everywhere it is shown. +""" + +from __future__ import annotations + +from typing import Any + +CHARS_PER_TOKEN = 4 + +# What a message costs beyond its text: the role, the delimiters a chat template +# wraps each turn in. Small, but a hundred-turn conversation is a hundred of them. +PER_MESSAGE_OVERHEAD = 4 + + +def estimate(text: str) -> int: + """Roughly how many tokens a piece of text is.""" + if not text: + return 0 + return max(1, round(len(text) / CHARS_PER_TOKEN)) + + +def estimate_content(content: Any) -> int: + """A message's content, whether it is a plain string or typed parts. + + An image part contributes nothing: its cost depends on the model's tiling, + and a number invented here would be worse than the omission. + """ + if isinstance(content, str): + return estimate(content) + if isinstance(content, list): + return sum( + estimate(part.get("text", "")) + for part in content + if isinstance(part, dict) and part.get("type") == "text" + ) + return 0 + + +def estimate_messages(messages: list[dict[str, Any]]) -> int: + total = 0 + for message in messages: + if not isinstance(message, dict): + continue + total += PER_MESSAGE_OVERHEAD + estimate_content(message.get("content")) + # A tool result the model reads back is part of the window too. + for call in message.get("tool_calls") or []: + function = (call or {}).get("function") or {} + total += estimate(str(function.get("name", ""))) + total += estimate(str(function.get("arguments", ""))) + return total + + +def estimate_request(payload: dict[str, Any]) -> int: + """What a whole request body costs, tools included. + + The tools array is sent on every request when tools are offered and is not + small -- thirteen schemas is a meaningful slice of a short window, and + leaving it out would make the percentage read low exactly when it matters. + """ + total = estimate_messages(payload.get("messages") or []) + for tool in payload.get("tools") or []: + function = (tool or {}).get("function") or {} + total += estimate(str(function.get("name", ""))) + total += estimate(str(function.get("description", ""))) + total += estimate(str(function.get("parameters", ""))) + return total diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 0d669cc..ecb855c 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -6,8 +6,9 @@ from fastapi.testclient import TestClient from sqlalchemy import select from lembas.db.models import Connection, Model +from lembas.services import tokens from lembas.services.crypto import encrypt -from lembas.services.llm.openai_client import context_from +from lembas.services.llm.openai_client import chunk_usage, context_from def _model(db, **kwargs) -> Model: @@ -120,3 +121,140 @@ def test_junk_in_the_context_length_field_is_ignored_not_a_500( assert response.status_code == 303 db.refresh(model) assert model.context_length == 4096 + + +# --- Usage off the wire ------------------------------------------------------- +def test_usage_is_read_from_a_usage_chunk(): + chunk = { + "choices": [], + "usage": {"prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120}, + } + assert chunk_usage(chunk) == { + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + } + + +def test_a_missing_total_is_worked_out(): + chunk = {"choices": [], "usage": {"prompt_tokens": 100, "completion_tokens": 20}} + assert chunk_usage(chunk)["total_tokens"] == 120 + + +def test_an_ordinary_chunk_carries_no_usage(): + assert chunk_usage({"choices": [{"delta": {"content": "hi"}}]}) is None + assert chunk_usage({}) is None + assert chunk_usage({"usage": "lots"}) is None + + +def test_an_all_zero_usage_object_is_not_an_answer(): + """Some servers attach zeros to every chunk and the real numbers only at the + end. Believing the zeros freezes the count at nothing.""" + chunk = {"choices": [], "usage": {"prompt_tokens": 0, "completion_tokens": 0}} + assert chunk_usage(chunk) is None + + +def test_the_other_accessors_still_ignore_a_usage_chunk(): + """They return early on `choices: []`, which is exactly the shape of one. + That is what lets a usage chunk through the loop untouched.""" + from lembas.services.llm.openai_client import ( + delta_reasoning, + delta_text, + delta_tool_calls, + finish_reason, + ) + + chunk = {"choices": [], "usage": {"prompt_tokens": 1, "completion_tokens": 1}} + assert delta_text(chunk) == "" + assert delta_reasoning(chunk) == "" + assert delta_tool_calls(chunk) == [] + assert finish_reason(chunk) == "" + + +async def test_stream_options_is_asked_for(mock_http): + import json as json_module + + import httpx + + from lembas.services.llm.openai_client import Endpoint, stream_chat + + seen: list[dict] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(json_module.loads(request.content)) + return httpx.Response(200, text="data: [DONE]\n\n") + + mock_http(handler) + async for _ in stream_chat(Endpoint("http://ask.test", "", {}), {"model": "m"}): + pass + + assert seen[0]["stream_options"] == {"include_usage": True} + + +async def test_an_endpoint_that_rejects_stream_options_is_asked_once(mock_http): + """A 400 for an unknown key is the same hazard as sending `tools` to an + endpoint without support. Retry without it, then stop asking.""" + import json as json_module + + import httpx + + from lembas.services.llm.openai_client import ( + _NO_STREAM_OPTIONS, + Endpoint, + stream_chat, + ) + + _NO_STREAM_OPTIONS.discard("http://fussy.test") + seen: list[dict] = [] + + def handler(request: httpx.Request) -> httpx.Response: + body = json_module.loads(request.content) + seen.append(body) + if "stream_options" in body: + return httpx.Response(400, json={"error": {"message": "unknown field"}}) + reply = 'data: {"choices":[{"delta":{"content":"hi"}}]}\n\ndata: [DONE]\n\n' + return httpx.Response(200, text=reply) + + mock_http(handler) + endpoint = Endpoint("http://fussy.test", "", {}) + + text = [c async for c in stream_chat(endpoint, {"model": "m"})] + assert text, "the retry should have produced the reply" + assert len(seen) == 2 + + # Second reply: it already knows, so one request and no stream_options. + async for _ in stream_chat(endpoint, {"model": "m"}): + pass + assert len(seen) == 3 + assert "stream_options" not in seen[2] + _NO_STREAM_OPTIONS.discard("http://fussy.test") + + +# --- The estimate ------------------------------------------------------------- +def test_the_estimate_is_about_four_characters_a_token(): + assert tokens.estimate("") == 0 + assert tokens.estimate("x" * 400) == 100 + + +def test_typed_content_parts_are_counted_and_images_are_not(): + """An image's cost depends on the model's tiling. A number invented here + would be worse than the omission.""" + content = [ + {"type": "text", "text": "x" * 40}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA" * 500}}, + ] + assert tokens.estimate_content(content) == 10 + + +def test_a_request_estimate_includes_the_tools_array(): + """Thirteen schemas is a meaningful slice of a short window; leaving them + out would read low exactly when it matters.""" + payload = { + "messages": [{"role": "user", "content": "x" * 40}], + "tools": [ + {"function": {"name": "web_search", "description": "y" * 400, "parameters": {}}} + ], + } + assert tokens.estimate_request(payload) > tokens.estimate_request( + {"messages": payload["messages"]} + )