test_openrouter.py

"""
Example script demonstrating OpenRouter client usage
for medical transcription correction
"""

import os
import sys
from pathlib import Path

# Add project root to path
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from corrector.openrouter_client import OpenRouterClient
from dotenv import load_dotenv

# Load environment variables
load_dotenv()


def test_basic_completion():
    """Test basic chat completion"""
    print("=" * 60)
    print("Testing Basic Chat Completion")
    print("=" * 60)
    
    client = OpenRouterClient()
    
    messages = [
        {
            "role": "user",
            "content": "How many r's are in the word 'strawberry'?"
        }
    ]
    
    response = client.chat_completion(messages=messages, reasoning_enabled=True)
    
    print(f"\nModel: {client.model}")
    print(f"Response: {client._extract_content(response)}")
    print(f"\nFull response:\n{response}")


def test_medical_correction():
    """Test medical transcription correction"""
    print("\n" + "=" * 60)
    print("Testing Medical Transcription Correction")
    print("=" * 60)
    
    client = OpenRouterClient()
    
    # Sample transcription with potential errors
    transcription = "Пациент жалуется на боль в животе, тошноту и рвоту. Диагноз апендицит."
    
    system_prompt = """Ты медицинский помощник. Исправь ошибки в медицинской транскрипции.
Используй правильную медицинскую терминологию.
Верни только исправленный текст без дополнительных пояснений."""
    
    corrected_text = client.correct_text(
        text=transcription,
        system_prompt=system_prompt,
        temperature=0.1
    )
    
    print(f"\nOriginal: {transcription}")
    print(f"Corrected: {corrected_text}")


def test_with_medical_terms():
    """Test correction with medical terms context"""
    print("\n" + "=" * 60)
    print("Testing with Medical Terms Context")
    print("=" * 60)
    
    client = OpenRouterClient()
    
    # Load medical terms
    medical_terms = "аппендицит, гастрит, энцефалопатия, кардиомиопатия"
    
    transcription = "У пациента подозрение на апендицит и гастрит"
    
    system_prompt = f"""Ты медицинский помощник. Исправь ошибки в транскрипции.
    
Используй следующие медицинские термины: {medical_terms}

Верни только исправленный текст."""
    
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": transcription}
    ]
    
    response = client.chat_completion(
        messages=messages,
        temperature=0.1,
        reasoning_enabled=True
    )
    
    corrected_text = client._extract_content(response)
    
    print(f"\nMedical terms: {medical_terms}")
    print(f"Original: {transcription}")
    print(f"Corrected: {corrected_text}")


def test_model_info():
    """Test model info retrieval"""
    print("\n" + "=" * 60)
    print("Model Information")
    print("=" * 60)
    
    client = OpenRouterClient()
    info = client.get_model_info()
    
    print("\nConfiguration:")
    for key, value in info.items():
        print(f"  {key}: {value}")


if __name__ == "__main__":
    try:
        # Run tests
        test_basic_completion()
        test_medical_correction()
        test_with_medical_terms()
        test_model_info()
        
        print("\n" + "=" * 60)
        print("All tests completed successfully!")
        print("=" * 60)
        
    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)