count_tokens.py 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. #!/usr/bin/env python3
  2. # vendored from bmad-workflow-builder/scripts; canonical source there
  3. # /// script
  4. # requires-python = ">=3.9"
  5. # dependencies = ["tiktoken"]
  6. # ///
  7. """count_tokens — the single length metric for skill authoring.
  8. Token counts replace line counts everywhere in the builder and eval-runner.
  9. This script reports the token length of a file or of text piped on stdin, using
  10. the tiktoken cl100k_base encoding. When tiktoken is not installed it falls back
  11. to a character-based estimate (len(text) // 4) and says so, so the script always
  12. runs under a bare python3 even with no third-party packages present.
  13. Usage:
  14. count_tokens.py <file> count the tokens in a file
  15. count_tokens.py --stdin count the tokens read from stdin
  16. Output (one line of JSON on stdout):
  17. {"tokens": <int>, "method": "tiktoken"} when tiktoken loaded
  18. {"tokens": <int>, "method": "fallback"} when it fell back to chars // 4
  19. Budgets this feeds: SKILL.md ~1500-2500, multi-branch reference ~4500,
  20. single-purpose reference ~9000.
  21. """
  22. import argparse
  23. import json
  24. import sys
  25. ENCODING = "cl100k_base"
  26. def count_tokens(text: str) -> tuple[int, str]:
  27. """Return (token_count, method).
  28. Tries tiktoken's cl100k_base encoding first. If tiktoken cannot be imported
  29. or initialized, estimates with len(text) // 4 and reports method "fallback".
  30. """
  31. try:
  32. import tiktoken
  33. except Exception:
  34. return len(text) // 4, "fallback"
  35. try:
  36. enc = tiktoken.get_encoding(ENCODING)
  37. except Exception:
  38. return len(text) // 4, "fallback"
  39. return len(enc.encode(text)), "tiktoken"
  40. def read_input(args) -> str:
  41. if args.stdin:
  42. return sys.stdin.read()
  43. with open(args.file, encoding="utf-8") as f:
  44. return f.read()
  45. def main(argv: list[str] | None = None) -> int:
  46. p = argparse.ArgumentParser(
  47. description=__doc__,
  48. formatter_class=argparse.RawDescriptionHelpFormatter,
  49. )
  50. p.add_argument("file", nargs="?", help="path to the file to count")
  51. p.add_argument("--stdin", action="store_true", help="read text from stdin instead of a file")
  52. args = p.parse_args(argv)
  53. if not args.stdin and not args.file:
  54. p.error("provide a file path or --stdin")
  55. if args.stdin and args.file:
  56. p.error("provide either a file path or --stdin, not both")
  57. text = read_input(args)
  58. tokens, method = count_tokens(text)
  59. print(json.dumps({"tokens": tokens, "method": method}))
  60. return 0
  61. if __name__ == "__main__":
  62. sys.exit(main())