From c94995dd8e5054d8adb28560b72d344fe5401eb0 Mon Sep 17 00:00:00 2001 From: Javinator9889 Date: Wed, 16 Sep 2026 23:15:26 +0200 Subject: [PATCH 1/3] fix: report full prompt length in OpenAI usage The prompt cache erases the matching prefix from `tokens` before prefill, so `meta_info.prompt_tokens = tokens.size()` recorded only the newly evaluated suffix. The first turn of a conversation has no cache and looked correct, but every later turn reused the whole history as a cached prefix and reported just the new message. OpenAI defines usage.prompt_tokens as the entire input, with cached tokens counted inside it and broken out separately. Reporting the delta instead made the context appear to reset on every request, so clients that size the context from usage never learn how full it is -- an agent front-end deciding when to compact from that figure never compacts. Add the erased prefix back into prompt_tokens, carry the cached count on chat_meta_info_t, and expose it as usage.prompt_tokens_details.cached_tokens on the OpenAI endpoints. prefill_speed_tps now divides by the tokens actually evaluated, so a cache hit no longer inflates it. The Ollama-compatible prompt_eval_count is left reporting evaluated tokens only, matching upstream Ollama. --- src/common/AutoModel/automodel.cpp | 6 +++++- src/common/AutoModel/modeling_qwen3_5_omni.cpp | 7 ++++++- src/common/AutoModel/modeling_qwen3vl.cpp | 3 +++ src/include/AutoModel/automodel.hpp | 5 +++-- src/server/rest_handler.cpp | 4 +++- src/server/streaming_ostream_openai.hpp | 4 +++- 6 files changed, 23 insertions(+), 6 deletions(-) diff --git a/src/common/AutoModel/automodel.cpp b/src/common/AutoModel/automodel.cpp index 78a4a0f88..4f794522e 100644 --- a/src/common/AutoModel/automodel.cpp +++ b/src/common/AutoModel/automodel.cpp @@ -215,7 +215,11 @@ bool AutoModel::_shared_insert(chat_meta_info_t& meta_info, std::vector& to auto prefill_end_time = this->profiler_list[PREFILL_TIME].stop(tokens.size()); meta_info.prefill_duration = (uint64_t)time_utils::duration_ns(prefill_start_time, prefill_end_time).first; - meta_info.prompt_tokens = tokens.size(); + // `tokens` was trimmed to the uncached suffix above, so add the prefix served + // from the KV cache back on: usage.prompt_tokens is the whole prompt, and the + // cached part is reported separately rather than subtracted. + meta_info.cached_prompt_tokens = static_cast(skip_count); + meta_info.prompt_tokens = static_cast(skip_count + tokens.size()); if (meta_info.stop_reason == CANCEL_DETECTED) { return false; diff --git a/src/common/AutoModel/modeling_qwen3_5_omni.cpp b/src/common/AutoModel/modeling_qwen3_5_omni.cpp index 28695786d..660b6ac66 100644 --- a/src/common/AutoModel/modeling_qwen3_5_omni.cpp +++ b/src/common/AutoModel/modeling_qwen3_5_omni.cpp @@ -551,7 +551,12 @@ bool Qwen3_5_Omni::insert(chat_meta_info_t& meta_info, lm_uniform_input_t& input auto prefill_end = this->profiler_list[PREFILL_TIME].stop(tokens.size()); meta_info.prefill_duration = (uint64_t)time_utils::duration_ns(prefill_start, prefill_end).first; - meta_info.prompt_tokens = tokens.size(); + // As in AutoModel::_shared_insert, `tokens` holds only the uncached suffix, so + // the cached prefix is added back to report the whole prompt. The relevant count + // is `skip_count`, which is what was erased from `tokens` above -- not the earlier + // `prefix_skip_count`, which only trims the multi-modal payload. + meta_info.cached_prompt_tokens = static_cast(skip_count); + meta_info.prompt_tokens = static_cast(skip_count + tokens.size()); this->total_tokens += tokens.size(); diff --git a/src/common/AutoModel/modeling_qwen3vl.cpp b/src/common/AutoModel/modeling_qwen3vl.cpp index bd69d6e67..3f0e7b70e 100644 --- a/src/common/AutoModel/modeling_qwen3vl.cpp +++ b/src/common/AutoModel/modeling_qwen3vl.cpp @@ -583,6 +583,9 @@ bool Qwen3VL_Flash::insert(chat_meta_info_t& meta_info, lm_uniform_input_t& inpu meta_info.prefill_duration = (uint64_t)time_utils::duration_ns(prefill_start_time, prefill_end_time).first; meta_info.prompt_tokens = static_cast(tokens.size()); // report full prompt length to caller + // The pinned system prefix is already in the kv cache, so it is part of the + // prompt but was not prefilled on this turn. + meta_info.cached_prompt_tokens = skip; if (meta_info.stop_reason == CANCEL_DETECTED) { return false; diff --git a/src/include/AutoModel/automodel.hpp b/src/include/AutoModel/automodel.hpp index 45a6a63b7..5b54ec3ba 100644 --- a/src/include/AutoModel/automodel.hpp +++ b/src/include/AutoModel/automodel.hpp @@ -108,7 +108,8 @@ typedef enum { struct chat_meta_info_t { int max_prefill_len; - int prompt_tokens; + int prompt_tokens; // whole prompt, cached prefix included + int cached_prompt_tokens; // subset of prompt_tokens served from the KV cache int generated_tokens; uint64_t total_duration; // in nanoseconds uint64_t load_duration; // in nanoseconds @@ -118,7 +119,7 @@ struct chat_meta_info_t { bool restore_allowed; tool_choice_t tool_choice; - chat_meta_info_t() : max_prefill_len(0), prompt_tokens(0), generated_tokens(0), total_duration(0), load_duration(0), prefill_duration(0), decoding_duration(0), stop_reason(EOT_DETECTED), restore_allowed(false), tool_choice(TOOL_CHOICE_AUTO) {} + chat_meta_info_t() : max_prefill_len(0), prompt_tokens(0), cached_prompt_tokens(0), generated_tokens(0), total_duration(0), load_duration(0), prefill_duration(0), decoding_duration(0), stop_reason(EOT_DETECTED), restore_allowed(false), tool_choice(TOOL_CHOICE_AUTO) {} }; typedef enum { diff --git a/src/server/rest_handler.cpp b/src/server/rest_handler.cpp index 7890a4b5b..b2fe90891 100644 --- a/src/server/rest_handler.cpp +++ b/src/server/rest_handler.cpp @@ -1303,13 +1303,14 @@ void RestHandler::handle_openai_chat_completion(const json& request, {"choices", choices}, {"usage", { {"prompt_tokens", meta_info.prompt_tokens}, + {"prompt_tokens_details", {{"cached_tokens", meta_info.cached_prompt_tokens}}}, {"completion_tokens", meta_info.generated_tokens}, {"total_tokens", meta_info.prompt_tokens + meta_info.generated_tokens}, {"kv_token_occupancy_rate_percentage", (float)this->auto_chat_engine->get_current_context_length() / (float)this->auto_chat_engine->get_max_length() * 100}, {"load_duration", static_cast(meta_info.load_duration) / 1'000'000'000}, {"prefill_duration_ttft", static_cast(meta_info.prefill_duration) / 1'000'000'000}, {"decoding_duration", static_cast(meta_info.decoding_duration) / 1'000'000'000}, - {"prefill_speed_tps", static_cast(meta_info.prompt_tokens) / static_cast(meta_info.prefill_duration) * 1'000'000'000}, + {"prefill_speed_tps", static_cast(meta_info.prompt_tokens - meta_info.cached_prompt_tokens) / static_cast(meta_info.prefill_duration) * 1'000'000'000}, {"decoding_speed_tps", static_cast(meta_info.generated_tokens) / static_cast(meta_info.decoding_duration) * 1'000'000'000}, }}, {"service_tier", "default"} @@ -1511,6 +1512,7 @@ void RestHandler::handle_openai_completion(const json& request, })}, {"usage", { {"prompt_tokens", meta_info.prompt_tokens}, + {"prompt_tokens_details", {{"cached_tokens", meta_info.cached_prompt_tokens}}}, {"completion_tokens", meta_info.generated_tokens}, {"total_tokens", meta_info.prompt_tokens + meta_info.generated_tokens} }} diff --git a/src/server/streaming_ostream_openai.hpp b/src/server/streaming_ostream_openai.hpp index 15cd1f9cf..ca93e4849 100644 --- a/src/server/streaming_ostream_openai.hpp +++ b/src/server/streaming_ostream_openai.hpp @@ -207,6 +207,7 @@ class streaming_buf_openai : public std::streambuf { })}, {"usage", { {"prompt_tokens", meta_info.prompt_tokens}, + {"prompt_tokens_details", {{"cached_tokens", meta_info.cached_prompt_tokens}}}, {"completion_tokens", meta_info.generated_tokens}, {"total_tokens", meta_info.prompt_tokens + meta_info.generated_tokens} }} @@ -497,6 +498,7 @@ class streaming_buf_openai_chat : public std::streambuf { })}, {"usage", { {"prompt_tokens", meta_info.prompt_tokens}, + {"prompt_tokens_details", {{"cached_tokens", meta_info.cached_prompt_tokens}}}, {"completion_tokens", meta_info.generated_tokens}, {"total_tokens", meta_info.prompt_tokens + meta_info.generated_tokens}, {"active_kv_tokens", this->auto_chat_engine->get_current_context_length()}, @@ -505,7 +507,7 @@ class streaming_buf_openai_chat : public std::streambuf { {"load_duration", static_cast(meta_info.load_duration) / 1'000'000'000}, {"prefill_duration_ttft", static_cast(meta_info.prefill_duration) / 1'000'000'000}, {"decoding_duration", static_cast(meta_info.decoding_duration) / 1'000'000'000}, - {"prefill_speed_tps", static_cast(meta_info.prompt_tokens) / static_cast(meta_info.prefill_duration) * 1'000'000'000}, + {"prefill_speed_tps", static_cast(meta_info.prompt_tokens - meta_info.cached_prompt_tokens) / static_cast(meta_info.prefill_duration) * 1'000'000'000}, {"decoding_speed_tps", static_cast(meta_info.generated_tokens) / static_cast(meta_info.decoding_duration) * 1'000'000'000}, }} }; From 2f8cb4f77f15ccfc260a684bb40565494583288f Mon Sep 17 00:00:00 2001 From: Javinator9889 Date: Thu, 17 Sep 2026 10:16:12 +0200 Subject: [PATCH 2/3] test: cover the OpenAI usage/cache contract Two-turn check, streaming and non-streaming, that prompt_tokens counts the cached prefix and cached_tokens reports it. Needs a live server. --- src/test/openai_usage/test_usage_contract.py | 413 +++++++++++++++++++ 1 file changed, 413 insertions(+) create mode 100755 src/test/openai_usage/test_usage_contract.py diff --git a/src/test/openai_usage/test_usage_contract.py b/src/test/openai_usage/test_usage_contract.py new file mode 100755 index 000000000..28719e99b --- /dev/null +++ b/src/test/openai_usage/test_usage_contract.py @@ -0,0 +1,413 @@ +#!/usr/bin/env python3 +"""Integration test for the OpenAI-compatible usage/cache contract. + +FastFlowLM's engine keeps a KV cache across turns of the same +conversation and only prefills the tokens that were not already in +that cache. A regression made ``usage.prompt_tokens`` report only the +newly prefilled tail instead of the whole prompt, so a client that +tracks context growth from ``usage`` saw it collapse to near-zero on +every turn after the first. + +This script guards the fix by driving a live, already-running +FastFlowLM server through two sequential chat turns -- once with a +plain JSON response and once with ``stream=True`` -- and checking +that, on both endpoints: + +* ``usage.prompt_tokens`` is the *whole* prompt (cached prefix + new + tail), so it strictly grows from turn 1 to turn 2. +* ``usage.prompt_tokens_details.cached_tokens`` is present on turn 2, + is a positive integer, and never exceeds ``prompt_tokens`` (it is a + subset of it, not an alternate count). +* ``usage.total_tokens == usage.prompt_tokens + + usage.completion_tokens`` on every turn. + +It deliberately does NOT assert anything about +``stream_options.include_usage`` or about ``usage`` being present (as +``null`` or otherwise) on intermediate streaming chunks -- that +behaviour belongs to a different, unmerged branch. + +Usage +----- +Start the FastFlowLM server yourself (e.g. ``flm serve ``), +then run:: + + python3 test_usage_contract.py + python3 test_usage_contract.py --endpoint URL --timeout T + +The script requires only the Python 3 standard library. It exits 0 +and prints a pass summary (with the observed token counts for every +turn) on success, or exits non-zero with a diagnosable +``AssertionError`` message on failure. +""" + +# Every check here reports on the server under test, not on a caller +# passing bad arguments, so a wrong type in the response is a test +# failure rather than a TypeError. main() catches AssertionError alone +# to turn any violation into the pass/fail summary and exit code. +# ruff: noqa: TRY004 + +import argparse +import copy +import json +import sys +import urllib.error +import urllib.request +from typing import Any + +DEFAULT_ENDPOINT = "http://127.0.0.1:52625" +DEFAULT_TIMEOUT = 300.0 +MAX_TOKENS = 64 +TEMPERATURE = 0.1 +PROMPT_1 = "In one short sentence, what is the capital of France?" +PROMPT_2 = "In one short sentence, name its most famous landmark." + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments for this integration test.""" + parser = argparse.ArgumentParser( + description=( + "Integration test for the OpenAI usage/cache contract: " + "usage.prompt_tokens must report the whole prompt " + "(cached prefix + new tail), and " + "usage.prompt_tokens_details.cached_tokens must report " + "the cached subset of it." + ), + ) + parser.add_argument( + "model", + help="model tag currently loaded by the FastFlowLM server", + ) + parser.add_argument( + "--endpoint", + default=DEFAULT_ENDPOINT, + help="server base URL (default: %(default)s)", + ) + parser.add_argument( + "--timeout", + type=float, + default=DEFAULT_TIMEOUT, + help="per-request timeout in seconds (default: %(default)s)", + ) + return parser.parse_args() + + +def chat_completions_url(endpoint: str) -> str: + """Build the chat completions URL from a server base ``endpoint``.""" + return endpoint.rstrip("/") + "/v1/chat/completions" + + +def build_turn1_messages() -> list[dict[str, str]]: + """Build the message list for the first turn of the conversation.""" + return [{"role": "user", "content": PROMPT_1}] + + +def build_turn2_messages( + turn1_messages: list[dict[str, str]], + assistant_reply: str, +) -> list[dict[str, str]]: + """Append the turn-1 assistant reply and a new user message. + + ``turn1_messages`` is not mutated; a new list is returned that + contains the original turn, the assistant's reply, and a follow + up user message, so the server sees a growing conversation. + """ + messages = copy.deepcopy(turn1_messages) + messages.append({"role": "assistant", "content": assistant_reply}) + messages.append({"role": "user", "content": PROMPT_2}) + return messages + + +def post_json( + url: str, + payload: dict[str, Any], + timeout: float, +) -> bytes: + """POST ``payload`` as JSON to ``url`` and return the raw body. + + Raises ``AssertionError`` (with the response body, if any) when + the request fails, so a failure is diagnosable without a + debugger. + """ + data = json.dumps(payload).encode("utf-8") + accept = ( + "text/event-stream" if payload.get("stream") else "application/json" + ) + request = urllib.request.Request( + url, + data=data, + method="POST", + headers={ + "Content-Type": "application/json", + "Accept": accept, + }, + ) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + return response.read() + except urllib.error.HTTPError as error: + body = error.read().decode("utf-8", errors="replace") + raise AssertionError( + f"POST {url} failed with HTTP {error.code}: {body}" + ) from error + except urllib.error.URLError as error: + raise AssertionError(f"POST {url} failed: {error.reason}") from error + + +def extract_message_and_usage( + response: dict[str, Any], +) -> tuple[str, dict[str, Any]]: + """Extract the assistant reply and usage from a JSON response. + + ``response`` is the decoded body of a non-streaming chat + completion response. + """ + choices = response.get("choices") + if not choices: + raise AssertionError(f"response has no 'choices': {response!r}") + message = choices[0].get("message") or {} + content = message.get("content") + if not isinstance(content, str): + raise AssertionError( + f"assistant message has no string 'content': {message!r}" + ) + usage = response.get("usage") + if not isinstance(usage, dict): + raise AssertionError(f"response has no 'usage' object: {response!r}") + return content, usage + + +def parse_sse_chunks(raw_body: bytes) -> list[dict[str, Any]]: + """Parse an SSE response body into a list of decoded JSON chunks. + + Lines that are not ``data: ...`` are ignored, and the terminal + literal ``data: [DONE]`` line is dropped rather than parsed as + JSON. + """ + chunks = [] + text = raw_body.decode("utf-8", errors="replace") + for raw_line in text.splitlines(): + line = raw_line.strip() + if not line.startswith("data:"): + continue + payload = line[len("data:") :].strip() + if payload == "[DONE]": + continue + if not payload: + continue + chunks.append(json.loads(payload)) + return chunks + + +def extract_streamed_reply_and_usage( + chunks: list[dict[str, Any]], +) -> tuple[str, dict[str, Any]]: + """Reassemble the assistant reply and find the final usage chunk. + + The assistant reply is reassembled by concatenating every + ``choices[0].delta.content`` piece in order. The final usage is + taken from the last chunk that carries a non-null ``usage`` + object, per the OpenAI streaming contract. + """ + content_parts = [] + final_usage = None + for chunk in chunks: + choices = chunk.get("choices") or [] + if choices: + delta = choices[0].get("delta") or {} + piece = delta.get("content") + if isinstance(piece, str): + content_parts.append(piece) + usage = chunk.get("usage") + if isinstance(usage, dict): + final_usage = usage + if final_usage is None: + raise AssertionError( + "no streamed chunk carried a non-null 'usage' object; " + f"chunks were: {chunks!r}" + ) + return "".join(content_parts), final_usage + + +def print_usage(turn_label: str, usage: dict[str, Any]) -> None: + """Print the observed token counts for one turn, for PR evidence.""" + details = usage.get("prompt_tokens_details") + cached_tokens = None + if isinstance(details, dict): + cached_tokens = details.get("cached_tokens") + prompt_tokens = usage.get("prompt_tokens") + completion_tokens = usage.get("completion_tokens") + total_tokens = usage.get("total_tokens") + print( + f" {turn_label}: prompt_tokens={prompt_tokens} " + f"cached_tokens={cached_tokens} " + f"completion_tokens={completion_tokens} " + f"total_tokens={total_tokens}" + ) + + +def assert_usage_consistency(usage: dict[str, Any], turn_label: str) -> None: + """Assert ``total_tokens == prompt_tokens + completion_tokens``. + + Raises ``AssertionError`` with the offending usage object when + fields are missing, not integers, or inconsistent. + """ + prompt_tokens = usage.get("prompt_tokens") + completion_tokens = usage.get("completion_tokens") + total_tokens = usage.get("total_tokens") + for name, value in ( + ("prompt_tokens", prompt_tokens), + ("completion_tokens", completion_tokens), + ("total_tokens", total_tokens), + ): + if not isinstance(value, int) or isinstance(value, bool): + raise AssertionError( + f"{turn_label}: usage.{name} must be an int, got {value!r} in " + f"{usage!r}" + ) + if total_tokens != prompt_tokens + completion_tokens: + raise AssertionError( + f"{turn_label}: total_tokens ({total_tokens}) != " + f"prompt_tokens ({prompt_tokens}) + " + f"completion_tokens ({completion_tokens})" + ) + + +def assert_cache_growth( + turn1_prompt_tokens: int, + turn2_usage: dict[str, Any], + turn_label: str, +) -> None: + """Assert turn 2 grew the prompt and reused a cached prefix. + + Checks that turn 2's ``prompt_tokens`` is strictly greater than + turn 1's, and that ``prompt_tokens_details.cached_tokens`` is a + positive integer no larger than turn 2's ``prompt_tokens``. + Nothing is asserted about turn 1's ``cached_tokens``: the server + is long-lived and may already hold a cached prefix from earlier + activity. + """ + turn2_prompt_tokens = turn2_usage.get("prompt_tokens") + if not isinstance(turn2_prompt_tokens, int): + raise AssertionError( + f"{turn_label}: turn 2 usage.prompt_tokens must be an int, got " + f"{turn2_prompt_tokens!r} in {turn2_usage!r}" + ) + if turn2_prompt_tokens <= turn1_prompt_tokens: + raise AssertionError( + f"{turn_label}: expected turn 2 prompt_tokens " + f"({turn2_prompt_tokens}) > turn 1 prompt_tokens " + f"({turn1_prompt_tokens}); the KV cache prefix reuse bug " + "would make usage collapse instead of grow" + ) + details = turn2_usage.get("prompt_tokens_details") + if not isinstance(details, dict) or "cached_tokens" not in details: + raise AssertionError( + f"{turn_label}: turn 2 usage is missing " + f"prompt_tokens_details.cached_tokens: {turn2_usage!r}" + ) + cached_tokens = details["cached_tokens"] + if not isinstance(cached_tokens, int) or isinstance(cached_tokens, bool): + raise AssertionError( + f"{turn_label}: cached_tokens must be an int, got " + f"{cached_tokens!r} ({type(cached_tokens).__name__})" + ) + if cached_tokens <= 0: + raise AssertionError( + f"{turn_label}: expected turn 2 cached_tokens > 0, " + f"got {cached_tokens}" + ) + if cached_tokens > turn2_prompt_tokens: + raise AssertionError( + f"{turn_label}: cached_tokens ({cached_tokens}) exceeds " + f"prompt_tokens ({turn2_prompt_tokens}); cached_tokens " + "must be a subset of prompt_tokens" + ) + + +def run_non_streaming_test(endpoint: str, model: str, timeout: float) -> None: + """Run Test 1: the two-turn, non-streaming usage contract check.""" + print("Test 1: non-streaming usage contract") + url = chat_completions_url(endpoint) + + turn1_messages = build_turn1_messages() + turn1_payload = { + "model": model, + "messages": turn1_messages, + "temperature": TEMPERATURE, + "max_tokens": MAX_TOKENS, + "stream": False, + } + turn1_body = post_json(url, turn1_payload, timeout) + turn1_response = json.loads(turn1_body) + turn1_reply, turn1_usage = extract_message_and_usage(turn1_response) + print_usage("turn 1", turn1_usage) + assert_usage_consistency(turn1_usage, "test 1 turn 1") + + turn2_messages = build_turn2_messages(turn1_messages, turn1_reply) + turn2_payload = copy.deepcopy(turn1_payload) + turn2_payload["messages"] = turn2_messages + turn2_body = post_json(url, turn2_payload, timeout) + turn2_response = json.loads(turn2_body) + _, turn2_usage = extract_message_and_usage(turn2_response) + print_usage("turn 2", turn2_usage) + assert_usage_consistency(turn2_usage, "test 1 turn 2") + assert_cache_growth( + turn1_usage["prompt_tokens"], + turn2_usage, + "test 1", + ) + print("Test 1 PASSED") + + +def run_streaming_test(endpoint: str, model: str, timeout: float) -> None: + """Run Test 2: the two-turn, streaming usage contract check.""" + print("\nTest 2: streaming usage contract") + url = chat_completions_url(endpoint) + + turn1_messages = build_turn1_messages() + turn1_payload = { + "model": model, + "messages": turn1_messages, + "temperature": TEMPERATURE, + "max_tokens": MAX_TOKENS, + "stream": True, + } + turn1_body = post_json(url, turn1_payload, timeout) + turn1_chunks = parse_sse_chunks(turn1_body) + turn1_reply, turn1_usage = extract_streamed_reply_and_usage( + turn1_chunks, + ) + print_usage("turn 1 (stream)", turn1_usage) + assert_usage_consistency(turn1_usage, "test 2 turn 1") + + turn2_messages = build_turn2_messages(turn1_messages, turn1_reply) + turn2_payload = copy.deepcopy(turn1_payload) + turn2_payload["messages"] = turn2_messages + turn2_body = post_json(url, turn2_payload, timeout) + turn2_chunks = parse_sse_chunks(turn2_body) + _, turn2_usage = extract_streamed_reply_and_usage(turn2_chunks) + print_usage("turn 2 (stream)", turn2_usage) + assert_usage_consistency(turn2_usage, "test 2 turn 2") + assert_cache_growth( + turn1_usage["prompt_tokens"], + turn2_usage, + "test 2", + ) + print("Test 2 PASSED") + + +def main() -> int: + """Run both usage-contract tests and print a pass/fail summary.""" + args = parse_args() + try: + run_non_streaming_test(args.endpoint, args.model, args.timeout) + run_streaming_test(args.endpoint, args.model, args.timeout) + except AssertionError as error: + print(f"\nFAIL: {error}", file=sys.stderr) + return 1 + print("\nAll usage contract checks PASSED.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From ce7eadbc12c7dd6694bb8f8da56eb46dd8641507 Mon Sep 17 00:00:00 2001 From: Javinator9889 Date: Thu, 17 Sep 2026 11:17:34 +0200 Subject: [PATCH 3/3] test: adapt usage contract test to single-turn models Single-turn models reject conversation history, so the two-turn shape does not apply. Detect the server's rejection and check the pinned system prefix across two one-shot requests instead. --- src/test/openai_usage/test_usage_contract.py | 139 +++++++++++++++++-- 1 file changed, 126 insertions(+), 13 deletions(-) diff --git a/src/test/openai_usage/test_usage_contract.py b/src/test/openai_usage/test_usage_contract.py index 28719e99b..f13e3fe60 100755 --- a/src/test/openai_usage/test_usage_contract.py +++ b/src/test/openai_usage/test_usage_contract.py @@ -61,6 +61,22 @@ PROMPT_1 = "In one short sentence, what is the capital of France?" PROMPT_2 = "In one short sentence, name its most famous landmark." +# Single-turn models reject conversation history, so their cache reuse +# is a pinned system prefix shared by successive one-shot requests. +SINGLE_TURN_SYSTEM = ( + "You are a terse assistant. Answer in one short sentence. " + "Never elaborate. Stay factual." +) +SINGLE_TURN_PROMPT_1 = "What is the capital of France?" +SINGLE_TURN_PROMPT_2 = "What is the capital of Japan?" + +# Substring of the server's rejection when history is not allowed. +SINGLE_TURN_MARKER = b"only supports single-turn requests" + + +class SingleTurnModel(Exception): + """The model under test rejects multi-turn conversation history.""" + def parse_args() -> argparse.Namespace: """Parse command-line arguments for this integration test.""" @@ -117,6 +133,14 @@ def build_turn2_messages( return messages +def build_single_turn_messages(prompt: str) -> list[dict[str, str]]: + """Build a one-shot exchange sharing the pinned system prefix.""" + return [ + {"role": "system", "content": SINGLE_TURN_SYSTEM}, + {"role": "user", "content": prompt}, + ] + + def post_json( url: str, payload: dict[str, Any], @@ -143,14 +167,22 @@ def post_json( ) try: with urllib.request.urlopen(request, timeout=timeout) as response: - return response.read() + body = response.read() except urllib.error.HTTPError as error: - body = error.read().decode("utf-8", errors="replace") + raw = error.read() + if SINGLE_TURN_MARKER in raw: + raise SingleTurnModel(raw.decode("utf-8", errors="replace")) + text = raw.decode("utf-8", errors="replace") raise AssertionError( - f"POST {url} failed with HTTP {error.code}: {body}" + f"POST {url} failed with HTTP {error.code}: {text}" ) from error except urllib.error.URLError as error: raise AssertionError(f"POST {url} failed: {error.reason}") from error + # The server answers 200 with an error body in some builds, so the + # marker has to be checked on the success path too. + if SINGLE_TURN_MARKER in body: + raise SingleTurnModel(body.decode("utf-8", errors="replace")) + return body def extract_message_and_usage( @@ -299,11 +331,22 @@ def assert_cache_growth( f"({turn1_prompt_tokens}); the KV cache prefix reuse bug " "would make usage collapse instead of grow" ) - details = turn2_usage.get("prompt_tokens_details") + assert_prefix_reused(turn2_usage, turn_label) + + +def assert_prefix_reused(usage: dict[str, Any], turn_label: str) -> None: + """Assert a cached prefix was reused and counted in the prompt. + + ``prompt_tokens_details.cached_tokens`` must be a positive integer + no larger than ``prompt_tokens``: the cached prefix is reported as + a subset of the whole prompt, never instead of it. + """ + prompt_tokens = usage.get("prompt_tokens") + details = usage.get("prompt_tokens_details") if not isinstance(details, dict) or "cached_tokens" not in details: raise AssertionError( - f"{turn_label}: turn 2 usage is missing " - f"prompt_tokens_details.cached_tokens: {turn2_usage!r}" + f"{turn_label}: usage is missing " + f"prompt_tokens_details.cached_tokens: {usage!r}" ) cached_tokens = details["cached_tokens"] if not isinstance(cached_tokens, int) or isinstance(cached_tokens, bool): @@ -313,17 +356,57 @@ def assert_cache_growth( ) if cached_tokens <= 0: raise AssertionError( - f"{turn_label}: expected turn 2 cached_tokens > 0, " - f"got {cached_tokens}" + f"{turn_label}: expected cached_tokens > 0, got {cached_tokens}" ) - if cached_tokens > turn2_prompt_tokens: + if cached_tokens > prompt_tokens: raise AssertionError( f"{turn_label}: cached_tokens ({cached_tokens}) exceeds " - f"prompt_tokens ({turn2_prompt_tokens}); cached_tokens " - "must be a subset of prompt_tokens" + f"prompt_tokens ({prompt_tokens}); cached_tokens must be a " + "subset of prompt_tokens" ) +def run_pinned_prefix_checks( + url: str, + model: str, + timeout: float, + stream: bool, + label: str, +) -> None: + """Check cache accounting on a single-turn model. + + Single-turn models reject conversation history, so the growing + transcript used elsewhere does not apply. Their cache reuse is the + pinned system prefix shared by successive one-shot requests: two + requests share one system message, and the second must report that + prefix as reused and counted inside ``prompt_tokens``. + """ + prompts = (SINGLE_TURN_PROMPT_1, SINGLE_TURN_PROMPT_2) + suffix = " (stream)" if stream else "" + usages = [] + for index, prompt in enumerate(prompts, start=1): + payload: dict[str, Any] = { + "model": model, + "messages": build_single_turn_messages(prompt), + "temperature": TEMPERATURE, + "max_tokens": MAX_TOKENS, + } + if stream: + payload["stream"] = True + body = post_json(url, payload, timeout) + if stream: + _, usage = extract_streamed_reply_and_usage(parse_sse_chunks(body)) + else: + _, usage = extract_message_and_usage(json.loads(body)) + print_usage(f"request {index}{suffix}", usage) + assert_usage_consistency(usage, f"{label} request {index}") + usages.append(usage) + # Only the second request is required to show a reused prefix: the + # first may or may not already sit on the pin, depending on whether + # the model was just loaded. + assert_prefix_reused(usages[1], f"{label} request 2") + + def run_non_streaming_test(endpoint: str, model: str, timeout: float) -> None: """Run Test 1: the two-turn, non-streaming usage contract check.""" print("Test 1: non-streaming usage contract") @@ -346,7 +429,22 @@ def run_non_streaming_test(endpoint: str, model: str, timeout: float) -> None: turn2_messages = build_turn2_messages(turn1_messages, turn1_reply) turn2_payload = copy.deepcopy(turn1_payload) turn2_payload["messages"] = turn2_messages - turn2_body = post_json(url, turn2_payload, timeout) + try: + turn2_body = post_json(url, turn2_payload, timeout) + except SingleTurnModel: + print( + " model rejects multi-turn history; checking the pinned " + "system prefix instead" + ) + run_pinned_prefix_checks( + url, + model, + timeout, + stream=False, + label="test 1", + ) + print("Test 1 PASSED") + return turn2_response = json.loads(turn2_body) _, turn2_usage = extract_message_and_usage(turn2_response) print_usage("turn 2", turn2_usage) @@ -383,7 +481,22 @@ def run_streaming_test(endpoint: str, model: str, timeout: float) -> None: turn2_messages = build_turn2_messages(turn1_messages, turn1_reply) turn2_payload = copy.deepcopy(turn1_payload) turn2_payload["messages"] = turn2_messages - turn2_body = post_json(url, turn2_payload, timeout) + try: + turn2_body = post_json(url, turn2_payload, timeout) + except SingleTurnModel: + print( + " model rejects multi-turn history; checking the pinned " + "system prefix instead" + ) + run_pinned_prefix_checks( + url, + model, + timeout, + stream=True, + label="test 2", + ) + print("Test 2 PASSED") + return turn2_chunks = parse_sse_chunks(turn2_body) _, turn2_usage = extract_streamed_reply_and_usage(turn2_chunks) print_usage("turn 2 (stream)", turn2_usage)