diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index 5f86d48a4..6d0f31606 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -27,6 +27,7 @@ if TYPE_CHECKING: import httpx from headroom.copilot_auth import apply_copilot_api_auth +from headroom.proxy.stream_output_tokens import estimate_output_tokens logger = logging.getLogger("headroom.proxy") @@ -952,11 +953,27 @@ class StreamingMixin: output_tokens = stream_state["output_tokens"] output_tokens_source = "provider" if output_tokens is None: - output_tokens = stream_state["total_bytes"] // 40 - output_tokens_source = "estimated_bytes" + # No usage chunk from the upstream. Count the stream's OWN TEXT + # rather than dividing the raw wire by a constant: `total_bytes` + # includes every `data:` prefix, JSON envelope and framing newline, + # so its error tracked how chattily the answer was chunked instead + # of how long the answer was. Falls back to the byte heuristic only + # when no text could be recovered at all. + output_tokens, output_tokens_source = estimate_output_tokens( + sse_text=full_sse_data, + total_bytes=stream_state["total_bytes"], + ) + # Name the actual basis. The old message always said "from N bytes" + # even though that is now only true for the fallback rung, and an + # operator reading it needs to know which estimate they are looking + # at before trusting the number. + basis = ( + "counted from stream text" + if output_tokens_source == "estimated_text" + else f"estimated from {stream_state['total_bytes']} raw SSE bytes" + ) logger.warning( - f"[{request_id}] Could not parse output_tokens from SSE, " - f"estimating {output_tokens} from {stream_state['total_bytes']} bytes" + f"[{request_id}] No usage chunk in SSE; output_tokens={output_tokens} ({basis})" ) outcome_tags = dict(tags or {}) diff --git a/headroom/proxy/stream_output_tokens.py b/headroom/proxy/stream_output_tokens.py new file mode 100644 index 000000000..96e410581 --- /dev/null +++ b/headroom/proxy/stream_output_tokens.py @@ -0,0 +1,141 @@ +"""Recover output-token counts from a finished SSE stream. + +When an upstream omits a usage chunk, the proxy has to estimate how many output +tokens the turn produced. The estimate was ``total_bytes // 40`` over the RAW +SSE WIRE — every ``data:`` prefix, every JSON envelope, every ``role``/ +``finish_reason``/``id``/``model`` field, blank-line framing included. On a +Copilot chat turn that produced a short answer the log read: + + Could not parse output_tokens from SSE, estimating 8 from 334 bytes + +334 bytes of wire is mostly envelope; the generated text inside it was a +fraction of that. The divisor 40 is a fudge for "bytes per token INCLUDING +framing overhead", so its error scales with how chatty the framing is rather +than with the answer — a stream split into many small deltas is punished for +the split, and one delivered in a few large chunks is not. + +The stream's own text is right there in the buffer, so extract it and count +that instead. ``bytes // 40`` survives only as the last resort for a stream +whose text could not be recovered at all. + +Pure and I/O-free so the parsing is testable without a proxy or a network. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterator +from typing import Any + +# Bytes per token when nothing better is available. Applied to the raw wire, so +# it must absorb SSE framing overhead as well as the text — which is exactly why +# it is a poor estimator and a last resort. +WIRE_BYTES_PER_TOKEN = 40 + +# Characters per token for extracted text. ~4 is the usual English/code +# approximation and is applied to generated text ONLY, with no framing in it. +TEXT_CHARS_PER_TOKEN = 4 + + +def _iter_sse_payloads(sse: str) -> Iterator[Any]: + """Yield each ``data:`` payload in an SSE stream as a parsed object. + + Tolerates the two things real streams do that a naive split does not: + an event whose ``data:`` is spread over multiple lines, and ``[DONE]``. + """ + for block in sse.split("\n\n"): + lines = [ln for ln in block.split("\n") if ln.startswith("data:")] + if not lines: + continue + # Multi-line data: fields concatenate, per the SSE spec. + raw = "".join(ln[5:].lstrip() for ln in lines).strip() + if not raw or raw == "[DONE]": + continue + try: + yield json.loads(raw) + except (ValueError, TypeError): + continue + + +def extract_stream_text(sse: str) -> str: + """Return the assistant text carried by a completed SSE stream. + + Handles the three surfaces this proxy forwards: + + * OpenAI chat completions — ``choices[].delta.content`` + * OpenAI responses — ``response.output_text.delta`` / ``delta`` + * Anthropic messages — ``content_block_delta.delta.text`` + + Reasoning/thinking deltas are counted too: the provider bills them as + output tokens, so omitting them would under-count exactly the turns where + output is most expensive. + """ + if not sse: + return "" + + parts: list[str] = [] + for obj in _iter_sse_payloads(sse): + if not isinstance(obj, dict): + continue + + # --- OpenAI chat completions -------------------------------------- # + choices = obj.get("choices") + if isinstance(choices, list): + for choice in choices: + if not isinstance(choice, dict): + continue + delta = choice.get("delta") + if not isinstance(delta, dict): + continue + for key in ("content", "reasoning_content", "refusal"): + value = delta.get(key) + if isinstance(value, str): + parts.append(value) + # Tool-call arguments stream as text and are billed as output. + tool_calls = delta.get("tool_calls") + if isinstance(tool_calls, list): + for call in tool_calls: + fn = call.get("function") if isinstance(call, dict) else None + args = fn.get("arguments") if isinstance(fn, dict) else None + if isinstance(args, str): + parts.append(args) + continue + + obj_type = obj.get("type") + + # --- Anthropic messages ------------------------------------------- # + if obj_type == "content_block_delta": + delta = obj.get("delta") + if isinstance(delta, dict): + for key in ("text", "thinking", "partial_json"): + value = delta.get(key) + if isinstance(value, str): + parts.append(value) + continue + + # --- OpenAI responses --------------------------------------------- # + if isinstance(obj_type, str) and obj_type.endswith(".delta"): + value = obj.get("delta") + if isinstance(value, str): + parts.append(value) + continue + + return "".join(parts) + + +def estimate_output_tokens(*, sse_text: str, total_bytes: int) -> tuple[int, str]: + """Return ``(tokens, source)`` for a stream with no provider usage chunk. + + ``source`` names which rung of the ladder produced the number so the caller + can log it honestly rather than implying the provider reported it: + + * ``estimated_text`` — counted from the generated text (good) + * ``estimated_bytes`` — the raw-wire fallback (poor, last resort) + """ + text = extract_stream_text(sse_text) + if text: + # At least one token for any non-empty answer: integer division would + # report 0 for a 1-3 character reply ("OK", "42"), and a turn that + # produced output must never be recorded as having produced none. + return max(1, len(text) // TEXT_CHARS_PER_TOKEN), "estimated_text" + return max(0, total_bytes) // WIRE_BYTES_PER_TOKEN, "estimated_bytes" diff --git a/tests/test_stream_output_tokens.py b/tests/test_stream_output_tokens.py new file mode 100644 index 000000000..0c4dbafe0 --- /dev/null +++ b/tests/test_stream_output_tokens.py @@ -0,0 +1,171 @@ +"""Output tokens must be counted from the stream's text, not its wire size. + +When an upstream sends no usage chunk, the proxy estimated output tokens as +``total_bytes // 40`` over the RAW SSE WIRE — ``data:`` prefixes, JSON +envelopes, ``role``/``finish_reason``/``id``/``model`` fields and blank-line +framing all included. From a field log (Copilot Chat, 0.36.x): + + Could not parse output_tokens from SSE, estimating 8 from 334 bytes + +The divisor is a fudge for "bytes per token including framing", so the error +tracked how chattily the answer was chunked rather than how long it was: the +same answer split into more deltas scores higher purely for being split. + +GitHub's Copilot CAPI is one of the upstreams that omits the usage chunk, so +this was every Copilot turn's output number — and output tokens feed the +output-shaping savings estimate and the cost model. +""" + +from __future__ import annotations + +import json + +import pytest + +from headroom.proxy.stream_output_tokens import ( + TEXT_CHARS_PER_TOKEN, + WIRE_BYTES_PER_TOKEN, + estimate_output_tokens, + extract_stream_text, +) + + +def _sse(*objs: dict, done: bool = True) -> str: + out = "".join(f"data: {json.dumps(o)}\n\n" for o in objs) + return out + ("data: [DONE]\n\n" if done else "") + + +def _chat_delta(text: str) -> dict: + return {"choices": [{"index": 0, "delta": {"content": text}}]} + + +# --------------------------------------------------------------------------- # +# Extraction, per surface +# --------------------------------------------------------------------------- # +def test_openai_chat_deltas() -> None: + sse = _sse( + {"choices": [{"delta": {"role": "assistant"}}]}, + _chat_delta("Hello "), + _chat_delta("world"), + {"choices": [{"delta": {}, "finish_reason": "stop"}]}, + ) + assert extract_stream_text(sse) == "Hello world" + + +def test_anthropic_content_block_deltas() -> None: + sse = _sse( + {"type": "message_start", "message": {"id": "msg_1"}}, + {"type": "content_block_delta", "delta": {"type": "text_delta", "text": "abc"}}, + {"type": "content_block_delta", "delta": {"type": "text_delta", "text": "def"}}, + {"type": "message_stop"}, + done=False, + ) + assert extract_stream_text(sse) == "abcdef" + + +def test_openai_responses_deltas() -> None: + sse = _sse( + {"type": "response.output_text.delta", "delta": "part one "}, + {"type": "response.output_text.delta", "delta": "part two"}, + ) + assert extract_stream_text(sse) == "part one part two" + + +def test_reasoning_and_tool_arguments_are_billed_output_too() -> None: + """Omitting these under-counts exactly the most expensive turns.""" + sse = _sse( + {"choices": [{"delta": {"reasoning_content": "thinking hard"}}]}, + {"choices": [{"delta": {"tool_calls": [{"function": {"arguments": '{"path":"a.py"}'}}]}}]}, + ) + text = extract_stream_text(sse) + assert "thinking hard" in text + assert '{"path":"a.py"}' in text + + +def test_anthropic_thinking_and_partial_json() -> None: + sse = _sse( + {"type": "content_block_delta", "delta": {"thinking": "plan"}}, + {"type": "content_block_delta", "delta": {"partial_json": '{"a":1}'}}, + done=False, + ) + assert extract_stream_text(sse) == 'plan{"a":1}' + + +# --------------------------------------------------------------------------- # +# Malformed input must never raise — this runs on the response path +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize( + "sse", + [ + "", + "data: not-json\n\n", + "data: [DONE]\n\n", + "garbage without data prefix\n\n", + 'data: {"choices": "not-a-list"}\n\n', + 'data: {"choices": [null]}\n\n', + 'data: {"choices": [{"delta": null}]}\n\n', + 'data: {"choices": [{"delta": {"content": 42}}]}\n\n', + 'data: {"type": "content_block_delta", "delta": "not-a-dict"}\n\n', + "data:\n\n", + ], +) +def test_malformed_streams_yield_empty_not_an_exception(sse: str) -> None: + assert extract_stream_text(sse) == "" + + +def test_multi_line_data_fields_concatenate() -> None: + """Per the SSE spec, and real streams do it.""" + payload = json.dumps(_chat_delta("joined")) + half = len(payload) // 2 + sse = f"data: {payload[:half]}\ndata: {payload[half:]}\n\n" + assert extract_stream_text(sse) == "joined" + + +# --------------------------------------------------------------------------- # +# The estimate itself +# --------------------------------------------------------------------------- # +def test_text_beats_the_wire_heuristic_on_the_reported_shape() -> None: + """A chunky stream: framing dominates the wire, so bytes//40 misreads it.""" + sse = _sse(*[_chat_delta(w) for w in ("The ", "quick ", "brown ", "fox ", "jumps")]) + total_bytes = len(sse.encode()) + + tokens, source = estimate_output_tokens(sse_text=sse, total_bytes=total_bytes) + + assert source == "estimated_text" + # "The quick brown fox jumps" is 25 chars -> 6 tokens. + assert tokens == len("The quick brown fox jumps") // TEXT_CHARS_PER_TOKEN + # The old estimator scored this stream far higher purely for its framing. + assert total_bytes // WIRE_BYTES_PER_TOKEN > tokens + + +def test_chunking_no_longer_changes_the_answer() -> None: + """Same text, different delta split — the count must not move.""" + text = "identical content across both streams" + one = _sse(_chat_delta(text)) + many = _sse(*[_chat_delta(c) for c in text]) + + a, _ = estimate_output_tokens(sse_text=one, total_bytes=len(one.encode())) + b, _ = estimate_output_tokens(sse_text=many, total_bytes=len(many.encode())) + + assert a == b + # And the wire-based estimator would have disagreed wildly. + assert len(one.encode()) // WIRE_BYTES_PER_TOKEN != len(many.encode()) // WIRE_BYTES_PER_TOKEN + + +def test_a_short_answer_is_never_recorded_as_zero() -> None: + sse = _sse(_chat_delta("OK")) + tokens, source = estimate_output_tokens(sse_text=sse, total_bytes=len(sse.encode())) + assert source == "estimated_text" + assert tokens == 1 + + +def test_falls_back_to_bytes_when_no_text_is_recoverable() -> None: + """The upstream-error path reaches here with no stream text at all.""" + tokens, source = estimate_output_tokens(sse_text="", total_bytes=800) + assert source == "estimated_bytes" + assert tokens == 800 // WIRE_BYTES_PER_TOKEN + + +def test_negative_or_zero_bytes_are_safe() -> None: + assert estimate_output_tokens(sse_text="", total_bytes=0) == (0, "estimated_bytes") + assert estimate_output_tokens(sse_text="", total_bytes=-5) == (0, "estimated_bytes")