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) <noreply@anthropic.com>
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user