| 1 | #!/usr/bin/env python3 |
| 2 | from __future__ import annotations |
| 3 | |
| 4 | import argparse |
| 5 | import json |
| 6 | import math |
| 7 | import re |
| 8 | import sys |
| 9 | from datetime import UTC, datetime |
| 10 | from pathlib import Path |
| 11 | |
| 12 | from scripts.model_pricing import estimate_cost_usd |
| 13 | |
| 14 | ROOT = Path(__file__).resolve().parent.parent |
| 15 | DEFAULT_USAGE_FILE = ROOT / "data" / "metrics" / "token-usage.jsonl" |
| 16 | CHARS_PER_TOKEN = 4 |
| 17 | |
| 18 | |
| 19 | def parse_args(argv: list[str] | None = None) -> argparse.Namespace: |
| 20 | parser = argparse.ArgumentParser( |
| 21 | description="Track token usage and estimated cost per pipeline run." |
| 22 | ) |
| 23 | parser.add_argument( |
| 24 | "--stage", required=True, help="Pipeline stage (for example: analysis, reskill)." |
| 25 | ) |
| 26 | parser.add_argument( |
| 27 | "--source", |
| 28 | required=True, |
| 29 | help="Execution source (for example: copilot-cli, github-models).", |
| 30 | ) |
| 31 | parser.add_argument("--model", required=True, help="Model name used for cost rates.") |
| 32 | parser.add_argument("--current-datetime", required=True, help="ISO-8601 timestamp for the run.") |
| 33 | parser.add_argument( |
| 34 | "--week", help="Week slug (YYYY-WNN). If omitted, inferred from current datetime." |
| 35 | ) |
| 36 | parser.add_argument( |
| 37 | "--prompt-file", type=Path, help="Prompt file used to estimate input tokens." |
| 38 | ) |
| 39 | parser.add_argument( |
| 40 | "--output-file", type=Path, help="Output file used to estimate output tokens." |
| 41 | ) |
| 42 | parser.add_argument("--input-tokens", type=int, help="Explicit input token count.") |
| 43 | parser.add_argument("--output-tokens", type=int, help="Explicit output token count.") |
| 44 | parser.add_argument( |
| 45 | "--transcript", |
| 46 | type=Path, |
| 47 | help="Copilot CLI --share transcript file for parsing token usage.", |
| 48 | ) |
| 49 | parser.add_argument( |
| 50 | "--api-response", |
| 51 | type=Path, |
| 52 | help="GitHub Models API response JSON for extracting usage data.", |
| 53 | ) |
| 54 | parser.add_argument( |
| 55 | "--input-manifest", |
| 56 | type=Path, |
| 57 | help="analysis-input-manifest JSON used to validate final prompt input tokens within 10%.", |
| 58 | ) |
| 59 | parser.add_argument( |
| 60 | "--usage-file", type=Path, default=DEFAULT_USAGE_FILE, help="JSONL path for usage ledger." |
| 61 | ) |
| 62 | return parser.parse_args(argv) |
| 63 | |
| 64 | |
| 65 | def parse_datetime(value: str) -> datetime: |
| 66 | candidate = value.strip() |
| 67 | if candidate.endswith("Z"): |
| 68 | candidate = f"{candidate[:-1]}+00:00" |
| 69 | parsed = datetime.fromisoformat(candidate) |
| 70 | return parsed if parsed.tzinfo else parsed.replace(tzinfo=UTC) |
| 71 | |
| 72 | |
| 73 | def week_slug(value: datetime) -> str: |
| 74 | year, week, _ = value.isocalendar() |
| 75 | return f"{year}-W{week:02d}" |
| 76 | |
| 77 | |
| 78 | def estimate_tokens_from_text(text: str) -> int: |
| 79 | stripped = text.strip() |
| 80 | if not stripped: |
| 81 | return 0 |
| 82 | return max(1, math.ceil(len(stripped) / CHARS_PER_TOKEN)) |
| 83 | |
| 84 | |
| 85 | def estimate_tokens_from_path(path: Path | None) -> int: |
| 86 | if path is None or not path.exists(): |
| 87 | return 0 |
| 88 | return estimate_tokens_from_text(path.read_text(encoding="utf-8")) |
| 89 | |
| 90 | |
| 91 | def parse_copilot_transcript(path: Path) -> tuple[int, int] | None: |
| 92 | """Parse a Copilot CLI --share transcript for token usage metadata. |
| 93 | |
| 94 | Searches for patterns like: |
| 95 | - "Input tokens: 1234" / "Output tokens: 567" |
| 96 | - "prompt_tokens: 1234" / "completion_tokens: 567" |
| 97 | - "Tokens used: 1234 input, 567 output" |
| 98 | Returns (input_tokens, output_tokens) or None if not found. |
| 99 | """ |
| 100 | if not path.exists(): |
| 101 | return None |
| 102 | try: |
| 103 | text = path.read_text(encoding="utf-8") |
| 104 | except (OSError, UnicodeDecodeError): |
| 105 | return None |
| 106 | |
| 107 | # Pattern: "Input tokens: N" and "Output tokens: N" |
| 108 | m_input = re.search(r"[Ii]nput[\s_]tokens[\s:]+(\d+)", text) |
| 109 | m_output = re.search(r"[Oo]utput[\s_]tokens[\s:]+(\d+)", text) |
| 110 | if m_input and m_output: |
| 111 | return int(m_input.group(1)), int(m_output.group(1)) |
| 112 | |
| 113 | # Pattern: "prompt_tokens: N" and "completion_tokens: N" |
| 114 | m_prompt = re.search(r"prompt_tokens[\"'\s:]+(\d+)", text) |
| 115 | m_completion = re.search(r"completion_tokens[\"'\s:]+(\d+)", text) |
| 116 | if m_prompt and m_completion: |
| 117 | return int(m_prompt.group(1)), int(m_completion.group(1)) |
| 118 | |
| 119 | # Pattern: "Tokens used: N input, N output" |
| 120 | m_combined = re.search(r"[Tt]okens\s+used[\s:]+(\d+)\s+input[,;\s]+(\d+)\s+output", text) |
| 121 | if m_combined: |
| 122 | return int(m_combined.group(1)), int(m_combined.group(2)) |
| 123 | |
| 124 | # Pattern: "Usage: N/N tokens (input/output)" |
| 125 | m_usage = re.search(r"[Uu]sage[\s:]+(\d+)\s*/\s*(\d+)\s*tokens", text) |
| 126 | if m_usage: |
| 127 | return int(m_usage.group(1)), int(m_usage.group(2)) |
| 128 | |
| 129 | return None |
| 130 | |
| 131 | |
| 132 | def parse_api_response(path: Path) -> tuple[int, int] | None: |
| 133 | """Parse a GitHub Models API response JSON for usage data. |
| 134 | |
| 135 | Expects OpenAI-compatible format with usage.prompt_tokens and |
| 136 | usage.completion_tokens fields. |
| 137 | Returns (input_tokens, output_tokens) or None if not found. |
| 138 | """ |
| 139 | if not path.exists(): |
| 140 | return None |
| 141 | try: |
| 142 | data = json.loads(path.read_text(encoding="utf-8")) |
| 143 | except (OSError, UnicodeDecodeError, json.JSONDecodeError): |
| 144 | return None |
| 145 | |
| 146 | usage = data.get("usage") if isinstance(data, dict) else None |
| 147 | if not isinstance(usage, dict): |
| 148 | return None |
| 149 | |
| 150 | prompt_tokens = usage.get("prompt_tokens") |
| 151 | completion_tokens = usage.get("completion_tokens") |
| 152 | if isinstance(prompt_tokens, int) and isinstance(completion_tokens, int): |
| 153 | return prompt_tokens, completion_tokens |
| 154 | |
| 155 | return None |
| 156 | |
| 157 | |
| 158 | def build_record(args: argparse.Namespace) -> dict[str, object]: |
| 159 | parsed_datetime = parse_datetime(args.current_datetime).astimezone(UTC) |
| 160 | week = args.week or week_slug(parsed_datetime) |
| 161 | |
| 162 | # Priority: 1) explicit flags, 2) transcript/api-response, 3) file-size estimate |
| 163 | estimated = True |
| 164 | input_tokens: int | None = None |
| 165 | output_tokens: int | None = None |
| 166 | |
| 167 | # Highest priority: explicit --input-tokens / --output-tokens |
| 168 | if args.input_tokens is not None and args.output_tokens is not None: |
| 169 | input_tokens = args.input_tokens |
| 170 | output_tokens = args.output_tokens |
| 171 | estimated = False |
| 172 | |
| 173 | # Second priority: parsed from transcript or API response |
| 174 | if input_tokens is None or output_tokens is None: |
| 175 | parsed = None |
| 176 | transcript_path = getattr(args, "transcript", None) |
| 177 | api_response_path = getattr(args, "api_response", None) |
| 178 | if transcript_path is not None: |
| 179 | parsed = parse_copilot_transcript(transcript_path) |
| 180 | if parsed is None and api_response_path is not None: |
| 181 | parsed = parse_api_response(api_response_path) |
| 182 | if parsed is not None: |
| 183 | input_tokens = parsed[0] |
| 184 | output_tokens = parsed[1] |
| 185 | estimated = False |
| 186 | |
| 187 | # Lowest priority: file-size estimation |
| 188 | if input_tokens is None: |
| 189 | input_tokens = estimate_tokens_from_path(args.prompt_file) |
| 190 | if output_tokens is None: |
| 191 | output_tokens = estimate_tokens_from_path(args.output_file) |
| 192 | |
| 193 | cost = estimate_cost_usd(args.model, input_tokens, output_tokens) |
| 194 | record: dict[str, object] = { |
| 195 | "timestamp": parsed_datetime.isoformat().replace("+00:00", "Z"), |
| 196 | "month": parsed_datetime.strftime("%Y-%m"), |
| 197 | "week": week, |
| 198 | "stage": args.stage, |
| 199 | "source": args.source, |
| 200 | "model": args.model, |
| 201 | "input_tokens": input_tokens, |
| 202 | "output_tokens": output_tokens, |
| 203 | "total_tokens": input_tokens + output_tokens, |
| 204 | "cost_usd": cost, |
| 205 | "estimated": estimated, |
| 206 | } |
| 207 | validation = validate_input_manifest(args.input_manifest, input_tokens) |
| 208 | if validation is not None: |
| 209 | record["input_manifest_validation"] = validation |
| 210 | return record |
| 211 | |
| 212 | |
| 213 | def _manifest_prompt_tokens(manifest: dict[str, object]) -> int | None: |
| 214 | rendered = manifest.get("rendered_prompt_estimate") |
| 215 | if isinstance(rendered, dict) and isinstance(rendered.get("tokens"), int): |
| 216 | return int(rendered["tokens"]) |
| 217 | value = manifest.get("prompt_tokens") |
| 218 | return int(value) if isinstance(value, int) else None |
| 219 | |
| 220 | |
| 221 | def validate_input_manifest(path: Path | None, input_tokens: int) -> dict[str, object] | None: |
| 222 | if path is None: |
| 223 | return None |
| 224 | manifest = json.loads(path.read_text(encoding="utf-8")) |
| 225 | if not isinstance(manifest, dict): |
| 226 | raise ValueError(f"Input manifest must be an object: {path}") |
| 227 | estimated_tokens = _manifest_prompt_tokens(manifest) |
| 228 | if estimated_tokens is None: |
| 229 | raise ValueError(f"Input manifest missing rendered prompt token estimate: {path}") |
| 230 | delta = abs(input_tokens - estimated_tokens) |
| 231 | ratio = delta / max(input_tokens, 1) |
| 232 | degraded = bool(manifest.get("degraded")) or not bool( |
| 233 | manifest.get("prompt_within_budget", True) |
| 234 | ) |
| 235 | passed = ratio <= 0.10 |
| 236 | reason = None |
| 237 | if not passed: |
| 238 | reason = ( |
| 239 | f"Final input usage differs from manifest by {ratio:.1%} " |
| 240 | f"({input_tokens} actual vs {estimated_tokens} estimated)." |
| 241 | ) |
| 242 | if degraded: |
| 243 | reason += ( |
| 244 | " Manifest is degraded/compacted, so the run is already marked candidate-only." |
| 245 | ) |
| 246 | return { |
| 247 | "manifest_path": path.as_posix(), |
| 248 | "estimated_input_tokens": estimated_tokens, |
| 249 | "actual_input_tokens": input_tokens, |
| 250 | "delta_tokens": delta, |
| 251 | "delta_ratio": round(ratio, 6), |
| 252 | "within_10_percent": passed, |
| 253 | "degraded_or_compacted": degraded, |
| 254 | "reason": reason, |
| 255 | } |
| 256 | |
| 257 | |
| 258 | def append_record(path: Path, record: dict[str, object]) -> None: |
| 259 | path.parent.mkdir(parents=True, exist_ok=True) |
| 260 | with path.open("a", encoding="utf-8") as handle: |
| 261 | handle.write(json.dumps(record, ensure_ascii=True)) |
| 262 | handle.write("\n") |
| 263 | |
| 264 | |
| 265 | def main(argv: list[str] | None = None) -> int: |
| 266 | args = parse_args(argv) |
| 267 | try: |
| 268 | record = build_record(args) |
| 269 | except (OSError, json.JSONDecodeError, ValueError) as exc: |
| 270 | print(f"::error::Token usage manifest validation failed: {exc}", file=sys.stderr) |
| 271 | return 1 |
| 272 | validation = record.get("input_manifest_validation") |
| 273 | if ( |
| 274 | isinstance(validation, dict) |
| 275 | and not validation.get("within_10_percent") |
| 276 | and not validation.get("degraded_or_compacted") |
| 277 | ): |
| 278 | print(f"::error::{validation.get('reason')}", file=sys.stderr) |
| 279 | return 1 |
| 280 | append_record(args.usage_file, record) |
| 281 | print(json.dumps(record, indent=2)) |
| 282 | return 0 |
| 283 | |
| 284 | |
| 285 | if __name__ == "__main__": |
| 286 | raise SystemExit(main()) |