502 lines
19 KiB
Python
502 lines
19 KiB
Python
#!/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())
|