cli.py

#!/usr/bin/env python3
"""
CLI for the LLM-powered data extraction system.

Usage:
    python cli.py extract --input invoice.txt --type invoice --output result.json
    python cli.py extract --input email.txt --type email
    python cli.py validate --schema invoice --file data.json
"""

import json
import logging
import sys
from pathlib import Path
from typing import Optional

import typer
from rich.console import Console
from rich.panel import Panel
from rich.syntax import Syntax
from rich.table import Table
from dotenv import load_dotenv
import os

from src import LLMExtractor, ExtractionResult, setup_logging, EXTRACTION_SCHEMAS


# Load environment variables
load_dotenv()

# Initialize CLI
app = typer.Typer(
    name="llm-extractor",
    help="Robust LLM-powered data extraction with validation and retry logic",
    add_completion=False
)

console = Console()


@app.command()
def extract(
    input_file: Path = typer.Option(
        ..., "--input", "-i",
        help="Path to input text file",
        exists=True,
        dir_okay=False
    ),
    schema_type: str = typer.Option(
        ..., "--type", "-t",
        help="Type of data to extract (invoice, email, support_ticket)"
    ),
    output_file: Optional[Path] = typer.Option(
        None, "--output", "-o",
        help="Path to save extracted JSON (optional)"
    ),
    model: str = typer.Option(
        None, "--model", "-m",
        help="Ollama model to use (default: from .env or llama3.2)",
        envvar="MODEL_NAME"
    ),
    temperature: float = typer.Option(
        None, "--temperature",
        help="Temperature for LLM (default: from .env or 0.1)",
        envvar="TEMPERATURE"
    ),
    max_retries: int = typer.Option(
        None, "--max-retries",
        help="Maximum retry attempts (default: from .env or 3)",
        envvar="MAX_RETRIES"
    ),
    verbose: bool = typer.Option(
        False, "--verbose", "-v",
        help="Enable verbose debug logging"
    ),
    show_attempts: bool = typer.Option(
        True, "--show-attempts/--no-attempts",
        help="Show all extraction attempts"
    )
):
    """Extract structured data from unstructured text."""
    
    # Setup logging
    log_level = "DEBUG" if verbose else "INFO"
    setup_logging(log_level=log_level)
    logger = logging.getLogger(__name__)
    
    # Validate schema type
    if schema_type not in EXTRACTION_SCHEMAS:
        console.print(
            f"[red]Error: Unknown schema type '{schema_type}'[/red]",
            style="bold"
        )
        console.print(
            f"Valid types: {', '.join(EXTRACTION_SCHEMAS.keys())}"
        )
        raise typer.Exit(1)
    
    # Get configuration
    ollama_host = os.getenv("OLLAMA_HOST", "http://localhost:11434")
    
    # Read input text
    try:
        input_text = input_file.read_text(encoding='utf-8')
    except Exception as e:
        console.print(f"[red]Error reading input file: {e}[/red]")
        raise typer.Exit(1)
    
    console.print(Panel.fit(
        f"[bold cyan]LLM Data Extractor[/bold cyan]\n\n"
        f"Input: {input_file}\n"
        f"Schema: {schema_type}\n"
        f"Model: {model or os.getenv('MODEL_NAME', 'llama3.2')}",
        border_style="cyan"
    ))
    
    # Initialize extractor
    extractor = LLMExtractor(
        model=model or os.getenv("MODEL_NAME", "llama3.2"),
        ollama_host=ollama_host,
        temperature=float(temperature or os.getenv("TEMPERATURE", 0.1)),
        max_retries=int(max_retries or os.getenv("MAX_RETRIES", 3))
    )
    
    # Extract data
    console.print("\n[yellow]⚙ Starting extraction...[/yellow]\n")
    
    try:
        result: ExtractionResult = extractor.extract(
            text=input_text,
            schema_type=schema_type
        )
    except Exception as e:
        logger.exception("Extraction failed with exception")
        console.print(f"\n[red]✗ Extraction failed: {e}[/red]\n")
        raise typer.Exit(1)
    
    # Display results
    _display_result(result, show_attempts)
    
    # Save output if requested
    if result.success and output_file:
        _save_output(result, output_file)
    
    # Exit with appropriate code
    sys.exit(0 if result.success else 1)


@app.command()
def validate(
    schema_type: str = typer.Option(
        ..., "--schema", "-s",
        help="Schema type to validate against (invoice, email, support_ticket)"
    ),
    input_file: Path = typer.Option(
        ..., "--file", "-f",
        help="JSON file to validate",
        exists=True,
        dir_okay=False
    )
):
    """Validate a JSON file against a schema."""
    
    # Validate schema type
    if schema_type not in EXTRACTION_SCHEMAS:
        console.print(
            f"[red]Error: Unknown schema type '{schema_type}'[/red]",
            style="bold"
        )
        console.print(
            f"Valid types: {', '.join(EXTRACTION_SCHEMAS.keys())}"
        )
        raise typer.Exit(1)
    
    # Load JSON
    try:
        data = json.loads(input_file.read_text())
    except json.JSONDecodeError as e:
        console.print(f"[red]Error: Invalid JSON: {e}[/red]")
        raise typer.Exit(1)
    
    # Validate
    schema_model = EXTRACTION_SCHEMAS[schema_type]["model"]
    
    try:
        validated = schema_model(**data)
        console.print(f"\n[green]✓ Validation passed![/green]\n")
        console.print(Panel(
            Syntax(
                validated.model_dump_json(indent=2),
                "json",
                theme="monokai"
            ),
            title="Validated Data",
            border_style="green"
        ))
        sys.exit(0)
        
    except Exception as e:
        console.print(f"\n[red]✗ Validation failed![/red]\n")
        console.print(f"[red]{str(e)}[/red]\n")
        sys.exit(1)


@app.command()
def list_schemas():
    """List all available extraction schemas."""
    
    table = Table(title="Available Extraction Schemas", show_header=True)
    table.add_column("Schema Type", style="cyan", no_wrap=True)
    table.add_column("Description", style="white")
    table.add_column("Model Class", style="yellow")
    
    for schema_type, info in EXTRACTION_SCHEMAS.items():
        table.add_row(
            schema_type,
            info["description"],
            info["model"].__name__
        )
    
    console.print("\n")
    console.print(table)
    console.print("\n")


def _display_result(result: ExtractionResult, show_attempts: bool):
    """Display extraction results with rich formatting."""
    
    if show_attempts:
        # Show attempt table
        table = Table(title="Extraction Attempts", show_header=True)
        table.add_column("#", style="cyan", no_wrap=True)
        table.add_column("Status", style="white")
        table.add_column("Details", style="white")
        
        for attempt in result.attempts:
            status = "✓ Success" if attempt.success else "✗ Failed"
            status_style = "green" if attempt.success else "red"
            
            details = ""
            if attempt.validation_errors:
                details = "; ".join(attempt.validation_errors[:2])
                if len(attempt.validation_errors) > 2:
                    details += "..."
            elif attempt.success:
                details = "Valid data extracted"
            
            table.add_row(
                str(attempt.attempt_number),
                f"[{status_style}]{status}[/{status_style}]",
                details
            )
        
        console.print("\n")
        console.print(table)
        console.print("\n")
    
    # Show final result
    if result.success:
        console.print(
            f"[green bold]✓ Extraction succeeded after {result.total_attempts} attempt(s)![/green bold]\n"
        )
        
        # Display extracted data
        console.print(Panel(
            Syntax(
                result.data.model_dump_json(indent=2),
                "json",
                theme="monokai",
                line_numbers=True
            ),
            title="[green]Extracted Data[/green]",
            border_style="green"
        ))
    else:
        console.print(
            f"[red bold]✗ Extraction failed after {result.total_attempts} attempt(s)[/red bold]\n"
        )
        console.print(f"[red]Error: {result.error_message}[/red]\n")


def _save_output(result: ExtractionResult, output_file: Path):
    """Save extraction result to JSON file."""
    try:
        output_file.parent.mkdir(parents=True, exist_ok=True)
        
        output_data = {
            "success": result.success,
            "total_attempts": result.total_attempts,
            "data": result.data.model_dump() if result.data else None,
            "timestamp": result.attempts[-1].timestamp if result.attempts else None
        }
        
        output_file.write_text(
            json.dumps(output_data, indent=2),
            encoding='utf-8'
        )
        
        console.print(f"\n[green]✓ Saved to {output_file}[/green]\n")
        
    except Exception as e:
        console.print(f"\n[red]✗ Failed to save output: {e}[/red]\n")


if __name__ == "__main__":
    app()