208 lines
8.4 KiB
Python
208 lines
8.4 KiB
Python
#!/usr/bin/env python3
|
|
import json
|
|
import unittest
|
|
|
|
from validate_model_budget import (
|
|
ModelBudgetError,
|
|
extract_response_usage,
|
|
validate_provider_response,
|
|
validate_request,
|
|
validate_response,
|
|
)
|
|
|
|
|
|
def packet(
|
|
profile="feature",
|
|
token_limit=4096,
|
|
packet_bytes=1200,
|
|
evidence_bytes=300,
|
|
byte_limit=16384,
|
|
evidence_byte_limit=None,
|
|
):
|
|
budget = {
|
|
"profile": profile,
|
|
"token_limit": token_limit,
|
|
"byte_limit": byte_limit,
|
|
"packet_bytes": packet_bytes,
|
|
"selected_evidence_bytes": evidence_bytes,
|
|
}
|
|
if evidence_byte_limit is not None:
|
|
budget["evidence_byte_limit"] = evidence_byte_limit
|
|
return {
|
|
"budget": budget
|
|
}
|
|
|
|
|
|
class ModelBudgetTests(unittest.TestCase):
|
|
def test_profile_reserves_output_and_safety_margin(self):
|
|
result = validate_request(packet())
|
|
self.assertEqual(result["max_output_tokens"], 1024)
|
|
self.assertEqual(result["reserved_total_tokens"], 2780)
|
|
|
|
def test_non_packet_request_overhead_is_reserved(self):
|
|
result = validate_request(packet(), overhead_tokens=700)
|
|
self.assertEqual(result["overhead_tokens"], 700)
|
|
self.assertEqual(result["reserved_total_tokens"], 3480)
|
|
|
|
def test_request_overhead_can_exhaust_context_window(self):
|
|
result = validate_request(packet(), context_window=4000, overhead_tokens=700)
|
|
self.assertEqual(result["estimated_request_input_tokens"], 2200)
|
|
|
|
def test_context_window_overflow_fails_before_request(self):
|
|
with self.assertRaisesRegex(ModelBudgetError, "context budget exceeded"):
|
|
validate_request(packet(), context_window=1700)
|
|
|
|
def test_estimated_input_must_fit_declared_packet_budget(self):
|
|
with self.assertRaisesRegex(ModelBudgetError, "input token budget exceeded"):
|
|
validate_request(packet(token_limit=400), input_tokens=401)
|
|
|
|
def test_byte_upper_bound_prevents_unsafe_heuristic_underestimate(self):
|
|
with self.assertRaisesRegex(ModelBudgetError, "input token budget exceeded"):
|
|
validate_request(packet(token_limit=1499))
|
|
|
|
def test_exact_tokenizer_count_can_be_supplied(self):
|
|
result = validate_request(packet(token_limit=600), input_tokens=500)
|
|
self.assertEqual(result["estimated_input_tokens"], 500)
|
|
|
|
def test_explicit_output_budget_can_be_smaller(self):
|
|
result = validate_request(packet(), context_window=5000, max_output_tokens=256)
|
|
self.assertEqual(result["max_output_tokens"], 256)
|
|
|
|
def test_output_budget_cannot_expand_profile(self):
|
|
with self.assertRaisesRegex(ModelBudgetError, "output token budget exceeds profile"):
|
|
validate_request(packet(), context_window=20000, max_output_tokens=2048)
|
|
|
|
def test_input_budget_cannot_expand_profile(self):
|
|
with self.assertRaisesRegex(ModelBudgetError, "packet token limit exceeds profile"):
|
|
validate_request(packet(token_limit=5000), context_window=20000)
|
|
|
|
def test_byte_budget_cannot_expand_profile(self):
|
|
with self.assertRaisesRegex(ModelBudgetError, "byte limit exceeds profile"):
|
|
validate_request(packet(byte_limit=16385), context_window=20000)
|
|
|
|
def test_evidence_budget_is_checked_even_without_envelope(self):
|
|
with self.assertRaisesRegex(ModelBudgetError, "selected evidence exceeds its evidence byte budget"):
|
|
validate_request(packet(evidence_bytes=301, evidence_byte_limit=300))
|
|
|
|
def test_nested_accumulated_history_is_rejected(self):
|
|
value = packet()
|
|
value["scope"] = {"state": {"messages": ["old"]}}
|
|
with self.assertRaisesRegex(ModelBudgetError, "forbidden accumulated context"):
|
|
validate_request(value)
|
|
|
|
def test_full_packet_rejects_unknown_top_level_context(self):
|
|
value = packet()
|
|
value.update({"context_id": "c1", "packet_id": "p1", "tool_output": "log"})
|
|
with self.assertRaisesRegex(ModelBudgetError, "unknown top-level context field"):
|
|
validate_request(value)
|
|
|
|
def test_packet_byte_overflow_is_rejected(self):
|
|
with self.assertRaisesRegex(ModelBudgetError, "byte budget"):
|
|
validate_request(packet(packet_bytes=16000, evidence_bytes=1000))
|
|
|
|
def test_full_packet_rejects_stale_serialized_size(self):
|
|
value = packet()
|
|
value.update({
|
|
"context_id": "c1",
|
|
"packet_id": "p1",
|
|
"evidence": {"selected": [], "excluded": []},
|
|
})
|
|
with self.assertRaisesRegex(ModelBudgetError, "does not match serialized packet"):
|
|
validate_request(value)
|
|
|
|
def test_packet_without_ids_cannot_lie_about_serialized_size(self):
|
|
value = packet()
|
|
value["task"] = "x" * 20000
|
|
with self.assertRaisesRegex(ModelBudgetError, "does not match serialized packet"):
|
|
validate_request(value)
|
|
|
|
def test_full_packet_rejects_inconsistent_evidence_totals(self):
|
|
value = packet(packet_bytes=215, evidence_bytes=5)
|
|
value.update({
|
|
"context_id": "c1",
|
|
"packet_id": "p1",
|
|
"evidence": {"selected": [{"path": "x", "bytes": 4}], "excluded": []},
|
|
})
|
|
with self.assertRaisesRegex(ModelBudgetError, "selected evidence byte metadata"):
|
|
validate_request(value)
|
|
|
|
def test_full_packet_rejects_accumulated_history(self):
|
|
value = packet()
|
|
value["history"] = ["old response"]
|
|
with self.assertRaisesRegex(ModelBudgetError, "forbidden accumulated context"):
|
|
validate_request(value)
|
|
|
|
def test_full_packet_rejects_unknown_exclusion_reason(self):
|
|
value = packet(packet_bytes=203, evidence_bytes=0)
|
|
value.update({
|
|
"context_id": "c1",
|
|
"packet_id": "p1",
|
|
"evidence": {"selected": [], "excluded": [{"path": "x", "reason_code": "old"}]},
|
|
})
|
|
# Stabilize the serialized size after adding the evidence envelope.
|
|
for _ in range(4):
|
|
value["budget"]["packet_bytes"] = len(
|
|
json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
|
|
)
|
|
with self.assertRaisesRegex(ModelBudgetError, "unknown reason code"):
|
|
validate_request(value)
|
|
|
|
def test_response_overrun_is_rejected(self):
|
|
with self.assertRaisesRegex(ModelBudgetError, "response exceeded output token budget"):
|
|
validate_response(1025, max_output_tokens=1024)
|
|
|
|
def test_response_length_finish_is_checkpointable(self):
|
|
result = validate_response(1024, max_output_tokens=1024, finish_reason="length")
|
|
self.assertTrue(result["truncated"])
|
|
|
|
def test_responses_api_incomplete_reason_is_checkpointable(self):
|
|
result = validate_response(
|
|
1024,
|
|
max_output_tokens=1024,
|
|
finish_reason="max_output_tokens",
|
|
)
|
|
self.assertTrue(result["truncated"])
|
|
|
|
def test_stream_envelope_usage_is_validated(self):
|
|
event = {
|
|
"type": "response.completed",
|
|
"response": {
|
|
"status": "completed",
|
|
"usage": {
|
|
"input_tokens": 300,
|
|
"output_tokens": 40,
|
|
"total_tokens": 340,
|
|
},
|
|
},
|
|
}
|
|
self.assertEqual(extract_response_usage(event)["output_tokens"], 40)
|
|
result = validate_provider_response(event, max_output_tokens=64)
|
|
self.assertEqual(result["status"], "completed")
|
|
|
|
def test_provider_usage_total_cannot_be_inconsistent(self):
|
|
response = {
|
|
"status": "completed",
|
|
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 14},
|
|
}
|
|
with self.assertRaisesRegex(ModelBudgetError, "total is smaller"):
|
|
validate_provider_response(response, max_output_tokens=16)
|
|
|
|
def test_failed_provider_response_is_not_reusable(self):
|
|
response = {
|
|
"status": "failed",
|
|
"usage": {"input_tokens": 10, "output_tokens": 0, "total_tokens": 10},
|
|
}
|
|
with self.assertRaisesRegex(ModelBudgetError, "not reusable"):
|
|
validate_provider_response(response, max_output_tokens=16)
|
|
|
|
def test_provider_response_without_completion_status_is_rejected(self):
|
|
response = {
|
|
"usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11},
|
|
}
|
|
with self.assertRaisesRegex(ModelBudgetError, "missing completion status"):
|
|
validate_provider_response(response, max_output_tokens=16)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|