trax/examples/basic_transcription.py

#!/usr/bin/env python3
"""Example: Basic transcription with retry logic using AI Assistant Library.

This example demonstrates:
- Using the TranscriptionService with automatic retry
- Leveraging the AI Assistant Library's retry decorators
- Error handling and logging
"""

import asyncio
import logging
from pathlib import Path
from typing import Dict, Any

# Add parent directory to path for imports
import sys
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.base.services import TranscriptionService, async_retry, RetryConfig

class RetryableError(Exception):
    """Error that can be retried."""
    def __init__(self, message, retry_after=None):
        super().__init__(message)
        self.retry_after = retry_after

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class WhisperTranscriptionService(TranscriptionService):
    """Example Whisper transcription service implementation."""
    
    async def _transcribe_impl(self, audio_path: Path) -> Dict[str, Any]:
        """Simulate Whisper transcription.
        
        In a real implementation, this would:
        1. Load the Whisper model
        2. Process the audio file
        3. Return structured transcript data
        """
        logger.info(f"Transcribing {audio_path} with Whisper...")
        
        # Simulate processing time
        await asyncio.sleep(2)
        
        # Simulate occasional failures for retry demonstration
        import random
        if random.random() < 0.3:  # 30% chance of failure
            raise RetryableError("Temporary transcription failure", retry_after=1)
        
        # Return mock transcript data
        return {
            "text": f"This is a sample transcript for {audio_path.name}",
            "segments": [
                {
                    "id": 0,
                    "start": 0.0,
                    "end": 5.0,
                    "text": "This is a sample transcript",
                    "confidence": 0.95
                },
                {
                    "id": 1,
                    "start": 5.0,
                    "end": 10.0,
                    "text": f"for {audio_path.name}",
                    "confidence": 0.92
                }
            ],
            "language": "en",
            "duration": 10.0,
            "model": "whisper-large-v3",
        }


# Custom retry configuration for critical operations
critical_retry_config = RetryConfig(
    max_attempts=5,
    initial_delay=1.0,
    backoff_factor=2.0,
    max_delay=30.0,
    jitter=True,
)


@async_retry(
    max_attempts=critical_retry_config.max_attempts,
    backoff_factor=critical_retry_config.backoff_factor
)
async def transcribe_with_custom_retry(service: TranscriptionService, audio_path: Path) -> Dict[str, Any]:
    """Transcribe with custom retry configuration.
    
    This demonstrates using the library's retry decorator with custom settings
    for critical transcription operations.
    """
    logger.info(f"Attempting transcription with custom retry logic...")
    result = await service.transcribe(audio_path)
    logger.info(f"Transcription successful!")
    return result


async def main():
    """Run transcription examples."""
    
    # Example audio file paths
    audio_files = [
        Path("example_audio.mp3"),
        Path("example_video.mp4"),
        Path("example_podcast.wav"),
    ]
    
    # Initialize service with configuration
    config = {
        "pipeline_version": "v1",
        "max_retries": 3,
        "supported_formats": [".mp3", ".mp4", ".wav", ".flac"],
    }
    
    service = WhisperTranscriptionService(config=config)
    
    # Initialize the service
    await service.initialize()
    
    try:
        # Example 1: Basic transcription with built-in retry
        print("\n" + "="*60)
        print("Example 1: Basic Transcription with Built-in Retry")
        print("="*60)
        
        for audio_path in audio_files[:1]:  # Process first file
            try:
                # Create a dummy file for demonstration
                audio_path.touch()
                
                if service.can_handle(audio_path):
                    result = await service.transcribe(audio_path)
                    print(f"✓ Transcribed {audio_path.name}:")
                    print(f"  Text: {result['text'][:50]}...")
                    print(f"  Duration: {result['duration']}s")
                    print(f"  Segments: {len(result['segments'])}")
                else:
                    print(f"✗ Cannot handle {audio_path.suffix} files")
                    
            except Exception as e:
                print(f"✗ Failed to transcribe {audio_path.name}: {e}")
            finally:
                # Clean up dummy file
                if audio_path.exists():
                    audio_path.unlink()
        
        # Example 2: Custom retry configuration
        print("\n" + "="*60)
        print("Example 2: Transcription with Custom Retry Logic")
        print("="*60)
        
        test_audio = Path("critical_audio.mp3")
        test_audio.touch()
        
        try:
            result = await transcribe_with_custom_retry(service, test_audio)
            print(f"✓ Successfully transcribed with custom retry")
            print(f"  Result: {result['text']}")
        except Exception as e:
            print(f"✗ Failed after {critical_retry_config.max_attempts} attempts: {e}")
        finally:
            if test_audio.exists():
                test_audio.unlink()
        
        # Example 3: Service health and capabilities
        print("\n" + "="*60)
        print("Example 3: Service Health and Capabilities")
        print("="*60)
        
        # Check service health
        health = service.get_health_status()
        print(f"Service Health:")
        print(f"  Status: {health['status']}")
        print(f"  Healthy: {health['is_healthy']}")
        print(f"  Uptime: {health.get('uptime_seconds', 0):.1f}s")
        
        # Get pipeline information
        pipeline_info = service.get_pipeline_info()
        print(f"\nPipeline Information:")
        print(f"  Version: {pipeline_info['version']}")
        print(f"  Capabilities: {', '.join(pipeline_info['capabilities'])}")
        
    finally:
        # Cleanup
        await service.shutdown()
        print("\n✓ Service shutdown complete")


if __name__ == "__main__":
    print("Trax Transcription Example")
    print("Using AI Assistant Library for retry and error handling")
    print("-" * 60)
    
    # Run the async main function
    asyncio.run(main())