Jaroslav Beneš
9e2caeac48
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>
261 lines
8.9 KiB
Python
261 lines
8.9 KiB
Python
"""How full the context is, what a reply cost, and how fast it arrived."""
|
|
|
|
from __future__ import annotations
|
|
|
|
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 chunk_usage, context_from
|
|
|
|
|
|
def _model(db, **kwargs) -> Model:
|
|
connection = Connection(
|
|
name="Test", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt("")
|
|
)
|
|
db.add(connection)
|
|
db.commit()
|
|
model = Model(connection_id=connection.id, model_id="test-model", **kwargs)
|
|
db.add(model)
|
|
db.commit()
|
|
return model
|
|
|
|
|
|
# --- Reading a context length off /v1/models ---------------------------------
|
|
def test_context_length_is_read_from_any_of_the_spellings():
|
|
assert context_from({"id": "m", "context_length": 8192}) == 8192
|
|
assert context_from({"id": "m", "max_model_len": 32768}) == 32768
|
|
assert context_from({"id": "m", "context_window": 4096}) == 4096
|
|
assert context_from({"id": "m", "meta": {"n_ctx": 2048}}) == 2048
|
|
|
|
|
|
def test_a_quoted_number_is_accepted_but_a_label_is_not():
|
|
"""Some servers quote it. "8192 tokens" is a label, not a measurement."""
|
|
assert context_from({"id": "m", "context_length": "8192"}) == 8192
|
|
assert context_from({"id": "m", "context_length": "8192 tokens"}) == 0
|
|
|
|
|
|
def test_an_absent_or_implausible_context_length_is_zero():
|
|
assert context_from({"id": "m"}) == 0
|
|
assert context_from({"id": "m", "context_length": 0}) == 0
|
|
assert context_from({"id": "m", "context_length": 64}) == 0
|
|
assert context_from({"id": "m", "context_length": 10**9}) == 0
|
|
# True is an int in Python, and it is not a context length.
|
|
assert context_from({"id": "m", "context_length": True}) == 0
|
|
|
|
|
|
# --- Discovery ---------------------------------------------------------------
|
|
async def test_discovery_fills_in_a_context_length(client: TestClient, db, registered, mock_http):
|
|
import httpx
|
|
|
|
mock_http(
|
|
lambda _r: httpx.Response(
|
|
200, json={"data": [{"id": "big-model", "context_length": 16384}]}
|
|
)
|
|
)
|
|
client.post(
|
|
"/admin/connections",
|
|
data={"name": "Local", "base_url": "http://x.test", "api_key": ""},
|
|
follow_redirects=False,
|
|
)
|
|
model = db.scalar(select(Model).where(Model.model_id == "big-model"))
|
|
assert model.context_length == 16384
|
|
|
|
|
|
async def test_discovery_never_overwrites_a_number_an_admin_typed(
|
|
client: TestClient, db, registered, mock_http
|
|
):
|
|
"""A refresh must not undo a correction. Administrators set this precisely
|
|
because the endpoint was wrong or silent."""
|
|
import httpx
|
|
|
|
mock_http(
|
|
lambda _r: httpx.Response(200, json={"data": [{"id": "m", "context_length": 4096}]})
|
|
)
|
|
client.post(
|
|
"/admin/connections",
|
|
data={"name": "Local", "base_url": "http://x.test", "api_key": ""},
|
|
follow_redirects=False,
|
|
)
|
|
model = db.scalar(select(Model).where(Model.model_id == "m"))
|
|
model.context_length = 131072
|
|
db.commit()
|
|
|
|
connection = db.scalar(select(Connection))
|
|
client.post(f"/admin/connections/{connection.id}/refresh", follow_redirects=False)
|
|
|
|
db.refresh(model)
|
|
assert model.context_length == 131072
|
|
|
|
|
|
# --- The admin field ---------------------------------------------------------
|
|
def test_an_admin_can_set_and_clear_the_context_length(client: TestClient, db, registered):
|
|
model = _model(db)
|
|
client.post(
|
|
f"/admin/models/{model.id}",
|
|
data={"context_length": "8192", "position": ""},
|
|
follow_redirects=False,
|
|
)
|
|
db.refresh(model)
|
|
assert model.context_length == 8192
|
|
|
|
client.post(
|
|
f"/admin/models/{model.id}", data={"context_length": "", "position": ""},
|
|
follow_redirects=False,
|
|
)
|
|
db.refresh(model)
|
|
assert model.context_length == 0
|
|
|
|
|
|
def test_junk_in_the_context_length_field_is_ignored_not_a_500(
|
|
client: TestClient, db, registered
|
|
):
|
|
model = _model(db, context_length=4096)
|
|
response = client.post(
|
|
f"/admin/models/{model.id}",
|
|
data={"context_length": "eight thousand", "position": ""},
|
|
follow_redirects=False,
|
|
)
|
|
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"]}
|
|
)
|