""" Token Tracker - Track token consumption for novel writing process Records every step's token usage and generates reports """ import json from datetime import datetime from pathlib import Path class TokenTracker: """Track and report token consumption for novel writing""" def __init__(self, book_dir): self.book_dir = Path(book_dir) self.report_file = self.book_dir / "token-report.json" self.report = self._load_report() def _load_report(self): """Load existing report or create new one""" if self.report_file.exists(): with open(self.report_file, 'r', encoding='utf-8') as f: return json.load(f) # Initialize new report config_file = self.book_dir / "0-book-config.json" book_title = "Unknown" if config_file.exists(): with open(config_file, 'r', encoding='utf-8') as f: config = json.load(f) book_title = config.get("title", "Unknown") return { "book_title": book_title, "book_dir": str(self.book_dir), "created_at": datetime.now().isoformat(), "total_tokens": { "prompt": 0, "completion": 0, "total": 0 }, "total_cost_usd": 0.0, "by_phase": {}, "by_chapter": {}, "steps": [] } def _save_report(self): """Save report to file""" self.report["updated_at"] = datetime.now().isoformat() with open(self.report_file, 'w', encoding='utf-8') as f: json.dump(self.report, f, indent=2, ensure_ascii=False) def record_step(self, phase, step_name, prompt_tokens, completion_tokens, chapter_num=None, notes=None): """ Record token usage for a step Args: phase: One of [outline, worldbuilding, chapter_writing, quality_check, merge] step_name: Description of the step prompt_tokens: Input tokens completion_tokens: Output tokens chapter_num: If applicable, which chapter notes: Additional notes """ step_record = { "timestamp": datetime.now().isoformat(), "phase": phase, "step_name": step_name, "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": prompt_tokens + completion_tokens, "chapter_num": chapter_num, "notes": notes } # Add to steps list self.report["steps"].append(step_record) # Update total self.report["total_tokens"]["prompt"] += prompt_tokens self.report["total_tokens"]["completion"] += completion_tokens self.report["total_tokens"]["total"] += prompt_tokens + completion_tokens # Update by_phase if phase not in self.report["by_phase"]: self.report["by_phase"][phase] = { "prompt": 0, "completion": 0, "total": 0, "steps": 0 } self.report["by_phase"][phase]["prompt"] += prompt_tokens self.report["by_phase"][phase]["completion"] += completion_tokens self.report["by_phase"][phase]["total"] += prompt_tokens + completion_tokens self.report["by_phase"][phase]["steps"] += 1 # Update by_chapter if applicable if chapter_num is not None: ch_key = f"chapter_{chapter_num}" if ch_key not in self.report["by_chapter"]: self.report["by_chapter"][ch_key] = { "prompt": 0, "completion": 0, "total": 0, "steps": 0 } self.report["by_chapter"][ch_key]["prompt"] += prompt_tokens self.report["by_chapter"][ch_key]["completion"] += completion_tokens self.report["by_chapter"][ch_key]["total"] += prompt_tokens + completion_tokens self.report["by_chapter"][ch_key]["steps"] += 1 # Estimate cost (approximate rates) step_cost = self._estimate_cost(prompt_tokens, completion_tokens) self.report["total_cost_usd"] += step_cost step_record["estimated_cost_usd"] = step_cost self._save_report() return step_record def _estimate_cost(self, prompt_tokens, completion_tokens): """Estimate cost in USD (approximate rates for Claude/GPT-4 class models)""" # Approximate: $3 per 1M input tokens, $15 per 1M output tokens prompt_cost = (prompt_tokens / 1000000) * 3.0 completion_cost = (completion_tokens / 1000000) * 15.0 return round(prompt_cost + completion_cost, 4) def get_report(self): """Get current report""" return self.report def print_summary(self): """Print human-readable summary""" print("=" * 60) print(f"《{self.report['book_title']}》Token消耗报告") print("=" * 60) print(f"总Token数: {self.report['total_tokens']['total']:,}") print(f" - 输入: {self.report['total_tokens']['prompt']:,}") print(f" - 输出: {self.report['total_tokens']['completion']:,}") print(f"预估成本: ${self.report['total_cost_usd']:.2f} USD") print() print("分阶段统计:") print("-" * 40) for phase, stats in self.report["by_phase"].items(): print(f" {phase}: {stats['total']:,} tokens ({stats['steps']} 步)") if self.report["by_chapter"]: print() print("按章节统计 (Top 5):") print("-" * 40) sorted_chapters = sorted( self.report["by_chapter"].items(), key=lambda x: x[1]['total'], reverse=True )[:5] for ch_key, stats in sorted_chapters: print(f" {ch_key}: {stats['total']:,} tokens") print("=" * 60) def create_tracker(book_dir): """Factory function to create tracker""" return TokenTracker(book_dir) def record_tokens(book_dir, phase, step_name, prompt_tokens, completion_tokens, chapter_num=None, notes=None): """Quick record function""" tracker = TokenTracker(book_dir) return tracker.record_step(phase, step_name, prompt_tokens, completion_tokens, chapter_num, notes) def get_report_for_book(book_dir): """Get report for a specific book""" tracker = TokenTracker(book_dir) return tracker.get_report() if __name__ == "__main__": import sys if len(sys.argv) < 2: print("Usage: token_tracker.py [summary]") print(" book_dir: Path to book directory") print(" summary: Print summary report") sys.exit(1) book_dir = sys.argv[1] tracker = TokenTracker(book_dir) if len(sys.argv) > 2 and sys.argv[2] == "summary": tracker.print_summary() else: # Example: Record a step tracker.record_step( phase="chapter_writing", step_name="Generate detailed outline for Chapter 1", prompt_tokens=3500, completion_tokens=2800, chapter_num=1, notes="First attempt" ) print("Token usage recorded.")