pipeline-app/app/model_selector.py

"""
Model selection interactive handler - allows customer to choose video generation model.

Creates a human_task after storyboard generation to:
1. Display cost estimates for all available models
2. Wait for customer selection
3. Store selected model in task context for downstream steps
"""

import json
import logging
from typing import Dict, Any

from .cost_estimator import estimate_all_models, format_cost_summary

logger = logging.getLogger("pipeline.model_selector")


async def handle_model_selecting(
    tenant_id: str,
    task_id: str,
    step_name: str,
    input_data: Dict[str, Any],
    config: Dict[str, Any]
) -> Dict[str, Any]:
    """
    Interactive handler: present model options and wait for customer selection.
    
    Args:
        tenant_id: Tenant identifier
        task_id: Task identifier
        step_name: Step name (should be "model_selecting")
        input_data: Input data containing storyboard
        config: Step configuration
    
    Returns:
        Dict with selected_model and cost_estimates
    """
    # Extract storyboard from input
    storyboard = None
    for dep_name, dep_output in input_data.items():
        if isinstance(dep_output, dict):
            if dep_output.get("storyboard"):
                storyboard = dep_output["storyboard"]
                break
    
    if not storyboard:
        raise ValueError("上游步骤未提供分镜脚本 (storyboard)")
    
    # Generate cost estimates
    fps = config.get("fps", 24)
    estimates = estimate_all_models(storyboard, fps)
    
    # Format summary
    summary = format_cost_summary(estimates)
    logger.info(f"Generated cost estimates for {len(estimates)} models")
    
    # Create human_task for customer selection
    task_data = {
        "title": "请选择视频生成模型",
        "description": summary,
        "options": [
            {
                "value": est["model_id"],
                "label": est["model_name"],
                "description": f"${est['total_cost']:.4f} - {est['description']}",
                "recommended": est.get("recommended", False),
            }
            for est in estimates
        ],
        "default_value": estimates[0]["model_id"] if estimates else "wan2.2",
        "task_type": "model_selection",
        "metadata": {
            "cost_estimates": estimates,
            "scene_count": len(storyboard),
            "total_frames": sum(
                int((s.get("end_time", 10) - s.get("start_time", 0)) * fps)
                for s in storyboard
            ),
        },
    }
    
    from pipeline_service.storage import create_human_task
    human_task_id = await create_human_task(
        task_id=task_id,
        step_name=step_name,
        version=config.get("version", 1),
        task_type="model_selection",
        form_schema=task_data,
        timeout_hours=config.get("timeout_hours", 24),
    )
    
    logger.info(f"Created model selection human_task: {human_task_id}")
    
    # Return placeholder (actual selection will be in human_task.result)
    return {
        "human_task_id": human_task_id,
        "status": "waiting_for_selection",
        "cost_estimates": estimates,
        "summary": summary,
    }


def get_selected_model(human_task_result: Dict[str, Any]) -> str:
    """
    Extract selected model from human_task result.
    
    Args:
        human_task_result: Result dict from completed human_task
    
    Returns:
        Selected model_id
    """
    if not human_task_result:
        logger.warning("No human_task result, using default model")
        return "wan2.2"
    
    selected = human_task_result.get("selected_model")
    if not selected:
        logger.warning("No model selected, using default")
        return "wan2.2"
    
    logger.info(f"Customer selected model: {selected}")
    return selected


def get_model_cost_for_selected(
    estimates: list,
    selected_model: str
) -> Dict[str, Any]:
    """
    Get cost estimate for the selected model.
    
    Args:
        estimates: List of cost estimate dicts
        selected_model: Selected model_id
    
    Returns:
        Cost estimate dict for selected model
    """
    for est in estimates:
        if est["model_id"] == selected_model:
            return est
    
    logger.warning(f"Selected model {selected_model} not found in estimates")
    return {}