""" 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 {}