/usr/bin/token-count

token-count.py

#!/usr/bin/env python3
"""
Simple script to count input tokens using Anthropic's count_tokens endpoint
Works both with piped input and command-line arguments
"""

import sys
import json
import os
import argparse
import requests


def get_content() -> str:
    """Get input either from stdin pipe or command line arguments"""
    # Check if anything is being piped in
    if not sys.stdin.isatty():
        return sys.stdin.read().rstrip()
    
    # Otherwise use command line arguments
    parser = argparse.ArgumentParser(description="Count tokens using Anthropic API")
    parser.add_argument("text", nargs="*", help="Text to count tokens for")
    args = parser.parse_args()
    
    if args.text:
        return " ".join(args.text)
    
    print("Error: No input provided", file=sys.stderr)
    print("Usage examples:", file=sys.stderr)
    print("  echo 'Hello world' | python count_tokens.py", file=sys.stderr)
    print("  python count_tokens.py 'Hello world how are you'", file=sys.stderr)
    sys.exit(1)


def main():
    api_key = os.environ.get("ANTHROPIC_API_KEY") or os.environ.get("ANTHROPIC_API_KEY_FIX")
    
    if not api_key:
        print("Error: ANTHROPIC_API_KEY environment variable not set", file=sys.stderr)
        sys.exit(1)

    content = get_content()

    payload = {
        "model": "claude-sonnet-4-6",           # ← change model name if needed
        "messages": [
            {"role": "user", "content": content}
        ]
    }

    headers = {
        "x-api-key": api_key,
        "anthropic-version": "2023-06-01",
        "content-type": "application/json"
    }

    try:
        response = requests.post(
            "https://api.anthropic.com/v1/messages/count_tokens",
            headers=headers,
            json=payload,
            timeout=10
        )
        
        response.raise_for_status()
        data = response.json()
        
        tokens = data.get("input_tokens")
        if tokens is not None:
            print(tokens)
        else:
            print("error", file=sys.stderr)
            print("Response:", json.dumps(data, indent=2), file=sys.stderr)
            sys.exit(1)

    except requests.RequestException as e:
        print("error", file=sys.stderr)
        print(f"Request failed: {e}", file=sys.stderr)
        if hasattr(e, 'response') and e.response is not None:
            print("Response content:", e.response.text, file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()