202 lines
7.3 KiB
Python
202 lines
7.3 KiB
Python
|
"""
|
||
|
|
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 <book_dir> [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.")
|