Implement WebAssembly runtime and kinematics coverage
This commit is contained in:
501
scripts/validate_model_budget.py
Normal file
501
scripts/validate_model_budget.py
Normal file
@@ -0,0 +1,501 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Check that a bounded context packet leaves room for model output.
|
||||
|
||||
This is deliberately provider-neutral. The caller supplies the model's
|
||||
context window; the packet's input budget and a fixed output reserve are
|
||||
checked before an API request is made.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
OUTPUT_TOKEN_RESERVE = {
|
||||
"doc": 768,
|
||||
"feature": 1024,
|
||||
"contract": 1536,
|
||||
"milestone": 2048,
|
||||
}
|
||||
PROFILE_INPUT_TOKEN_LIMIT = {
|
||||
"doc": 2048,
|
||||
"feature": 4096,
|
||||
"contract": 6144,
|
||||
"milestone": 8192,
|
||||
}
|
||||
PROFILE_BYTE_LIMIT = {
|
||||
"doc": 8192,
|
||||
"feature": 16384,
|
||||
"contract": 24576,
|
||||
"milestone": 32768,
|
||||
}
|
||||
PROFILE_EVIDENCE_BYTE_LIMIT = {
|
||||
"doc": 2048,
|
||||
"feature": 6144,
|
||||
"contract": 10240,
|
||||
"milestone": 14336,
|
||||
}
|
||||
PROFILE_EVIDENCE_FILE_LIMIT = {
|
||||
"doc": 8,
|
||||
"feature": 16,
|
||||
"contract": 24,
|
||||
"milestone": 32,
|
||||
}
|
||||
DEFAULT_CONTEXT_WINDOW = 16384
|
||||
DEFAULT_SAFETY_MARGIN = 256
|
||||
# Without a provider tokenizer, UTF-8 bytes are the only portable hard upper
|
||||
# bound. Any multi-byte-per-token heuristic can under-count provider tokens
|
||||
# for punctuation-heavy or non-ASCII input and is unsafe for a request gate.
|
||||
TOKEN_BYTES_PER_TOKEN = 1
|
||||
PACKET_FIELDS = {
|
||||
"context_id",
|
||||
"packet_id",
|
||||
"task",
|
||||
"header",
|
||||
"baseline",
|
||||
"scope",
|
||||
"contracts",
|
||||
"evidence",
|
||||
"risks",
|
||||
"next",
|
||||
"forbidden",
|
||||
"budget",
|
||||
}
|
||||
FORBIDDEN_CONTEXT_FIELDS = {
|
||||
"history",
|
||||
"chat",
|
||||
"conversation",
|
||||
"messages",
|
||||
"transcript",
|
||||
"previous_packet",
|
||||
"prior_context",
|
||||
}
|
||||
EXCLUSION_REASON_CODES = {"E1", "E2", "E3", "E4", "E5", "E6"}
|
||||
|
||||
|
||||
class ModelBudgetError(ValueError):
|
||||
"""Raised when a request cannot fit input, output, and safety budgets."""
|
||||
|
||||
|
||||
def _response_object(value: Any) -> dict[str, Any]:
|
||||
"""Return a provider response object from a response or SSE envelope.
|
||||
|
||||
The Responses API exposes usage on the response itself, while streaming
|
||||
``response.completed`` events wrap it below a ``response`` key. Keeping
|
||||
this normalization here prevents callers from accidentally validating the
|
||||
event envelope instead of the actual provider usage.
|
||||
"""
|
||||
|
||||
if not isinstance(value, dict):
|
||||
raise ModelBudgetError("provider response must be an object")
|
||||
response = value.get("response")
|
||||
if isinstance(response, dict):
|
||||
return response
|
||||
return value
|
||||
|
||||
|
||||
def extract_response_usage(value: Any) -> dict[str, int]:
|
||||
"""Extract authoritative token usage from a provider response envelope."""
|
||||
|
||||
response = _response_object(value)
|
||||
usage = response.get("usage")
|
||||
if not isinstance(usage, dict):
|
||||
raise ModelBudgetError("provider response is missing usage")
|
||||
try:
|
||||
input_tokens = int(usage["input_tokens"])
|
||||
output_tokens = int(usage["output_tokens"])
|
||||
total_tokens = int(usage.get("total_tokens", input_tokens + output_tokens))
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise ModelBudgetError("provider usage has invalid token counts") from exc
|
||||
if min(input_tokens, output_tokens, total_tokens) < 0:
|
||||
raise ModelBudgetError("provider usage token counts must be non-negative")
|
||||
if total_tokens < input_tokens + output_tokens:
|
||||
raise ModelBudgetError("provider usage total is smaller than input plus output")
|
||||
return {
|
||||
"input_tokens": input_tokens,
|
||||
"output_tokens": output_tokens,
|
||||
"total_tokens": total_tokens,
|
||||
}
|
||||
|
||||
|
||||
def _find_forbidden(value: Any, path: str = "") -> str | None:
|
||||
"""Find accumulated conversation state, including when nested."""
|
||||
|
||||
if isinstance(value, dict):
|
||||
for key, child in value.items():
|
||||
key_text = str(key)
|
||||
child_path = f"{path}.{key_text}" if path else key_text
|
||||
if key_text in FORBIDDEN_CONTEXT_FIELDS:
|
||||
return child_path
|
||||
found = _find_forbidden(child, child_path)
|
||||
if found:
|
||||
return found
|
||||
elif isinstance(value, list):
|
||||
for index, child in enumerate(value):
|
||||
found = _find_forbidden(child, f"{path}[{index}]")
|
||||
if found:
|
||||
return found
|
||||
return None
|
||||
|
||||
|
||||
def _packet_input_bytes(packet: dict[str, Any]) -> int:
|
||||
budget = packet.get("budget")
|
||||
if not isinstance(budget, dict):
|
||||
raise ModelBudgetError("packet budget is missing")
|
||||
try:
|
||||
packet_bytes = int(budget["packet_bytes"])
|
||||
evidence_bytes = int(budget["selected_evidence_bytes"])
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise ModelBudgetError("packet budget lacks input byte totals") from exc
|
||||
if packet_bytes < 0 or evidence_bytes < 0:
|
||||
raise ModelBudgetError("packet input byte totals must be non-negative")
|
||||
return packet_bytes + evidence_bytes
|
||||
|
||||
|
||||
def _canonical_packet_bytes(packet: dict[str, Any]) -> int:
|
||||
"""Return the exact UTF-8 size used by the packet generator."""
|
||||
|
||||
return len(
|
||||
json.dumps(packet, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
|
||||
)
|
||||
|
||||
|
||||
def _validate_packet_metadata(packet: dict[str, Any], budget: dict[str, Any]) -> None:
|
||||
"""Reject stale or inconsistent budget metadata before a model request.
|
||||
|
||||
Small unit-test fixtures may contain only a budget object. Full packets
|
||||
generated for a request always carry ``context_id``/``packet_id`` and an
|
||||
evidence envelope, so strict consistency checks are applied to those.
|
||||
"""
|
||||
|
||||
forbidden = _find_forbidden(packet)
|
||||
if forbidden:
|
||||
raise ModelBudgetError(
|
||||
"packet contains forbidden accumulated context "
|
||||
f"field: {forbidden}"
|
||||
)
|
||||
|
||||
# Generated packets have a closed envelope. Reject ad-hoc top-level
|
||||
# fields that could smuggle full logs, tool output or a previous request
|
||||
# beside the audited eight segments.
|
||||
if "budget" in packet:
|
||||
unknown = sorted(set(packet) - PACKET_FIELDS)
|
||||
if unknown:
|
||||
raise ModelBudgetError(
|
||||
"packet contains unknown top-level context field: "
|
||||
f"{unknown[0]}"
|
||||
)
|
||||
|
||||
if set(packet) != {"budget"}:
|
||||
actual = _canonical_packet_bytes(packet)
|
||||
declared = int(budget["packet_bytes"])
|
||||
if declared != actual:
|
||||
raise ModelBudgetError(
|
||||
"packet byte metadata does not match serialized packet "
|
||||
f"(declared={declared}, actual={actual})"
|
||||
)
|
||||
|
||||
evidence = packet.get("evidence")
|
||||
if not isinstance(evidence, dict):
|
||||
return
|
||||
selected = evidence.get("selected", [])
|
||||
excluded = evidence.get("excluded", [])
|
||||
if not isinstance(selected, list) or not isinstance(excluded, list):
|
||||
raise ModelBudgetError("packet evidence must contain selected/excluded lists")
|
||||
for item in excluded:
|
||||
if not isinstance(item, dict) or item.get("reason_code") not in EXCLUSION_REASON_CODES:
|
||||
raise ModelBudgetError("excluded evidence has an unknown reason code")
|
||||
selected_bytes = 0
|
||||
for item in selected:
|
||||
if not isinstance(item, dict):
|
||||
raise ModelBudgetError("selected evidence entries must be objects")
|
||||
try:
|
||||
item_bytes = int(item["bytes"])
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise ModelBudgetError("selected evidence entry lacks bytes") from exc
|
||||
if item_bytes < 0:
|
||||
raise ModelBudgetError("selected evidence bytes must be non-negative")
|
||||
selected_bytes += item_bytes
|
||||
if selected_bytes != int(budget["selected_evidence_bytes"]):
|
||||
raise ModelBudgetError(
|
||||
"selected evidence byte metadata does not match entries "
|
||||
f"(declared={budget['selected_evidence_bytes']}, actual={selected_bytes})"
|
||||
)
|
||||
if "selected_evidence_files" in budget and int(budget["selected_evidence_files"]) != len(selected):
|
||||
raise ModelBudgetError("selected evidence file count does not match entries")
|
||||
if "excluded_evidence_files" in budget and int(budget["excluded_evidence_files"]) != len(excluded):
|
||||
raise ModelBudgetError("excluded evidence file count does not match entries")
|
||||
if "evidence_byte_limit" in budget and selected_bytes > int(budget["evidence_byte_limit"]):
|
||||
raise ModelBudgetError("selected evidence exceeds its evidence byte budget")
|
||||
if "evidence_file_limit" in budget and len(selected) > int(budget["evidence_file_limit"]):
|
||||
raise ModelBudgetError("selected evidence exceeds its file budget")
|
||||
|
||||
|
||||
def estimate_input_tokens(packet: dict[str, Any]) -> int:
|
||||
"""Return a safe provider-neutral upper bound from UTF-8 bytes.
|
||||
|
||||
A tokenizer supplied by the API client may replace this estimate. The
|
||||
byte-for-token bound is deliberately strict: it may reject a packet that
|
||||
would fit a particular provider, but it cannot silently approve a packet
|
||||
that that provider tokenizes larger than expected.
|
||||
"""
|
||||
|
||||
return math.ceil(_packet_input_bytes(packet) / TOKEN_BYTES_PER_TOKEN)
|
||||
|
||||
|
||||
def validate_request(
|
||||
packet: dict[str, Any],
|
||||
*,
|
||||
context_window: int = DEFAULT_CONTEXT_WINDOW,
|
||||
max_output_tokens: int | None = None,
|
||||
safety_margin: int = DEFAULT_SAFETY_MARGIN,
|
||||
input_tokens: int | None = None,
|
||||
overhead_tokens: int = 0,
|
||||
) -> dict[str, int]:
|
||||
budget = packet.get("budget")
|
||||
if not isinstance(budget, dict):
|
||||
raise ModelBudgetError("packet budget is missing")
|
||||
try:
|
||||
input_limit = int(budget["token_limit"])
|
||||
byte_limit = int(budget["byte_limit"])
|
||||
packet_bytes = int(budget["packet_bytes"])
|
||||
selected_evidence_bytes = int(budget["selected_evidence_bytes"])
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise ModelBudgetError("packet budget lacks required fields") from exc
|
||||
profile = str(budget.get("profile", ""))
|
||||
try:
|
||||
profile_output_limit = OUTPUT_TOKEN_RESERVE[profile]
|
||||
profile_input_limit = PROFILE_INPUT_TOKEN_LIMIT[profile]
|
||||
profile_byte_limit = PROFILE_BYTE_LIMIT[profile]
|
||||
profile_evidence_limit = PROFILE_EVIDENCE_BYTE_LIMIT[profile]
|
||||
profile_file_limit = PROFILE_EVIDENCE_FILE_LIMIT[profile]
|
||||
except KeyError as exc:
|
||||
raise ModelBudgetError(f"unknown packet profile: {profile}") from exc
|
||||
if max_output_tokens is None:
|
||||
max_output_tokens = profile_output_limit
|
||||
try:
|
||||
context_window = int(context_window)
|
||||
max_output_tokens = int(max_output_tokens)
|
||||
safety_margin = int(safety_margin)
|
||||
overhead_tokens = int(overhead_tokens)
|
||||
declared_evidence_limit = int(budget.get("evidence_byte_limit", profile_evidence_limit))
|
||||
declared_file_limit = int(budget.get("evidence_file_limit", profile_file_limit))
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ModelBudgetError("model budget fields must be integers") from exc
|
||||
values = (
|
||||
context_window,
|
||||
max_output_tokens,
|
||||
safety_margin,
|
||||
input_limit,
|
||||
byte_limit,
|
||||
packet_bytes,
|
||||
selected_evidence_bytes,
|
||||
overhead_tokens,
|
||||
declared_evidence_limit,
|
||||
declared_file_limit,
|
||||
)
|
||||
if any(value < 0 for value in values):
|
||||
raise ModelBudgetError("model budgets must be non-negative")
|
||||
if input_limit > profile_input_limit:
|
||||
raise ModelBudgetError(
|
||||
"packet token limit exceeds profile "
|
||||
f"(declared={input_limit}, profile={profile_input_limit})"
|
||||
)
|
||||
if byte_limit > profile_byte_limit:
|
||||
raise ModelBudgetError(
|
||||
"packet byte limit exceeds profile "
|
||||
f"(declared={byte_limit}, profile={profile_byte_limit})"
|
||||
)
|
||||
if declared_evidence_limit > profile_evidence_limit:
|
||||
raise ModelBudgetError(
|
||||
"evidence byte limit exceeds profile "
|
||||
f"(declared={declared_evidence_limit}, profile={profile_evidence_limit})"
|
||||
)
|
||||
if declared_file_limit > profile_file_limit:
|
||||
raise ModelBudgetError(
|
||||
"evidence file limit exceeds profile "
|
||||
f"(declared={declared_file_limit}, profile={profile_file_limit})"
|
||||
)
|
||||
if max_output_tokens > profile_output_limit:
|
||||
raise ModelBudgetError(
|
||||
"output token budget exceeds profile "
|
||||
f"(requested={max_output_tokens}, profile={profile_output_limit})"
|
||||
)
|
||||
try:
|
||||
_validate_packet_metadata(packet, budget)
|
||||
except ModelBudgetError:
|
||||
raise
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise ModelBudgetError("packet budget metadata is malformed") from exc
|
||||
context_bytes = packet_bytes + selected_evidence_bytes
|
||||
if packet_bytes > byte_limit or context_bytes > byte_limit:
|
||||
raise ModelBudgetError("packet exceeds its declared byte budget")
|
||||
if selected_evidence_bytes > declared_evidence_limit:
|
||||
raise ModelBudgetError("selected evidence exceeds its evidence byte budget")
|
||||
if selected_evidence_bytes > profile_evidence_limit:
|
||||
raise ModelBudgetError("selected evidence exceeds profile byte budget")
|
||||
try:
|
||||
estimated_input = estimate_input_tokens(packet) if input_tokens is None else int(input_tokens)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ModelBudgetError("input token count must be an integer") from exc
|
||||
if estimated_input < 0:
|
||||
raise ModelBudgetError("input token count must be non-negative")
|
||||
if estimated_input > input_limit:
|
||||
raise ModelBudgetError(
|
||||
"packet input token budget exceeded "
|
||||
f"(estimated={estimated_input}, limit={input_limit})"
|
||||
)
|
||||
# Reserve the input actually sent for this request plus fixed provider
|
||||
# overhead. Reserving the full profile here rejects small packets for no
|
||||
# reason, while the packet's own input_limit still prevents later growth.
|
||||
effective_input = estimated_input + overhead_tokens
|
||||
available = context_window - safety_margin - max_output_tokens
|
||||
if effective_input > available:
|
||||
raise ModelBudgetError(
|
||||
"model context budget exceeded "
|
||||
f"(input={effective_input}, output={max_output_tokens}, "
|
||||
f"safety={safety_margin}, window={context_window})"
|
||||
)
|
||||
return {
|
||||
"context_window": context_window,
|
||||
"input_token_limit": input_limit,
|
||||
"estimated_input_tokens": estimated_input,
|
||||
"overhead_tokens": overhead_tokens,
|
||||
"estimated_request_input_tokens": estimated_input + overhead_tokens,
|
||||
"max_output_tokens": max_output_tokens,
|
||||
"safety_margin": safety_margin,
|
||||
"reserved_total_tokens": effective_input + max_output_tokens + safety_margin,
|
||||
"available_input_tokens": max(0, available - overhead_tokens),
|
||||
}
|
||||
|
||||
|
||||
def validate_response(
|
||||
output_tokens: int,
|
||||
*,
|
||||
max_output_tokens: int,
|
||||
finish_reason: str | None = None,
|
||||
) -> dict[str, int | str | None]:
|
||||
"""Validate the completion actually returned by a provider.
|
||||
|
||||
The preflight check protects the request, while this postflight check
|
||||
catches SDKs that ignore or reinterpret ``max_output_tokens``. A length
|
||||
finish is reported as truncation so callers can persist a checkpoint and
|
||||
continue with a fresh bounded request instead of replaying the old one.
|
||||
"""
|
||||
|
||||
try:
|
||||
actual = int(output_tokens)
|
||||
limit = int(max_output_tokens)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ModelBudgetError("response token counts must be integers") from exc
|
||||
if actual < 0 or limit < 0:
|
||||
raise ModelBudgetError("response token counts must be non-negative")
|
||||
if actual > limit:
|
||||
raise ModelBudgetError(
|
||||
"model response exceeded output token budget "
|
||||
f"(actual={actual}, limit={limit})"
|
||||
)
|
||||
return {
|
||||
"output_tokens": actual,
|
||||
"max_output_tokens": limit,
|
||||
"finish_reason": finish_reason,
|
||||
"truncated": finish_reason in {
|
||||
"length",
|
||||
"max_tokens",
|
||||
"max_output_tokens",
|
||||
"max_completion_tokens",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def validate_provider_response(
|
||||
response: Any,
|
||||
*,
|
||||
max_output_tokens: int,
|
||||
) -> dict[str, int | str | None | bool]:
|
||||
"""Validate usage and completion state returned by a provider.
|
||||
|
||||
``response`` may be a normal Responses API object or a streaming event
|
||||
envelope. Failed/incomplete responses are rejected as non-reusable so a
|
||||
caller cannot append partial output to the next request. A length-limit
|
||||
completion is returned as ``truncated`` and must be resumed from a fresh
|
||||
checkpoint.
|
||||
"""
|
||||
|
||||
response_object = _response_object(response)
|
||||
usage = extract_response_usage(response_object)
|
||||
status = response_object.get("status")
|
||||
incomplete = response_object.get("incomplete_details")
|
||||
finish_reason = None
|
||||
if isinstance(incomplete, dict):
|
||||
reason = incomplete.get("reason")
|
||||
if reason is not None:
|
||||
finish_reason = str(reason)
|
||||
if finish_reason is None and response_object.get("finish_reason") is not None:
|
||||
finish_reason = str(response_object["finish_reason"])
|
||||
if status is None and finish_reason is None:
|
||||
raise ModelBudgetError("provider response is missing completion status")
|
||||
result = validate_response(
|
||||
usage["output_tokens"],
|
||||
max_output_tokens=max_output_tokens,
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
if status in {"failed", "cancelled"}:
|
||||
raise ModelBudgetError(
|
||||
f"provider response is not reusable (status={status})"
|
||||
)
|
||||
if status == "incomplete" and not result["truncated"]:
|
||||
raise ModelBudgetError(
|
||||
"provider response is incomplete without a length checkpoint reason"
|
||||
)
|
||||
result.update(usage)
|
||||
result["status"] = str(status) if status is not None else None
|
||||
return result
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("packet", type=Path)
|
||||
parser.add_argument("--context-window", type=int, default=DEFAULT_CONTEXT_WINDOW)
|
||||
parser.add_argument("--max-output-tokens", type=int)
|
||||
parser.add_argument("--safety-margin", type=int, default=DEFAULT_SAFETY_MARGIN)
|
||||
parser.add_argument("--input-tokens", type=int, help="exact count from the provider tokenizer")
|
||||
parser.add_argument(
|
||||
"--overhead-tokens",
|
||||
type=int,
|
||||
required=True,
|
||||
help="system/tool/history tokens outside the context packet",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
try:
|
||||
with args.packet.open("r", encoding="utf-8") as handle:
|
||||
packet = json.load(handle)
|
||||
if not isinstance(packet, dict):
|
||||
raise ModelBudgetError("packet must be a JSON object")
|
||||
result = validate_request(
|
||||
packet,
|
||||
context_window=args.context_window,
|
||||
max_output_tokens=args.max_output_tokens,
|
||||
safety_margin=args.safety_margin,
|
||||
input_tokens=args.input_tokens,
|
||||
overhead_tokens=args.overhead_tokens,
|
||||
)
|
||||
except (OSError, UnicodeError, json.JSONDecodeError, ModelBudgetError) as exc:
|
||||
parser.error(str(exc))
|
||||
print(
|
||||
"model_budget=valid "
|
||||
f"input_tokens={result['input_token_limit']} "
|
||||
f"estimated_input={result['estimated_input_tokens']} "
|
||||
f"overhead={result['overhead_tokens']} "
|
||||
f"max_output={result['max_output_tokens']} "
|
||||
f"reserved_total={result['reserved_total_tokens']}"
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user