cost_estimator.py

#!/usr/bin/env python3
"""
Cost Estimator for EPUB LLM Cleaner

Estimates API costs before processing based on:
- Token counts in the EPUB
- Number of chapters/API calls
- Selected model and workflow
"""

import zipfile
import re
from pathlib import Path
from bs4 import BeautifulSoup
from typing import Optional

# Pricing per million tokens (as of 2024)
# Update these as Anthropic adjusts pricing
MODEL_PRICING = {
    "claude-sonnet-4-5-20250929": {
        "input": 3.00,    # $3 per 1M input tokens
        "output": 15.00,  # $15 per 1M output tokens
        "name": "Claude Sonnet 4.5"
    },
    "claude-opus-4-5-20250929": {
        "input": 15.00,   # $15 per 1M input tokens
        "output": 75.00,  # $75 per 1M output tokens
        "name": "Claude Opus 4.5"
    },
    "claude-haiku-3-5-20241022": {
        "input": 0.80,    # $0.80 per 1M input tokens
        "output": 4.00,   # $4 per 1M output tokens
        "name": "Claude Haiku 3.5"
    },
    # Aliases for convenience
    "sonnet": {
        "input": 3.00,
        "output": 15.00,
        "name": "Claude Sonnet"
    },
    "opus": {
        "input": 15.00,
        "output": 75.00,
        "name": "Claude Opus"
    },
    "haiku": {
        "input": 0.80,
        "output": 4.00,
        "name": "Claude Haiku"
    }
}

# Workflow multipliers (how many API calls per chapter)
WORKFLOW_MULTIPLIERS = {
    "cleanup": {
        "analysis_calls": 1,      # Quick analysis pass
        "cleaning_calls": 0.3,    # ~30% of chapters need cleaning
        "output_ratio": 0.1,      # Output is ~10% of input (just changed paragraphs)
        "description": "OCR/formatting fixes - minimal API usage"
    },
    "filter": {
        "analysis_calls": 1,
        "cleaning_calls": 0.5,    # ~50% of chapters may need filtering
        "output_ratio": 0.15,
        "description": "Content filtering - moderate API usage"
    },
    "modernize": {
        "analysis_calls": 1,
        "cleaning_calls": 0.8,    # Most chapters need some modernization
        "output_ratio": 0.3,
        "description": "Language modernization - higher API usage"
    },
    "transform": {
        "analysis_calls": 1,
        "cleaning_calls": 1.0,    # All chapters need transformation
        "output_ratio": 0.8,      # Nearly full rewrite
        "description": "Major transformation - highest API usage"
    },
    "annotate": {
        "analysis_calls": 1,
        "cleaning_calls": 1.0,    # Generate annotations for all chapters
        "output_ratio": 0.2,      # Annotations are ~20% of original length
        "description": "Add commentary - moderate API usage"
    }
}

# Additional feature costs
FEATURE_COSTS = {
    "style_profile": {
        "input_tokens": 50000,    # Sample ~50k tokens for profiling
        "output_tokens": 5000,    # Profile output
        "description": "Build author style profile"
    },
    "book_analysis": {
        "input_tokens": 100000,   # Full book for analysis
        "output_tokens": 10000,   # Book model output
        "description": "Extract book model (characters, plot, timeline)"
    },
    "change_planning": {
        "input_tokens": 20000,    # Book model + change request
        "output_tokens": 5000,    # Change plan
        "description": "Generate modification plan"
    },
    "consistency_check": {
        "input_tokens": 30000,    # Sample chapters + changes
        "output_tokens": 3000,    # Consistency report
        "description": "Cross-chapter validation"
    },
    "drift_validation": {
        "input_tokens": 10000,    # Per-chapter validation
        "output_tokens": 1000,    # Drift scores
        "calls_per_chapter": 1,
        "description": "Style drift measurement"
    }
}


def estimate_tokens(text: str) -> int:
    """
    Estimate token count for text.
    Rule of thumb: ~4 characters per token for English text.
    Claude's actual tokenizer may vary, but this is close enough for estimates.
    """
    # More accurate: ~0.75 tokens per word, or ~4 chars per token
    return len(text) // 4


def extract_epub_stats(epub_path: str) -> dict:
    """Extract statistics from an EPUB file."""
    stats = {
        "total_chars": 0,
        "total_words": 0,
        "estimated_tokens": 0,
        "chapter_count": 0,
        "chapters": []
    }

    with zipfile.ZipFile(epub_path, 'r') as zf:
        for name in zf.namelist():
            if name.endswith(('.html', '.xhtml', '.htm')):
                content = zf.read(name).decode('utf-8', errors='ignore')
                soup = BeautifulSoup(content, 'html.parser')

                # Extract text
                text = soup.get_text(separator=' ', strip=True)
                if len(text) < 100:  # Skip near-empty files
                    continue

                char_count = len(text)
                word_count = len(text.split())
                token_estimate = estimate_tokens(text)

                stats["total_chars"] += char_count
                stats["total_words"] += word_count
                stats["estimated_tokens"] += token_estimate
                stats["chapter_count"] += 1
                stats["chapters"].append({
                    "file": name,
                    "chars": char_count,
                    "words": word_count,
                    "tokens": token_estimate
                })

    return stats


def estimate_cost(
    epub_path: str,
    model: str = "sonnet",
    workflow: str = "cleanup",
    with_profile: bool = False,
    with_book_analysis: bool = False,
    with_change_plan: bool = False,
    with_consistency_check: bool = False,
    with_drift_validation: bool = False,
    verbose: bool = True
) -> dict:
    """
    Estimate the API cost for processing an EPUB.

    Returns a dict with detailed cost breakdown.
    """
    # Get model pricing
    if model not in MODEL_PRICING:
        # Try to match partial model name
        for key in MODEL_PRICING:
            if model.lower() in key.lower():
                model = key
                break
        else:
            model = "sonnet"  # Default

    pricing = MODEL_PRICING[model]
    workflow_config = WORKFLOW_MULTIPLIERS.get(workflow, WORKFLOW_MULTIPLIERS["cleanup"])

    # Extract EPUB stats
    stats = extract_epub_stats(epub_path)

    if verbose:
        print(f"\n{'='*60}")
        print(f"COST ESTIMATE: {Path(epub_path).name}")
        print(f"{'='*60}")
        print(f"\nBook Statistics:")
        print(f"  Chapters: {stats['chapter_count']}")
        print(f"  Words: {stats['total_words']:,}")
        print(f"  Estimated tokens: {stats['estimated_tokens']:,}")
        print(f"\nModel: {pricing['name']}")
        print(f"Workflow: {workflow} - {workflow_config['description']}")

    # Calculate base processing costs
    chapters = stats['chapter_count']
    input_tokens = stats['estimated_tokens']

    # Analysis pass: send each chapter
    analysis_input = input_tokens * workflow_config['analysis_calls']
    analysis_output = chapters * 50  # ~50 tokens per analysis response (FILTER/CLEAN)

    # Cleaning pass: send chapters that need cleaning
    cleaning_input = input_tokens * workflow_config['cleaning_calls']
    cleaning_output = input_tokens * workflow_config['output_ratio']

    total_input = analysis_input + cleaning_input
    total_output = analysis_output + cleaning_output

    # Calculate costs
    base_input_cost = (total_input / 1_000_000) * pricing['input']
    base_output_cost = (total_output / 1_000_000) * pricing['output']
    base_cost = base_input_cost + base_output_cost

    cost_breakdown = {
        "model": pricing['name'],
        "workflow": workflow,
        "book_stats": {
            "chapters": chapters,
            "words": stats['total_words'],
            "tokens": input_tokens
        },
        "base_processing": {
            "input_tokens": int(total_input),
            "output_tokens": int(total_output),
            "cost": round(base_cost, 2)
        },
        "features": {},
        "total_cost": base_cost
    }

    if verbose:
        print(f"\n{'─'*40}")
        print(f"Base Processing Cost:")
        print(f"  Input tokens:  {int(total_input):,} (${base_input_cost:.2f})")
        print(f"  Output tokens: {int(total_output):,} (${base_output_cost:.2f})")
        print(f"  Subtotal: ${base_cost:.2f}")

    # Add optional features
    feature_cost = 0

    if with_profile:
        fc = FEATURE_COSTS["style_profile"]
        cost = ((fc['input_tokens'] / 1_000_000) * pricing['input'] +
                (fc['output_tokens'] / 1_000_000) * pricing['output'])
        feature_cost += cost
        cost_breakdown["features"]["style_profile"] = round(cost, 2)
        if verbose:
            print(f"\n  + Style Profile: ${cost:.2f}")

    if with_book_analysis:
        fc = FEATURE_COSTS["book_analysis"]
        cost = ((fc['input_tokens'] / 1_000_000) * pricing['input'] +
                (fc['output_tokens'] / 1_000_000) * pricing['output'])
        feature_cost += cost
        cost_breakdown["features"]["book_analysis"] = round(cost, 2)
        if verbose:
            print(f"  + Book Analysis: ${cost:.2f}")

    if with_change_plan:
        fc = FEATURE_COSTS["change_planning"]
        cost = ((fc['input_tokens'] / 1_000_000) * pricing['input'] +
                (fc['output_tokens'] / 1_000_000) * pricing['output'])
        feature_cost += cost
        cost_breakdown["features"]["change_planning"] = round(cost, 2)
        if verbose:
            print(f"  + Change Planning: ${cost:.2f}")

    if with_consistency_check:
        fc = FEATURE_COSTS["consistency_check"]
        cost = ((fc['input_tokens'] / 1_000_000) * pricing['input'] +
                (fc['output_tokens'] / 1_000_000) * pricing['output'])
        feature_cost += cost
        cost_breakdown["features"]["consistency_check"] = round(cost, 2)
        if verbose:
            print(f"  + Consistency Check: ${cost:.2f}")

    if with_drift_validation:
        fc = FEATURE_COSTS["drift_validation"]
        per_chapter_cost = ((fc['input_tokens'] / 1_000_000) * pricing['input'] +
                           (fc['output_tokens'] / 1_000_000) * pricing['output'])
        cost = per_chapter_cost * chapters
        feature_cost += cost
        cost_breakdown["features"]["drift_validation"] = round(cost, 2)
        if verbose:
            print(f"  + Drift Validation: ${cost:.2f} ({chapters} chapters)")

    total_cost = base_cost + feature_cost
    cost_breakdown["total_cost"] = round(total_cost, 2)

    # Add estimates for different models
    cost_breakdown["model_comparison"] = {}
    for model_key in ["haiku", "sonnet", "opus"]:
        mp = MODEL_PRICING[model_key]
        alt_input_cost = (total_input / 1_000_000) * mp['input']
        alt_output_cost = (total_output / 1_000_000) * mp['output']
        alt_total = alt_input_cost + alt_output_cost
        # Add feature costs scaled to this model
        if feature_cost > 0:
            scale = mp['input'] / pricing['input']
            alt_total += feature_cost * scale
        cost_breakdown["model_comparison"][model_key] = round(alt_total, 2)

    if verbose:
        print(f"\n{'─'*40}")
        print(f"ESTIMATED TOTAL: ${total_cost:.2f}")
        print(f"\nCost with different models:")
        print(f"  Haiku:  ${cost_breakdown['model_comparison']['haiku']:.2f}")
        print(f"  Sonnet: ${cost_breakdown['model_comparison']['sonnet']:.2f}")
        print(f"  Opus:   ${cost_breakdown['model_comparison']['opus']:.2f}")
        print(f"\n{'='*60}")
        print("NOTE: These are estimates. Actual costs may vary based on")
        print("content complexity and how many chapters need modifications.")
        print(f"{'='*60}\n")

    return cost_breakdown


def main():
    """CLI for cost estimation."""
    import argparse

    parser = argparse.ArgumentParser(
        description="Estimate API costs for EPUB processing"
    )
    parser.add_argument("input", help="Input EPUB file")
    parser.add_argument("--model", "-m", default="sonnet",
                       choices=["haiku", "sonnet", "opus"],
                       help="Model to use (default: sonnet)")
    parser.add_argument("--workflow", "-w", default="cleanup",
                       choices=list(WORKFLOW_MULTIPLIERS.keys()),
                       help="Workflow type (default: cleanup)")
    parser.add_argument("--with-profile", action="store_true",
                       help="Include style profiling cost")
    parser.add_argument("--with-analysis", action="store_true",
                       help="Include book analysis cost")
    parser.add_argument("--with-plan", action="store_true",
                       help="Include change planning cost")
    parser.add_argument("--with-consistency", action="store_true",
                       help="Include consistency checking cost")
    parser.add_argument("--with-drift", action="store_true",
                       help="Include drift validation cost")
    parser.add_argument("--all-features", action="store_true",
                       help="Include all optional features")
    parser.add_argument("--json", action="store_true",
                       help="Output as JSON")

    args = parser.parse_args()

    if args.all_features:
        args.with_profile = True
        args.with_analysis = True
        args.with_plan = True
        args.with_consistency = True
        args.with_drift = True

    result = estimate_cost(
        args.input,
        model=args.model,
        workflow=args.workflow,
        with_profile=args.with_profile,
        with_book_analysis=args.with_analysis,
        with_change_plan=args.with_plan,
        with_consistency_check=args.with_consistency,
        with_drift_validation=args.with_drift,
        verbose=not args.json
    )

    if args.json:
        import json
        print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()