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"]} + )