audiobook/backend/test_ai_segmentation.py

#!/usr/bin/env python3
"""
Test script for AI-powered text segmentation.
Tests the integration between OllamaService and TTSService for intelligent text splitting.
"""
import asyncio
import sys
from services.ollama_service import OllamaService
from services.tts_service import TTSService


async def test_ai_segmentation():
    """Test AI segmentation with a sample narrative text."""
    print("Testing AI-Powered Text Segmentation\n" + "="*50)
    
    # Sample text without explicit role markers
    text = """The old wizard gazed into his crystal ball. "I see danger approaching," he whispered.
A young knight heard the warning and drew his sword. "I will protect the kingdom," he declared boldly.
The princess watched from her tower window. She knew the prophecy would come true."""
    
    print(f"\nInput text:\n{text}\n")
    
    # Initialize services
    ollama_service = OllamaService()
    tts_service = TTSService(ollama_service=ollama_service)
    
    # Check if Ollama is available
    is_available = await ollama_service.is_available()
    if not is_available:
        print("❌ Ollama is not available. Please start Ollama service.")
        return False
    
    print("✓ Ollama service is available\n")
    
    # Define role-voice mapping (simulating what user would assign)
    role_voice_map = {
        "Narrator": "af_bella",
        "The old wizard": "am_adam",
        "A young knight": "am_michael",
        "The princess": "af_sarah"
    }
    
    print(f"Role-voice mapping: {role_voice_map}\n")
    
    # Test AI segmentation
    try:
        print("Requesting AI segmentation...")
        segments = await tts_service.segment_text_with_ai(text, role_voice_map)
        
        if segments is None:
            print("❌ AI segmentation returned None")
            return False
        
        print(f"\n✓ AI segmentation successful! Generated {len(segments)} segments:\n")
        
        for i, (speaker, voice, content) in enumerate(segments, 1):
            print(f"Segment {i}:")
            print(f"  Speaker: {speaker}")
            print(f"  Voice: {voice}")
            print(f"  Text: {content[:80]}{'...' if len(content) > 80 else ''}")
            print()
        
        # Verify segments are properly assigned
        speakers_used = set(seg[0] for seg in segments)
        print(f"Speakers used in segmentation: {speakers_used}")
        
        # Check if AI intelligently assigned different speakers
        if len(speakers_used) > 1:
            print("\n✓ SUCCESS: AI intelligently assigned multiple speakers!")
            return True
        else:
            print("\n⚠ WARNING: Only one speaker used. AI may not have detected dialogue.")
            return True  # Still a success, just not ideal
            
    except Exception as e:
        print(f"\n❌ ERROR: {e}")
        import traceback
        traceback.print_exc()
        return False


async def test_fallback_to_pattern():
    """Test that pattern-based segmentation works as fallback."""
    print("\n" + "="*50)
    print("Testing Fallback to Pattern-Based Segmentation\n")
    
    # Text with explicit markers
    text = """Alice: Hello there!
Bob: Hi Alice, how are you?
[Narrator] They continued their conversation in the garden."""
    
    print(f"Input text:\n{text}\n")
    
    tts_service = TTSService()
    
    role_voice_map = {
        "Alice": "af_bella",
        "Bob": "am_adam",
        "Narrator": "af_sarah"
    }
    
    print(f"Role-voice mapping: {role_voice_map}\n")
    
    # Use pattern-based segmentation
    segments = tts_service.segment_text(text, role_voice_map)
    
    print(f"✓ Pattern segmentation generated {len(segments)} segments:\n")
    
    for i, (speaker, voice, content) in enumerate(segments, 1):
        print(f"Segment {i}: [{speaker}] {content}")
    
    # Verify all speakers were detected
    speakers_used = set(seg[0] for seg in segments)
    expected_speakers = {"Alice", "Bob", "Narrator"}
    
    if speakers_used == expected_speakers:
        print(f"\n✓ SUCCESS: All expected speakers detected: {speakers_used}")
        return True
    else:
        print(f"\n❌ FAIL: Expected {expected_speakers}, got {speakers_used}")
        return False


async def main():
    """Run all tests."""
    print("\n" + "="*70)
    print("AI-POWERED TEXT SEGMENTATION TEST SUITE")
    print("="*70 + "\n")
    
    # Test 1: AI segmentation
    test1_passed = await test_ai_segmentation()
    
    # Test 2: Pattern fallback
    test2_passed = await test_fallback_to_pattern()
    
    # Summary
    print("\n" + "="*70)
    print("TEST SUMMARY")
    print("="*70)
    print(f"AI Segmentation: {'✓ PASSED' if test1_passed else '❌ FAILED'}")
    print(f"Pattern Fallback: {'✓ PASSED' if test2_passed else '❌ FAILED'}")
    print("="*70 + "\n")
    
    return test1_passed and test2_passed


if __name__ == '__main__':
    success = asyncio.run(main())
    sys.exit(0 if success else 1)