test_extraction.py

#!/usr/bin/env python3
"""Test script for two-stage fact extraction."""

import json
from pathlib import Path

from src.parsers.csv_parser import parse_iconik_csv
from src.extractors.statement_extractor import StatementExtractor

CSV_PATH = Path("assets/Trump Speech Davos Transcription.csv")


def main():
    print(f"Parsing {CSV_PATH}...")
    segments = parse_iconik_csv(CSV_PATH)
    print(f"Loaded {len(segments)} segments\n")

    extractor = StatementExtractor(use_two_stage=True)

    # Show token/cost estimates
    print("=" * 60)
    print("TOKEN & COST ESTIMATES")
    print("=" * 60)
    estimates = extractor.estimate_tokens(segments)
    print(f"Total segments: {estimates['total_segments']}")
    print(f"Avg tokens/segment: {estimates['avg_tokens_per_segment']}")
    print()
    print("Single-stage (Sonnet only):")
    print(f"  API calls: {estimates['single_stage']['api_calls']}")
    print(f"  Input tokens: {estimates['single_stage']['input_tokens']:,}")
    print(f"  Output tokens: {estimates['single_stage']['output_tokens']:,}")
    print(f"  Estimated cost: ${estimates['single_stage']['estimated_cost']:.4f}")
    print()
    print("Two-stage (Haiku + Sonnet):")
    print(f"  Filter API calls: {estimates['two_stage']['filter_api_calls']}")
    print(f"  Expected candidates: {estimates['two_stage']['expected_candidates']}")
    print(f"  Analysis API calls: {estimates['two_stage']['analysis_api_calls']}")
    print(f"  Total input tokens: {estimates['two_stage']['total_input_tokens']:,}")
    print(f"  Total output tokens: {estimates['two_stage']['total_output_tokens']:,}")
    print(f"  Estimated cost: ${estimates['two_stage']['estimated_cost']:.4f}")
    print()
    print(f"ESTIMATED SAVINGS: {estimates['savings_percent']}%")
    print("=" * 60)
    print()

    # Run extraction
    print("Running two-stage extraction...")
    print("  Stage 1: Filtering with Haiku...")
    result = extractor.extract(segments, source_file=CSV_PATH.name)

    print("  Stage 2: Analysis complete!")
    print()
    print("=" * 60)
    print("RESULTS")
    print("=" * 60)
    print(f"Statements found: {len(result.statements)}")
    print()

    # Show first 10 statements
    for i, stmt in enumerate(result.statements[:10], 1):
        print(f"{i}. [{stmt.start_timecode} - {stmt.end_timecode}] ({stmt.category})")
        print(f"   \"{stmt.statement}\"")
        print(f"   Confidence: {stmt.confidence}")
        print()

    if len(result.statements) > 10:
        print(f"... and {len(result.statements) - 10} more statements")

    # Save full results to JSON
    output_path = Path("extraction_results.json")
    with open(output_path, "w") as f:
        json.dump(result.model_dump(mode="json"), f, indent=2, default=str)
    print(f"\nFull results saved to {output_path}")


if __name__ == "__main__":
    main()