359 lines
14 KiB
Python
359 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate bounded, auditable task context packets.
|
|
|
|
The selector applies byte and count limits before writing a packet. It never
|
|
creates an over-budget packet for a later validator to reject.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
|
|
PROFILES = {
|
|
"doc": {"tokens": 2048, "bytes": 8192, "evidence_bytes": 2048, "files": 8},
|
|
"feature": {"tokens": 4096, "bytes": 16384, "evidence_bytes": 6144, "files": 16},
|
|
"contract": {"tokens": 6144, "bytes": 24576, "evidence_bytes": 10240, "files": 24},
|
|
"milestone": {"tokens": 8192, "bytes": 32768, "evidence_bytes": 14336, "files": 32},
|
|
}
|
|
|
|
SEGMENTS = ("header", "baseline", "scope", "contracts", "evidence", "risks", "next", "forbidden")
|
|
|
|
# Compact codes keep the audit trail inside the packet budget. The human
|
|
# explanation is stable and can be rendered by the reader without repeating
|
|
# large strings in every packet.
|
|
EXCLUSION_REASONS = {
|
|
"E1": "duplicate_path",
|
|
"E2": "evidence_item_bytes_exceed_budget",
|
|
"E3": "evidence_file_count_exceed_budget",
|
|
"E4": "evidence_bytes_exceed_budget",
|
|
"E5": "context_remaining_bytes_insufficient",
|
|
"E6": "token_budget_exceeded",
|
|
}
|
|
|
|
# Without a provider tokenizer, one UTF-8 byte per token is the only portable
|
|
# hard upper bound. The request validator can replace it with an exact count.
|
|
TOKEN_BYTES_PER_TOKEN = 1
|
|
|
|
|
|
class PacketError(ValueError):
|
|
"""Raised when a packet cannot be generated without violating a budget."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Candidate:
|
|
path: str
|
|
size: int
|
|
sha256: str | None
|
|
kind: str
|
|
priority: int
|
|
required: bool
|
|
|
|
|
|
def _json_bytes(value: Any) -> bytes:
|
|
return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
|
|
|
|
|
|
def _sha256_file(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _candidate(raw: dict[str, Any], base_dir: Path) -> Candidate:
|
|
path_value = str(raw.get("path", ""))
|
|
if not path_value:
|
|
raise PacketError("candidate is missing path")
|
|
path = Path(path_value)
|
|
if "size" in raw:
|
|
size = int(raw["size"])
|
|
else:
|
|
file_path = path if path.is_absolute() else base_dir / path
|
|
try:
|
|
size = file_path.stat().st_size
|
|
except OSError as exc:
|
|
raise PacketError(f"cannot stat evidence {path_value}: {exc}") from exc
|
|
if size < 0:
|
|
raise PacketError(f"evidence size must be non-negative: {path_value}")
|
|
sha256 = raw.get("sha256")
|
|
if sha256 is None and raw.get("hash_file", False):
|
|
file_path = path if path.is_absolute() else base_dir / path
|
|
sha256 = _sha256_file(file_path)
|
|
return Candidate(
|
|
path=path_value,
|
|
size=size,
|
|
sha256=str(sha256) if sha256 is not None else None,
|
|
kind=str(raw.get("kind", "artifact")),
|
|
priority=int(raw.get("priority", 0)),
|
|
required=bool(raw.get("required", False)),
|
|
)
|
|
|
|
|
|
def _evidence_item(candidate: Candidate) -> dict[str, Any]:
|
|
item: dict[str, Any] = {"path": candidate.path, "bytes": candidate.size, "kind": candidate.kind}
|
|
if candidate.sha256:
|
|
item["sha256"] = candidate.sha256
|
|
return item
|
|
|
|
|
|
def _packet_bytes(packet: dict[str, Any]) -> int:
|
|
return len(_json_bytes(packet))
|
|
|
|
|
|
def _context_bytes(packet: dict[str, Any], evidence_used: int) -> int:
|
|
"""Return the bytes consumed by the envelope and referenced evidence.
|
|
|
|
Evidence is referenced by path in the JSON envelope, but the referenced
|
|
payload is read into the task context by the consumer. Counting both
|
|
here prevents a packet that is small on disk from expanding past the
|
|
context limit when it is loaded.
|
|
"""
|
|
|
|
return _packet_bytes(packet) + evidence_used
|
|
|
|
|
|
def _estimate_tokens(packet_bytes: int, evidence_bytes: int) -> int:
|
|
"""Return a conservative, provider-neutral input token estimate."""
|
|
|
|
return math.ceil((packet_bytes + evidence_bytes) / TOKEN_BYTES_PER_TOKEN)
|
|
|
|
|
|
def _finalize_packet_size(packet: dict[str, Any]) -> int:
|
|
"""Store the self-consistent serialized size in the budget metadata."""
|
|
|
|
packet["budget"]["packet_bytes"] = 0
|
|
for _ in range(4):
|
|
packet_bytes = _packet_bytes(packet)
|
|
if packet["budget"]["packet_bytes"] == packet_bytes:
|
|
return packet_bytes
|
|
packet["budget"]["packet_bytes"] = packet_bytes
|
|
raise PacketError("could not stabilize serialized packet size")
|
|
|
|
|
|
def _base_packet(input_data: dict[str, Any], context_id: str, packet_id: str) -> dict[str, Any]:
|
|
packet: dict[str, Any] = {
|
|
"context_id": context_id,
|
|
"packet_id": packet_id,
|
|
"task": str(input_data.get("task", "")),
|
|
"header": dict(input_data.get("header", {})),
|
|
"baseline": dict(input_data.get("baseline", {})),
|
|
"scope": dict(input_data.get("scope", {})),
|
|
"contracts": list(input_data.get("contracts", [])),
|
|
"evidence": {"selected": [], "excluded": []},
|
|
"risks": list(input_data.get("risks", [])),
|
|
"next": list(input_data.get("next", []))[:3],
|
|
"forbidden": list(input_data.get("forbidden", [])),
|
|
}
|
|
return packet
|
|
|
|
|
|
def _budget_metadata(limits: dict[str, int], selected: list[dict[str, Any]], excluded: list[dict[str, Any]], evidence_used: int) -> dict[str, Any]:
|
|
return {
|
|
"profile": limits["profile"],
|
|
"token_limit": limits["tokens"],
|
|
"byte_limit": limits["bytes"],
|
|
"evidence_byte_limit": limits["evidence_bytes"],
|
|
"evidence_file_limit": limits["files"],
|
|
"selected_evidence_bytes": evidence_used,
|
|
"selected_evidence_files": len(selected),
|
|
"excluded_evidence_files": len(excluded),
|
|
# Reserve the complete decimal width of the configured limit while
|
|
# selecting. This keeps the selection check conservative.
|
|
"packet_bytes": limits["bytes"],
|
|
}
|
|
|
|
|
|
def generate_packet(
|
|
input_data: dict[str, Any],
|
|
*,
|
|
context_id: str,
|
|
packet_id: str,
|
|
profile: str = "feature",
|
|
token_limit: int | None = None,
|
|
total_bytes: int | None = None,
|
|
evidence_bytes: int | None = None,
|
|
max_files: int | None = None,
|
|
base_dir: Path = Path("."),
|
|
) -> dict[str, Any]:
|
|
"""Select evidence under all budgets and return a packet within its limit."""
|
|
|
|
if profile not in PROFILES:
|
|
raise PacketError(f"unknown profile: {profile}")
|
|
limits = dict(PROFILES[profile])
|
|
limits["profile"] = profile
|
|
# Overrides are intentionally one-way: a task may ask for a smaller
|
|
# packet, but it must not enlarge a profile and silently bypass the
|
|
# context-window policy enforced by the request validator.
|
|
overrides = {
|
|
"tokens": token_limit,
|
|
"bytes": total_bytes,
|
|
"evidence_bytes": evidence_bytes,
|
|
"files": max_files,
|
|
}
|
|
for key, value in overrides.items():
|
|
if value is None:
|
|
continue
|
|
value = int(value)
|
|
if value < 0:
|
|
raise PacketError("budgets must be non-negative")
|
|
if value > limits[key]:
|
|
raise PacketError(
|
|
f"{key} budget override exceeds profile "
|
|
f"(requested={value}, profile={limits[key]})"
|
|
)
|
|
limits[key] = value
|
|
if any(limits[key] < 0 for key in ("tokens", "bytes", "evidence_bytes", "files")):
|
|
raise PacketError("budgets must be non-negative")
|
|
|
|
raw_candidates = input_data.get("candidates", [])
|
|
if not isinstance(raw_candidates, list):
|
|
raise PacketError("candidates must be a list")
|
|
candidates = []
|
|
for raw in raw_candidates:
|
|
if not isinstance(raw, dict):
|
|
raise PacketError("each candidate must be an object")
|
|
candidates.append(_candidate(raw, base_dir))
|
|
candidates.sort(key=lambda item: (-int(item.required), -item.priority, item.path))
|
|
packet = _base_packet(input_data, context_id, packet_id)
|
|
skeleton_bytes = _packet_bytes(packet)
|
|
if skeleton_bytes > limits["bytes"]:
|
|
raise PacketError(
|
|
f"context skeleton exceeds byte budget ({skeleton_bytes}>{limits['bytes']}); "
|
|
"reduce task metadata or use a larger profile"
|
|
)
|
|
selected: list[dict[str, Any]] = []
|
|
excluded: list[dict[str, Any]] = []
|
|
evidence_used = 0
|
|
seen: set[str] = set()
|
|
|
|
for candidate in candidates:
|
|
if candidate.path in seen:
|
|
excluded.append({"path": candidate.path, "reason_code": "E1"})
|
|
continue
|
|
seen.add(candidate.path)
|
|
item = _evidence_item(candidate)
|
|
reason: str | None = None
|
|
if candidate.size > limits["evidence_bytes"]:
|
|
reason = "E2"
|
|
elif len(selected) >= limits["files"]:
|
|
reason = "E3"
|
|
elif evidence_used + candidate.size > limits["evidence_bytes"]:
|
|
reason = "E4"
|
|
else:
|
|
trial = dict(packet)
|
|
trial_selected = selected + [item]
|
|
trial["evidence"] = {"selected": trial_selected, "excluded": excluded}
|
|
trial["budget"] = _budget_metadata(
|
|
limits, trial_selected, excluded, evidence_used + candidate.size
|
|
)
|
|
# Finalize the trial before checking budgets. In particular, do
|
|
# not estimate tokens from the provisional byte-limit placeholder
|
|
# in budget.packet_bytes.
|
|
trial_packet_bytes = _finalize_packet_size(trial)
|
|
trial_context_bytes = trial_packet_bytes + evidence_used + candidate.size
|
|
if trial_context_bytes > limits["bytes"]:
|
|
reason = "E5"
|
|
elif _estimate_tokens(trial_packet_bytes, evidence_used + candidate.size) > limits["tokens"]:
|
|
reason = "E6"
|
|
if reason:
|
|
excluded.append({"path": candidate.path, "reason_code": reason})
|
|
if candidate.required:
|
|
raise PacketError(f"required evidence cannot fit ({reason}): {candidate.path}")
|
|
continue
|
|
selected.append(item)
|
|
evidence_used += candidate.size
|
|
|
|
# The exclusion audit trail is part of the packet, so it can consume the
|
|
# remaining context space. Remove optional selections until the complete
|
|
# packet (selected evidence, exclusions and budget metadata) fits.
|
|
while True:
|
|
packet["evidence"] = {"selected": selected, "excluded": excluded}
|
|
packet["budget"] = _budget_metadata(limits, selected, excluded, evidence_used)
|
|
packet_bytes = _finalize_packet_size(packet)
|
|
context_bytes = packet_bytes + evidence_used
|
|
over_bytes = packet_bytes > limits["bytes"] or context_bytes > limits["bytes"]
|
|
over_tokens = _estimate_tokens(packet_bytes, evidence_used) > limits["tokens"]
|
|
if (
|
|
not over_bytes
|
|
and not over_tokens
|
|
):
|
|
break
|
|
removable_index = next(
|
|
(index for index in range(len(selected) - 1, -1, -1)
|
|
if not any(item["path"] == selected[index]["path"] and candidate.required
|
|
for candidate in candidates)),
|
|
None,
|
|
)
|
|
if removable_index is None:
|
|
budget_kind = "byte" if over_bytes else "token"
|
|
raise PacketError(
|
|
f"context skeleton and required evidence exceed {budget_kind} budget"
|
|
)
|
|
removed = selected.pop(removable_index)
|
|
evidence_used -= int(removed["bytes"])
|
|
excluded.append({
|
|
"path": removed["path"],
|
|
"reason_code": "E5" if over_bytes else "E6",
|
|
})
|
|
return packet
|
|
|
|
|
|
def _load(path: Path) -> dict[str, Any]:
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
value = json.load(handle)
|
|
if not isinstance(value, dict):
|
|
raise PacketError("input must be a JSON object")
|
|
return value
|
|
|
|
|
|
def main(argv: Iterable[str] | None = None) -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--input", required=True, type=Path, help="candidate/context JSON")
|
|
parser.add_argument("--output", required=True, type=Path)
|
|
parser.add_argument("--context-id", required=True)
|
|
parser.add_argument("--packet-id", required=True)
|
|
parser.add_argument("--profile", choices=sorted(PROFILES), default="feature")
|
|
parser.add_argument("--base-dir", type=Path, default=Path("."))
|
|
parser.add_argument("--token-limit", type=int)
|
|
parser.add_argument("--total-bytes", type=int)
|
|
parser.add_argument("--evidence-bytes", type=int)
|
|
parser.add_argument("--max-files", type=int)
|
|
args = parser.parse_args(list(argv) if argv is not None else None)
|
|
try:
|
|
packet = generate_packet(
|
|
_load(args.input),
|
|
context_id=args.context_id,
|
|
packet_id=args.packet_id,
|
|
profile=args.profile,
|
|
token_limit=args.token_limit,
|
|
total_bytes=args.total_bytes,
|
|
evidence_bytes=args.evidence_bytes,
|
|
max_files=args.max_files,
|
|
base_dir=args.base_dir,
|
|
)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
payload = _json_bytes(packet)
|
|
# Do not expose a partially written packet to a concurrent validator.
|
|
temporary = args.output.with_name(f".{args.output.name}.tmp")
|
|
temporary.write_bytes(payload)
|
|
if len(payload) != packet["budget"]["packet_bytes"]:
|
|
raise PacketError("serialized packet size changed before write")
|
|
temporary.replace(args.output)
|
|
except (PacketError, OSError, TypeError, ValueError, json.JSONDecodeError) as exc:
|
|
print(f"context packet generation failed: {exc}", file=sys.stderr)
|
|
return 2
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|