Files
yao-meta-skill/scripts/context_sizer.py
T
2026-04-01 09:46:51 +08:00

136 lines
4.1 KiB
Python

#!/usr/bin/env python3
import argparse
import json
from pathlib import Path
TEXT_EXTS = {".md", ".txt", ".yaml", ".yml", ".json", ".py", ".sh", ".js", ".ts"}
IGNORED_RELATIVE_DIRS = {
Path("reports") / "release_snapshots",
Path("tests") / "tmp",
Path("tests") / "tmp_snapshot",
Path("tests") / "tmp_cli",
}
PACKAGE_PATHS = (
"SKILL.md",
"manifest.json",
"agents",
"references",
"scripts",
"assets",
"evals",
"templates",
"reports",
"failures",
"tests",
"input",
"outputs",
)
def estimate_tokens(text: str) -> int:
# Fast heuristic suitable for local gating.
return max(1, len(text) // 4)
def read_text(path: Path) -> str:
try:
return path.read_text(encoding="utf-8")
except UnicodeDecodeError:
return path.read_text(encoding="utf-8", errors="ignore")
def classify(path: Path) -> str:
parts = set(path.parts)
if path.name == "SKILL.md":
return "skill_body"
if "agents" in parts:
return "interface"
if "references" in parts:
return "reference"
if "scripts" in parts:
return "script"
if "assets" in parts:
return "asset"
if path.suffix in TEXT_EXTS:
return "other_text"
return "binary_or_other"
def should_ignore(path: Path, skill_dir: Path) -> bool:
rel = path.relative_to(skill_dir)
return any(rel == ignored or ignored in rel.parents for ignored in IGNORED_RELATIVE_DIRS)
def summarize(skill_dir: Path) -> dict:
files = []
total_tokens = 0
initial_tokens = 0
candidate_files = []
for entry in PACKAGE_PATHS:
path = skill_dir / entry
if path.is_file():
candidate_files.append(path)
elif path.is_dir():
candidate_files.extend(sorted(file for file in path.rglob("*") if file.is_file()))
for path in candidate_files:
if should_ignore(path, skill_dir):
continue
if path.suffix not in TEXT_EXTS and path.name != "SKILL.md":
size = path.stat().st_size
files.append({"path": str(path.relative_to(skill_dir)), "kind": "binary_or_other", "bytes": size})
continue
kind = classify(path)
if kind in {"binary_or_other", "asset"} and path.suffix not in TEXT_EXTS:
size = path.stat().st_size
files.append({"path": str(path.relative_to(skill_dir)), "kind": kind, "bytes": size})
continue
text = read_text(path)
tokens = estimate_tokens(text)
record = {
"path": str(path.relative_to(skill_dir)),
"kind": kind,
"chars": len(text),
"estimated_tokens": tokens,
}
files.append(record)
total_tokens += tokens
rel = path.relative_to(skill_dir)
if rel == Path("SKILL.md") or (rel.parts and rel.parts[0] == "agents"):
initial_tokens += tokens
return {
"skill_dir": str(skill_dir),
"estimated_initial_load_tokens": initial_tokens,
"estimated_total_text_tokens": total_tokens,
"warning": initial_tokens > 2000,
"files": files,
}
def main() -> None:
parser = argparse.ArgumentParser(description="Estimate context size for a skill package.")
parser.add_argument("skill_dir", help="Path to the skill directory")
parser.add_argument("--json", action="store_true", help="Emit machine-readable JSON")
args = parser.parse_args()
report = summarize(Path(args.skill_dir).resolve())
if args.json:
print(json.dumps(report, ensure_ascii=False, indent=2))
return
print(f"Skill: {report['skill_dir']}")
print(f"Estimated initial-load tokens: {report['estimated_initial_load_tokens']}")
print(f"Estimated total text tokens: {report['estimated_total_text_tokens']}")
print(f"Initial-load warning (>2000): {'YES' if report['warning'] else 'NO'}")
print("")
for file in report["files"]:
if "estimated_tokens" in file:
print(f"{file['kind']:12} {file['estimated_tokens']:6}t {file['path']}")
else:
print(f"{file['kind']:12} {file['bytes']:6}b {file['path']}")
if __name__ == "__main__":
main()