yaojingang
138 lines
4.2 KiB
Python
138 lines
4.2 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 == Path("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)
|
|
if any(rel == ignored or ignored in rel.parents for ignored in IGNORED_RELATIVE_DIRS):
|
|
return True
|
|
return len(rel.parts) >= 2 and rel.parts[0] == "tests" and rel.parts[1].startswith("tmp_")
|
|
|
|
|
|
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
|
|
rel = path.relative_to(skill_dir)
|
|
if path.suffix not in TEXT_EXTS and path.name != "SKILL.md":
|
|
size = path.stat().st_size
|
|
files.append({"path": str(rel), "kind": "binary_or_other", "bytes": size})
|
|
continue
|
|
kind = classify(rel)
|
|
if kind in {"binary_or_other", "asset"} and path.suffix not in TEXT_EXTS:
|
|
size = path.stat().st_size
|
|
files.append({"path": str(rel), "kind": kind, "bytes": size})
|
|
continue
|
|
text = read_text(path)
|
|
tokens = estimate_tokens(text)
|
|
record = {
|
|
"path": str(rel),
|
|
"kind": kind,
|
|
"chars": len(text),
|
|
"estimated_tokens": tokens,
|
|
}
|
|
files.append(record)
|
|
total_tokens += tokens
|
|
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()
|