Genarrative/server-rs/crates/api-server/src/editor_agent/pricing.rs

use std::collections::HashMap;
use platform_editor_agent::framework::tool::Tool;
use platform_image::GPT_IMAGE_2_MODEL;
use platform_editor_agent::agent::tools::context::EditorToolContext;
#[cfg(test)]
use platform_editor_agent::agent::asset::ImageId;
use platform_editor_agent::agent::tools::edit_image::{EditImageTool, EditImageToolArgs};
use platform_editor_agent::agent::tools::generate_background_music::{
  GenerateBackgroundMusicTool, GenerateBackgroundMusicToolArgs,
};
use platform_editor_agent::agent::tools::generate_character::GenerateCharacterTool;
use platform_editor_agent::agent::tools::generate_icon_spritesheet::{
  GenerateIconSpritesheetTool, GenerateIconSpritesheetToolArgs,
};
use platform_editor_agent::agent::tools::generate_image::{GenerateImageTool, GenerateImageToolArgs};
use platform_editor_agent::agent::tools::generate_sound_effect::{
  GenerateSoundEffectTool, GenerateSoundEffectToolArgs,
};
use platform_editor_agent::agent::tools::generate_ui_design::GenerateUiDesignTool;
use platform_editor_agent::agent::tools::generate_video::{GenerateVideoTool, GenerateVideoToolArgs};
use crate::editor_generation_config::{load_editor_generation_pricing_from_paths, EditorGenerationPricingConfig};

#[test]
fn every_editor_agent_tool_exposes_argument_based_pricing() {
    let pricing = load_editor_generation_pricing_from_paths(None)
        .expect("default editor pricing should load");
    let context = context();

    assert_eq!(
        EditImageTool {
            context: context.clone(),
        }
        .pricing(
            &pricing,
            &EditImageToolArgs {
                object_image_id: ImageId {
                    id: "image-1".to_string(),
                },
                reference_image_ids: Vec::new(),
                prompt: "改成蓝色".to_string(),
            },
        ),
        3
    );
    assert_eq!(
        GenerateImageTool {
            context: context.clone(),
        }
        .pricing(&pricing, &image_args(Some("2K"))),
        5
    );
    assert_eq!(
        GenerateCharacterTool {
            context: context.clone(),
        }
        .pricing(&pricing, &image_args(None)),
        3
    );
    assert_eq!(
        GenerateUiDesignTool {
            context: context.clone(),
        }
        .pricing(&pricing, &image_args(Some("2K"))),
        5
    );
    assert_eq!(
        GenerateIconSpritesheetTool {
            context: context.clone(),
        }
        .pricing(
            &pricing,
            &GenerateIconSpritesheetToolArgs {
                reference_image_id: ImageId {
                    id: "image-1".to_string(),
                },
                reference_image_ids: Vec::new(),
                icon_descriptions: vec!["背包".to_string(), "地图".to_string()],
                aspect_ratio: Some("1:1".to_string()),
                image_size: Some("2K".to_string()),
            },
        ),
        5
    );
    assert_eq!(
        GenerateVideoTool {
            context: context.clone(),
        }
        .pricing(
            &pricing,
            &GenerateVideoToolArgs {
                prompt: "镜头缓慢推进".to_string(),
                reference_image_ids: Vec::new(),
                aspect_ratio: None,
                duration_seconds: Some(6),
                model: Some("seedance2.0".to_string()),
                resolution: Some("720p".to_string()),
                sound: None,
            },
        ),
        144
    );
    assert_eq!(
        GenerateSoundEffectTool.pricing(
            &pricing,
            &GenerateSoundEffectToolArgs {
                prompt: "按钮点击声".to_string(),
                duration: None,
                model: None,
            },
        ),
        5
    );
    assert_eq!(
        GenerateBackgroundMusicTool.pricing(
            &pricing,
            &GenerateBackgroundMusicToolArgs {
                prompt: "轻松背景音乐".to_string(),
                make_instrumental: true,
            },
        ),
        12
    );
}

#[test]
fn pricing_uses_the_supplied_runtime_snapshot() {
    let mut pricing = load_editor_generation_pricing_from_paths(None)
        .expect("default editor pricing should load");
    pricing
        .models
        .get_mut("gpt-image-2")
        .expect("gpt image pricing should exist")
        .prices
        .insert("2K".to_string(), 37);
    pricing
        .models
        .get_mut("seedance2.0-fast")
        .expect("video pricing should exist")
        .prices
        .insert("720p".to_string(), 7);
    pricing
        .models
        .get_mut("audio1.0")
        .expect("sound pricing should exist")
        .price = Some(19);

    assert_eq!(
        GenerateImageTool { context: context() }.pricing(&pricing, &image_args(Some("2K"))),
        37
    );
    assert_eq!(
        GenerateVideoTool { context: context() }.pricing(
            &pricing,
            &GenerateVideoToolArgs {
                prompt: "four seconds".to_string(),
                reference_image_ids: Vec::new(),
                aspect_ratio: None,
                duration_seconds: None,
                model: None,
                resolution: None,
                sound: None,
            },
        ),
        28
    );
    assert_eq!(
        GenerateSoundEffectTool.pricing(
            &pricing,
            &GenerateSoundEffectToolArgs {
                prompt: "sound".to_string(),
                duration: None,
                model: None,
            },
        ),
        19
    );
}

impl EditorAgentPricedTool for GenerateVideoTool {
    fn pricing(
        &self,
        pricing: &EditorGenerationPricingConfig,
        args: &GenerateVideoToolArgs,
    ) -> u32 {
        let model = args
            .model
            .as_deref()
            .map(str::trim)
            .filter(|value| !value.is_empty())
            .unwrap_or(Self::DEFAULT_VIDEO_MODEL);
        let resolution = args
            .resolution
            .as_deref()
            .map(str::trim)
            .filter(|value| !value.is_empty())
            .unwrap_or(Self::DEFAULT_VIDEO_RESOLUTION);
        let duration_seconds = args
            .duration_seconds
            .unwrap_or(Self::DEFAULT_VIDEO_DURATION_SECONDS);
        pricing.video_model_mud_points(Some(model), resolution, duration_seconds)
    }
}

impl EditorAgentPricedTool for GenerateUiDesignTool {
    fn pricing(
        &self,
        pricing: &EditorGenerationPricingConfig,
        args: &GenerateImageToolArgs,
    ) -> u32 {
        editor_agent_image_mud_points(pricing, Some("ui-design"), args.image_size.as_deref())
    }
}

impl EditorAgentPricedTool for GenerateSoundEffectTool {
    fn pricing(
        &self,
        pricing: &EditorGenerationPricingConfig,
        args: &GenerateSoundEffectToolArgs,
    ) -> u32 {
        pricing.sound_effect_model_mud_points(None)
    }
}

impl EditorAgentPricedTool for GenerateImageTool {
    fn pricing(
        &self,
        pricing: &EditorGenerationPricingConfig,
        args: &GenerateImageToolArgs,
    ) -> u32 {
        editor_agent_image_mud_points(pricing, None, args.image_size.as_deref())
    }
}

impl EditorAgentPricedTool for GenerateIconSpritesheetTool {
    fn pricing(
        &self,
        pricing: &EditorGenerationPricingConfig,
        args: &GenerateIconSpritesheetToolArgs,
    ) -> u32 {
        editor_agent_image_mud_points(pricing, Some("icon"), args.image_size.as_deref())
    }
}

impl EditorAgentPricedTool for GenerateCharacterTool {
    fn pricing(
        &self,
        pricing: &EditorGenerationPricingConfig,
        args: &GenerateImageToolArgs,
    ) -> u32 {
        editor_agent_image_mud_points(pricing, Some("character"), args.image_size.as_deref())
    }
}

/// api-server 侧的画布 Agent 工具计价扩展。
///
/// 通用 `Tool` 仍只负责参数校验；价格依赖 api-server 的运行时配置，不能下沉到
/// `module-editor-agent`。实际执行和扣费仍由既有生成 BFF 负责。
pub(crate) trait EditorAgentPricedTool: Tool {
    fn pricing(&self, pricing: &EditorGenerationPricingConfig, args: &<Self as Tool>::Args) -> u32;
}

pub(crate) fn editor_agent_image_mud_points(
    pricing: &EditorGenerationPricingConfig,
    kind: Option<&str>,
    image_size: Option<&str>,
) -> u32 {
    // 这些 Agent 工具当前向既有 BFF 传 model=None；BFF 会先归一为 gpt-image-2。
    // 尺寸同样只把精确的 2K 识别为 2K，其余值回落到 1K。
    let normalized_image_size = match image_size.map(str::trim) {
        Some("2K") => "2K",
        _ => "1K",
    };
    pricing.image_generation_mud_points(kind, Some(GPT_IMAGE_2_MODEL), Some(normalized_image_size))
}

impl EditorAgentPricedTool for GenerateBackgroundMusicTool {
    fn pricing(
        &self,
        pricing: &EditorGenerationPricingConfig,
        _args: &GenerateBackgroundMusicToolArgs,
    ) -> u32 {
        pricing.background_music_model_mud_points(None)
    }
}

impl EditorAgentPricedTool for EditImageTool {
    fn pricing(&self, pricing: &EditorGenerationPricingConfig, _args: &EditImageToolArgs) -> u32 {
        editor_agent_image_mud_points(pricing, Some("quick-edit"), Some("1K"))
    }
}