utils.py

# Utility functions for BLEnD-VIS
"""
Utility functions for:
1. OpenAI Batch API operations
2. Async image generation
3. File handling and sanitization
"""

import asyncio
import base64
import json
import logging
import os
import random
import re
import time
from io import BytesIO
from pathlib import Path
from typing import Dict, List, Optional

import backoff
from dotenv import load_dotenv
from openai import AsyncOpenAI, OpenAI
from PIL import Image
from tqdm import tqdm
from tqdm.asyncio import tqdm as async_tqdm


# =============================================================================
# FILE UTILITIES
# =============================================================================

def sanitize_filename(text: str, max_length: int = 50) -> str:
    """
    Create a safe filename by removing illegal characters and truncating.
    
    Args:
        text: Input text to sanitize
        max_length: Maximum length of output
        
    Returns:
        Sanitized filename string
    """
    # Remove illegal characters
    text = re.sub(r'[<>:"/\\|?*\n\t]', '', text)
    # Replace whitespace with underscore
    text = re.sub(r'\s+', '_', text)
    # Collapse multiple underscores
    text = re.sub(r'__+', '_', text)
    # Remove leading/trailing underscores
    text = text.strip('_.-')
    return text[:max_length]


def load_json(filepath: Path) -> dict:
    """Load JSON file with error handling."""
    with open(filepath, 'r', encoding='utf-8') as f:
        return json.load(f)


def save_json(data: dict, filepath: Path, indent: int = 2):
    """Save data to JSON file."""
    filepath.parent.mkdir(parents=True, exist_ok=True)
    with open(filepath, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=indent, ensure_ascii=False)


# =============================================================================
# OPENAI BATCH API
# =============================================================================

def submit_and_retrieve_batch(
    tasks: List[dict], 
    batch_name: str, 
    client: OpenAI,
    output_dir: str = "batch_results",
    poll_interval: int = 30
) -> List[dict]:
    """
    Submit tasks to OpenAI Batch API, poll until completion, retrieve results.
    
    Args:
        tasks: List of task dictionaries for batch processing
        batch_name: Unique identifier for this batch
        client: OpenAI client instance
        output_dir: Directory to save batch files
        poll_interval: Seconds between status checks
        
    Returns:
        List of result dictionaries from the batch
    """
    if not tasks:
        print("No tasks provided for batch submission.")
        return []

    os.makedirs(output_dir, exist_ok=True)
    
    # Save tasks to JSONL
    jsonl_path = os.path.join(output_dir, f"{batch_name}_batch_tasks.jsonl")
    with open(jsonl_path, 'w', encoding='utf-8') as f:
        for task in tasks:
            f.write(json.dumps(task) + '\n')
    
    # Upload and create batch
    batch_file = client.files.create(
        file=open(jsonl_path, "rb"),
        purpose="batch"
    )
    
    batch_job = client.batches.create(
        input_file_id=batch_file.id,
        completion_window="24h",
        endpoint="/v1/chat/completions",
    )
    
    print(f"Submitted batch job: {batch_job.id}")
    
    # Poll for completion
    while True:
        batch_job = client.batches.retrieve(batch_job.id)
        if batch_job.status == "completed":
            print("Batch job completed.")
            break
        elif batch_job.status == "failed":
            print("Batch job failed.")
            return []
        else:
            if hasattr(batch_job, 'request_counts'):
                completed = batch_job.request_counts.completed
                total = batch_job.request_counts.total
                print(f"Progress: {completed}/{total} completed. Waiting...")
            else:
                print("Batch job in progress. Waiting...")
            time.sleep(poll_interval)
    
    # Retrieve results
    result_file_id = batch_job.output_file_id
    result = client.files.content(result_file_id).content
    
    result_path = os.path.join(output_dir, f"{batch_name}_batch_results.jsonl")
    with open(result_path, 'wb') as file:
        file.write(result)
    
    # Parse results
    results = []
    with open(result_path, 'r', encoding='utf-8') as f:
        for line in f:
            results.append(json.loads(line))
    
    return results


# =============================================================================
# ASYNC IMAGE GENERATION (OpenRouter)
# =============================================================================

def parse_base64_data_url(data_url: str) -> bytes:
    """Extract and decode Base64 data from a data URL."""
    try:
        _header, encoded = data_url.split(",", 1)
        return base64.b64decode(encoded)
    except (ValueError, IndexError, TypeError) as e:
        raise ValueError(f"Invalid Base64 data URL format: {e}")


@backoff.on_exception(
    backoff.expo,
    Exception,
    max_tries=5,
    on_giveup=lambda details: logging.error(
        f"API call for {details['args'][0].get('output_filename', 'unknown')} "
        f"failed after {details['tries']} tries. Error: {details['exception']}"
    )
)
async def get_image_completion_async(
    task: Dict, 
    model: str, 
    client: AsyncOpenAI
) -> Dict:
    """
    Make an async API call to OpenRouter for image generation.
    
    Args:
        task: Dictionary with 'prompt', 'output_path', 'output_filename'
        model: Model identifier
        client: AsyncOpenAI client
        
    Returns:
        Dictionary with status and output info
    """
    prompt_text = task['prompt']
    output_path = Path(task['output_path'])
    output_filename = task['output_filename']
    
    logging.info(f"Generating image: {output_filename}")
    
    completion = await client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt_text}],
        extra_body={"modalities": ["image", "text"]}
    )
    
    if not (completion.choices and completion.choices[0].message):
        raise ValueError("Invalid API response: No choices or message")
    
    message = completion.choices[0].message
    message_dict = message.model_dump()
    images_data = message_dict.get("images")
    
    if not images_data or not isinstance(images_data, list):
        raise ValueError("No valid image data in response")
    
    image_url = images_data[0].get("image_url", {}).get("url")
    if not image_url:
        raise ValueError("No image URL in response")
    
    image_bytes = parse_base64_data_url(image_url)
    image = Image.open(BytesIO(image_bytes))
    
    if image.mode in ('RGBA', 'P'):
        image = image.convert('RGB')
    
    output_path.parent.mkdir(parents=True, exist_ok=True)
    image.save(output_path, format='PNG')
    
    logging.info(f"SUCCESS: Saved {output_path}")
    return {
        'status': 'success', 
        'output_filename': output_filename, 
        'output_path': str(output_path)
    }


async def generate_images_asynchronously(
    tasks: List[Dict], 
    model: str, 
    client: AsyncOpenAI,
    concurrency_limit: int = 10
) -> List[Dict]:
    """
    Run image generation for all tasks asynchronously.
    
    Args:
        tasks: List of task dictionaries
        model: Model identifier
        client: AsyncOpenAI client
        concurrency_limit: Max concurrent requests
        
    Returns:
        List of result dictionaries
    """
    if not client:
        raise ValueError("AsyncOpenAI client not initialized")
    
    semaphore = asyncio.Semaphore(concurrency_limit)
    
    async def process_with_semaphore(task: Dict, pbar: async_tqdm) -> Dict:
        async with semaphore:
            try:
                await asyncio.sleep(random.uniform(0.1, 0.5))
                result = await get_image_completion_async(task, model, client)
                pbar.update(1)
                return result
            except Exception as e:
                logging.error(f"FAILED: {task['output_filename']} - {e}")
                pbar.update(1)
                return {
                    'status': 'failure', 
                    'output_filename': task['output_filename'], 
                    'error': str(e)
                }
    
    with async_tqdm(total=len(tasks), desc="Generating images") as pbar:
        async_tasks = [process_with_semaphore(task, pbar) for task in tasks]
        results = await asyncio.gather(*async_tasks)
    
    return results


# =============================================================================
# RESULT SCORING
# =============================================================================

def extract_answer_choice(output: str) -> Optional[str]:
    """
    Extract answer choice from model output.
    
    Args:
        output: Model output string
        
    Returns:
        Extracted choice letter or None
    """
    if not output:
        return None
    
    # Try JSON parsing first
    try:
        # Handle markdown code blocks
        cleaned = output.strip()
        if cleaned.startswith("```"):
            cleaned = re.sub(r'```\w*\n?', '', cleaned).strip()
        
        data = json.loads(cleaned)
        if isinstance(data, dict) and 'answer_choice' in data:
            return data['answer_choice'].strip().upper()
    except (json.JSONDecodeError, AttributeError):
        pass
    
    # Fallback: look for letter pattern
    match = re.search(r'\b([A-D])\b', output.upper())
    if match:
        return match.group(1)
    
    return None


def score_results(
    results: List[Dict], 
    mcq_data: List[Dict]
) -> List[Dict]:
    """
    Score model results against ground truth.
    
    Args:
        results: List of result dicts with 'MCQID' and 'output'
        mcq_data: List of MCQ dicts with 'MCQID' and 'answer_idx'
        
    Returns:
        List of scored results with 'correct' field added
    """
    # Build lookup
    answer_lookup = {mcq['MCQID']: mcq['answer_idx'] for mcq in mcq_data}
    
    scored = []
    for result in results:
        mcqid = result.get('MCQID')
        output = result.get('output', '')
        
        predicted = extract_answer_choice(output)
        correct_answer = answer_lookup.get(mcqid)
        
        scored.append({
            'MCQID': mcqid,
            'output': output,
            'predicted': predicted,
            'answer_idx': correct_answer,
            'correct': predicted == correct_answer if predicted and correct_answer else False
        })
    
    return scored