Jaroslav Beneš
e185edc9e1
Tokens, how full the context is, and tokens per second -- as chips under each assistant bubble, updating while the reply streams and still there when it finishes. The numbers come from one Metrics object built either from the generation still being written or from the row it left behind. That is the point rather than tidiness: the finished bubble is re-rendered from the database the instant the stream ends, so two code paths would make the figures visibly jump at exactly the moment someone is watching them. Here the only thing that changes is that an estimate may become exact. Message.usage_json has existed and been dead since the schema was written. It is the store. Two counts that look like one. prompt and completion are summed across tool rounds -- what the reply cost. context_tokens is overwritten each round with that round's prompt plus completion -- what the window actually holds. A three-round reply pays for its prompt three times and only ever occupies the window once, so a single number would be wrong for one of the two questions. Generation gains started_at as a field rather than a local in _run, because _follow is a different function that sees only the Generation and otherwise has nothing to compute a live speed against. It also carries a prompt estimate taken before the first chunk, since real usage arrives in one chunk at the very end and a percentage that appears only after the reply is useless. Everything is marked with a tilde when the endpoint reported nothing, and the percentage is simply absent when no context length is set: unknown has to stay tellable from small, and a percentage of an unknown total is a made-up number in a place people trust numbers. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
453 lines
16 KiB
Python
453 lines
16 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 metrics, 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"]}
|
|
)
|
|
|
|
|
|
# --- The Metrics object -------------------------------------------------------
|
|
def test_the_percentage_is_zero_when_nobody_said_how_big_the_window_is():
|
|
"""Unknown must stay tellable from small. A percentage of an unknown total
|
|
is a made-up number in a place people trust numbers."""
|
|
assert metrics.Metrics(context_tokens=5000, context_limit=0).percent == 0
|
|
assert metrics.Metrics(context_tokens=5000, context_limit=10000).percent == 50
|
|
|
|
|
|
def test_pressure_marks_the_bar_only_when_it_is_worth_noticing():
|
|
assert metrics.Metrics(context_tokens=50, context_limit=100).pressure == ""
|
|
assert metrics.Metrics(context_tokens=85, context_limit=100).pressure == "warning"
|
|
assert metrics.Metrics(context_tokens=96, context_limit=100).pressure == "danger"
|
|
assert metrics.Metrics(context_tokens=0, context_limit=0).pressure == ""
|
|
|
|
|
|
def test_speed_needs_both_a_count_and_a_clock():
|
|
assert metrics.Metrics(completion_tokens=100, elapsed_ms=2000).tokens_per_second == 50.0
|
|
assert metrics.Metrics(completion_tokens=100, elapsed_ms=0).tokens_per_second == 0.0
|
|
assert metrics.Metrics(completion_tokens=0, elapsed_ms=2000).tokens_per_second == 0.0
|
|
|
|
|
|
def test_metrics_survive_a_round_trip_through_the_row():
|
|
original = metrics.Metrics(
|
|
prompt_tokens=100,
|
|
completion_tokens=20,
|
|
total_tokens=120,
|
|
context_tokens=120,
|
|
context_limit=8192,
|
|
estimated=True,
|
|
elapsed_ms=1500,
|
|
rounds=2,
|
|
)
|
|
assert metrics.from_message(metrics.to_json(original)) == original
|
|
|
|
|
|
def test_a_row_with_no_usage_reads_as_nothing_rather_than_failing():
|
|
assert metrics.from_message(None).has_anything is False
|
|
assert metrics.from_message({}).has_anything is False
|
|
assert metrics.from_message({"prompt_tokens": "lots"}).prompt_tokens == 0
|
|
|
|
|
|
# --- Through a generation -----------------------------------------------------
|
|
def _chunks(*frames: str) -> str:
|
|
return "".join(f"data: {frame}\n\n" for frame in frames) + "data: [DONE]\n\n"
|
|
|
|
|
|
async def test_usage_is_summed_across_tool_rounds(db, registered, make_chat, monkeypatch):
|
|
"""Prompt and completion are what the reply cost; context_tokens is what the
|
|
window holds. A three-round reply pays for its prompt three times and only
|
|
ever occupies the window once."""
|
|
from lembas.services import generation as generation_service
|
|
|
|
generation = generation_service.Generation(chat_id="c", message_id="m")
|
|
for prompt_tokens, completion_tokens in ((100, 10), (250, 20)):
|
|
counts = {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens}
|
|
generation.prompt_tokens += counts["prompt_tokens"]
|
|
generation.completion_tokens += counts["completion_tokens"]
|
|
generation.context_tokens = prompt_tokens + completion_tokens
|
|
|
|
assert generation.prompt_tokens == 350
|
|
assert generation.completion_tokens == 30
|
|
assert generation.context_tokens == 270
|
|
|
|
|
|
async def test_a_reply_records_what_it_cost(client: TestClient, db, registered, mock_http):
|
|
import httpx
|
|
|
|
from lembas.db.models import Message
|
|
|
|
def handler(_request: httpx.Request) -> httpx.Response:
|
|
if _request.url.path.endswith("/models"):
|
|
return httpx.Response(200, json={"data": [{"id": "m", "context_length": 1000}]})
|
|
return httpx.Response(
|
|
200,
|
|
text=_chunks(
|
|
'{"choices":[{"delta":{"content":"Waybread."}}]}',
|
|
'{"choices":[],"usage":{"prompt_tokens":120,"completion_tokens":8}}',
|
|
),
|
|
)
|
|
|
|
mock_http(handler)
|
|
client.post(
|
|
"/admin/connections",
|
|
data={"name": "Local", "base_url": "http://x.test", "api_key": ""},
|
|
follow_redirects=False,
|
|
)
|
|
client.post("/api/chats/start", data={"content": "what is lembas?"})
|
|
reply = db.scalar(select(Message).where(Message.role == "assistant"))
|
|
|
|
# Following the stream is what runs the generation to completion.
|
|
with client.stream(
|
|
"GET", f"/api/chats/{reply.chat_id}/messages/{reply.id}/stream"
|
|
) as response:
|
|
list(response.iter_lines())
|
|
|
|
db.refresh(reply)
|
|
assert reply.usage_json["prompt_tokens"] == 120
|
|
assert reply.usage_json["completion_tokens"] == 8
|
|
assert reply.usage_json["context_limit"] == 1000
|
|
assert reply.usage_json["estimated"] is False
|
|
assert reply.usage_json["elapsed_ms"] >= 0
|
|
|
|
|
|
def test_the_chips_render_from_a_stored_row(client: TestClient, db, registered, make_chat):
|
|
from lembas.db.models import Message
|
|
|
|
_model(db)
|
|
chat_id = make_chat()
|
|
db.add(
|
|
Message(
|
|
chat_id=chat_id,
|
|
role="assistant",
|
|
content="Waybread.",
|
|
usage_json={
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 20,
|
|
"total_tokens": 120,
|
|
"context_tokens": 120,
|
|
"context_limit": 1000,
|
|
"estimated": False,
|
|
"elapsed_ms": 2000,
|
|
"rounds": 1,
|
|
},
|
|
)
|
|
)
|
|
db.commit()
|
|
|
|
page = client.get(f"/chat/{chat_id}").text
|
|
assert "120 tokens" in page
|
|
assert "12%" in page
|
|
assert "10.0 tok/s" in page
|
|
assert "~" not in page.split("msg__metrics")[1][:400]
|
|
|
|
|
|
def test_an_estimated_row_says_so(client: TestClient, db, registered, make_chat):
|
|
from lembas.db.models import Message
|
|
|
|
_model(db)
|
|
chat_id = make_chat()
|
|
db.add(
|
|
Message(
|
|
chat_id=chat_id,
|
|
role="assistant",
|
|
content="Waybread.",
|
|
usage_json={
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 20,
|
|
"total_tokens": 120,
|
|
"context_tokens": 120,
|
|
"context_limit": 1000,
|
|
"estimated": True,
|
|
"elapsed_ms": 2000,
|
|
"rounds": 1,
|
|
},
|
|
)
|
|
)
|
|
db.commit()
|
|
|
|
page = client.get(f"/chat/{chat_id}").text
|
|
assert "~120 tokens" in page
|
|
assert "reports no token counts" in page
|
|
|
|
|
|
def test_no_context_length_means_no_percentage(client: TestClient, db, registered, make_chat):
|
|
from lembas.db.models import Message
|
|
|
|
_model(db)
|
|
chat_id = make_chat()
|
|
db.add(
|
|
Message(
|
|
chat_id=chat_id,
|
|
role="assistant",
|
|
content="Waybread.",
|
|
usage_json={
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 20,
|
|
"total_tokens": 120,
|
|
"context_tokens": 120,
|
|
"context_limit": 0,
|
|
"estimated": False,
|
|
"elapsed_ms": 2000,
|
|
"rounds": 1,
|
|
},
|
|
)
|
|
)
|
|
db.commit()
|
|
|
|
page = client.get(f"/chat/{chat_id}").text
|
|
assert "120 tokens" in page
|
|
assert "metric--context" not in page
|