From 10d7f09cbf126cbbba6440ed7f38eedf70ff457a Mon Sep 17 00:00:00 2001 From: kvtodev Date: Fri, 10 Jul 2026 16:13:40 +0800 Subject: [PATCH] impl image gen tool --- .../crates/api-server/src/editor_agent/api.rs | 76 ++++- .../editor_tools/generate_background_music.rs | 0 .../editor_tools/generate_character.rs | 0 .../editor_tools/generate_icon_spritesheet.rs | 0 .../editor_tools/generate_image.rs | 290 ++++++++++++++++++ .../editor_tools/generate_sound_effect.rs | 0 .../editor_tools/generate_ui_design.rs | 0 .../editor_tools/generate_video.rs | 0 .../src/editor_agent/editor_tools/mod.rs | 8 +- 9 files changed, 365 insertions(+), 9 deletions(-) create mode 100644 server-rs/crates/api-server/src/editor_agent/editor_tools/generate_background_music.rs create mode 100644 server-rs/crates/api-server/src/editor_agent/editor_tools/generate_character.rs create mode 100644 server-rs/crates/api-server/src/editor_agent/editor_tools/generate_icon_spritesheet.rs create mode 100644 server-rs/crates/api-server/src/editor_agent/editor_tools/generate_image.rs create mode 100644 server-rs/crates/api-server/src/editor_agent/editor_tools/generate_sound_effect.rs create mode 100644 server-rs/crates/api-server/src/editor_agent/editor_tools/generate_ui_design.rs create mode 100644 server-rs/crates/api-server/src/editor_agent/editor_tools/generate_video.rs diff --git a/server-rs/crates/api-server/src/editor_agent/api.rs b/server-rs/crates/api-server/src/editor_agent/api.rs index 251e5c378..bf72db7ec 100644 --- a/server-rs/crates/api-server/src/editor_agent/api.rs +++ b/server-rs/crates/api-server/src/editor_agent/api.rs @@ -32,6 +32,7 @@ use crate::auth::AuthenticatedAccessToken; use crate::editor_agent::agent::LlmChatAgentBuilder; use crate::editor_agent::editor_tools::common::EditorToolContext; use crate::editor_agent::editor_tools::edit_image::{EditImageTool, EditImageToolArgs}; +use crate::editor_agent::editor_tools::generate_image::{GenerateImageTool, GenerateImageToolArgs}; use crate::editor_agent::utils::{ EditorAgentMessageResponse, ImageId, ImageMetadata, IntoImageId, build_editor_agent_canvas_completion, conversation_detail_from_record, @@ -146,7 +147,10 @@ pub async fn editor_agent_message( let mut agent = LlmChatAgentBuilder::new() .with_client(llm_client) - .tool(EditImageTool { context: tool_context }) + .tool(EditImageTool { + context: tool_context.clone(), + }) + .tool(GenerateImageTool { context: tool_context }) .max_turns(3) .memory(memory) .build(); @@ -500,8 +504,76 @@ pub async fn confirm_editor_agent_tool_call( let tool_name = tc.tool_name.clone(); let tool_args = tc.args.clone(); - // TODO dont use match + // Each confirmed tool reuses its existing editor-generation entry point. match tool_name.as_str() { + GenerateImageTool::NAME => { + let args: GenerateImageToolArgs = serde_json::from_value(tc.args.clone()) + .map_err(|e| editor_agent_bad_request(format!("invalid tool call args: {e}")))?; + let project = state + .spacetime_client() + .get_editor_project(EditorProjectGetRecordInput { + project_id: conversation.project_id.clone(), + owner_user_id: conversation.owner_user_id.clone(), + }) + .await + .map_err(|e| { + AppError::from_status(axum::http::StatusCode::NOT_FOUND) + .with_details(json!({ "message": format!("project not found: {e}") })) + })?; + let caller = EditorGenerationCaller { + owner_user_id: conversation.owner_user_id.clone(), + audit_subject_user_id: Some(conversation.owner_user_id.clone()), + audit_project_id: Some(conversation.project_id.clone()), + }; + let title = args.prompt.clone(); + let canvas_completion = + build_editor_agent_canvas_completion(&project, GenerateImageTool::NAME, &title); + let tool_context = build_tool_context(&document); + let generate_tool = GenerateImageTool { context: tool_context }; + let result = generate_tool + .execute( + &state, + &request_context, + caller, + args.clone(), + conversation.project_id.clone(), + Some(json!({ + "source": "editor-agent", + "conversationId": conversation.conversation_id, + "toolCallMessageId": message_id, + "fields": [{ "title": "用户指令", "value": args.prompt }], + })), + Some(title), + canvas_completion, + ) + .await?; + let generated_images = vec![EditorAgentGeneratedImage { + resource_id: result + .resource + .as_ref() + .and_then(|resource| resource.get("resourceId")) + .and_then(Value::as_str) + .map(ToOwned::to_owned), + object_key: result.object_key.clone(), + asset_object_id: result.asset_object_id.clone(), + image_src: result.image_src.clone(), + thumbnail_src: None, + width: Some(result.width), + height: Some(result.height), + }]; + let formatted_text = generate_tool.format_execute_message(&args, &result); + let msg = &mut document.messages[message_id]; + msg.text = formatted_text; + if let Some(tool_call) = &mut msg.tool_call { + tool_call.status = EditorAgentToolCallStatus::Completed; + tool_call.images = generated_images; + } + write_messages_document(&state, &conversation, &document).await?; + Ok(json_success_body( + Some(&request_context), + &document.messages[message_id], + )) + } EditImageTool::NAME => { let args: EditImageToolArgs = serde_json::from_value(tool_args) .map_err(|e| editor_agent_bad_request(format!("invalid tool call args: {e}")))?; diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_background_music.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_background_music.rs new file mode 100644 index 000000000..e69de29bb diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_character.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_character.rs new file mode 100644 index 000000000..e69de29bb diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_icon_spritesheet.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_icon_spritesheet.rs new file mode 100644 index 000000000..e69de29bb diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_image.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_image.rs new file mode 100644 index 000000000..7a0850923 --- /dev/null +++ b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_image.rs @@ -0,0 +1,290 @@ +use crate::editor_agent::editor_tools::common::EditorToolContext; +use crate::editor_agent::utils::ImageId; +use crate::editor_project::{ + EditorGenerationCaller, EditorImageGenerationRequest, generate_editor_image_for_owner, +}; +use crate::http_error::AppError; +use crate::request_context::RequestContext; +use crate::state::AppState; +use axum::http::StatusCode; +use module_editor_agent::agent::tool::{Tool, ToolFailure, ToolFailureKind}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use shared_contracts::api::ApiSuccessEnvelope; +use shared_contracts::assets::EditorCanvasGenerationCompletionPayload; +use std::error::Error; +use std::fmt::Display; + +pub struct GenerateImageTool { + pub context: EditorToolContext, +} + +#[derive(Debug, Clone)] +pub enum GenerateImageError { + PromptNotProvided, + AssetNotFound(ImageId), +} + +impl Display for GenerateImageError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::PromptNotProvided => write!(f, "prompt not provided"), + Self::AssetNotFound(image_id) => write!(f, "asset {image_id} not found in context"), + } + } +} + +impl Error for GenerateImageError {} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GenerateImageToolArgs { + pub prompt: String, + #[serde(default)] + pub reference_image_ids: Vec, + #[serde(default)] + // TODO restrict to a set of possible values + pub aspect_ratio: Option, + #[serde(default)] + pub image_size: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GenerateImageToolOutput { + pub message: String, +} + +impl Tool for GenerateImageTool { + const NAME: &'static str = "generate-image"; + type Error = GenerateImageError; + type Args = GenerateImageToolArgs; + type Output = GenerateImageToolOutput; + + fn description(&self) -> String { + "根据文字描述生成一张新图片；可选参考图仅用于借鉴画风、构图或元素，不会修改参考图本身。".to_string() + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "完整的生图提示词，包含主体、场景、风格、构图和背景。" + }, + "reference_image_ids": { + "type": "array", + "items": { "type": "string" }, + "description": "可选参考图 ID 列表，用于提供画风或元素参考。" + }, + "aspect_ratio": { + "type": "string", + "description": "可选画面比例，例如 1:1、16:9、9:16。" + }, + "image_size": { + "type": "string", + "description": "可选图片清晰度，例如 1K、2K、4K。" + } + }, + "required": ["prompt"], + "additionalProperties": false + }) + } + + fn call( + &self, + args: Self::Args, + ) -> impl Future> + Send { + async move { + self.validate_args(&args)?; + Ok(GenerateImageToolOutput { + message: "this tool call is pending user confirmation. if all is pending, just end this turn".to_string(), + }) + } + } + + fn classify_error(&self, error: &Self::Error) -> ToolFailure { + match error { + GenerateImageError::AssetNotFound(_) => { + ToolFailure::new(ToolFailureKind::NotFound, error.to_string()) + } + GenerateImageError::PromptNotProvided => ToolFailure::invalid_args(error.to_string()), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct EditorImageGenerationResult { + pub image_src: String, + pub object_key: Option, + pub asset_object_id: Option, + pub width: u32, + pub height: u32, + pub source_type: String, + pub prompt: String, + pub actual_prompt: Option, + pub model: String, + pub provider: String, + pub task_id: String, + pub resource: Option, + pub asset: Option, + pub project: Option, +} + +impl GenerateImageTool { + fn validate_args(&self, args: &GenerateImageToolArgs) -> Result<(), GenerateImageError> { + if args.prompt.trim().is_empty() { + return Err(GenerateImageError::PromptNotProvided); + } + for image_id in &args.reference_image_ids { + if !self.context.contains_image(image_id) { + return Err(GenerateImageError::AssetNotFound(image_id.clone())); + } + } + Ok(()) + } + + pub fn format_execute_message( + &self, + args: &GenerateImageToolArgs, + result: &EditorImageGenerationResult, + ) -> String { + let args = serde_json::to_string(args).unwrap_or_default(); + let image_id = result + .object_key + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|value| value.trim_start_matches('/').to_string()) + .unwrap_or_else(|| result.image_src.clone()); + format!( + "[tool_call:{}] args: {args} output: generated result saved as image: {image_id}", + Self::NAME, + ) + } + + pub async fn execute( + &self, + state: &AppState, + request_context: &RequestContext, + caller: EditorGenerationCaller, + args: GenerateImageToolArgs, + project_id: String, + generation_inputs: Option, + asset_label: Option, + canvas_completion: EditorCanvasGenerationCompletionPayload, + ) -> Result { + self.validate_args(&args).map_err(|error| { + AppError::from_status(StatusCode::BAD_REQUEST).with_details(json!({ + "provider": "editor-agent", + "message": error.to_string(), + })) + })?; + + let result = generate_editor_image_for_owner( + state, + request_context, + caller, + EditorImageGenerationRequest { + prompt: args.prompt, + size: None, + kind: None, + model: None, + screen_color: None, + seg_model: None, + aspect_ratio: args.aspect_ratio, + image_size: args.image_size, + reference_image_srcs: Some( + args.reference_image_ids + .into_iter() + .map(|image_id| image_id.id) + .collect(), + ), + project_id: Some(project_id), + asset_kind: Some("editor_agent_generated_image".to_string()), + generation_inputs, + asset_folder_id: Some("project".to_string()), + asset_label, + source_resource_id: None, + canvas_completion: Some(canvas_completion), + }, + ) + .await? + .0; + + let data = if request_context.wants_envelope() { + serde_json::from_value::>(result) + .map_err(|error| { + AppError::from_status(StatusCode::INTERNAL_SERVER_ERROR).with_details(json!({ + "message": format!("failed to parse success envelope: {error}"), + })) + })? + .data + } else { + result + }; + + serde_json::from_value::(data).map_err(|error| { + AppError::from_status(StatusCode::INTERNAL_SERVER_ERROR).with_details(json!({ + "message": format!("failed to deserialize image generation result: {error}"), + })) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + fn tool() -> GenerateImageTool { + GenerateImageTool { + context: EditorToolContext { + images: HashMap::from([( + ImageId { id: "reference.png".to_string() }, + crate::editor_agent::utils::ImageMetadata { tag: String::new() }, + )]), + }, + } + } + + #[tokio::test] + async fn proposal_accepts_prompt_and_known_references() { + let output = tool() + .call(GenerateImageToolArgs { + prompt: "雨夜的霓虹街道".to_string(), + reference_image_ids: vec![ImageId { id: "reference.png".to_string() }], + aspect_ratio: Some("16:9".to_string()), + image_size: Some("2K".to_string()), + }) + .await + .expect("proposal should be valid"); + + assert!(output.message.contains("pending user confirmation")); + } + + #[tokio::test] + async fn proposal_rejects_missing_prompt_and_unknown_reference() { + let prompt_error = tool() + .call(GenerateImageToolArgs { + prompt: " ".to_string(), + reference_image_ids: Vec::new(), + aspect_ratio: None, + image_size: None, + }) + .await + .expect_err("blank prompt must fail"); + assert!(matches!(prompt_error, GenerateImageError::PromptNotProvided)); + + let reference_error = tool() + .call(GenerateImageToolArgs { + prompt: "一棵树".to_string(), + reference_image_ids: vec![ImageId { id: "missing.png".to_string() }], + aspect_ratio: None, + image_size: None, + }) + .await + .expect_err("unknown reference must fail"); + assert!(matches!(reference_error, GenerateImageError::AssetNotFound(_))); + } +} diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_sound_effect.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_sound_effect.rs new file mode 100644 index 000000000..e69de29bb diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_ui_design.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_ui_design.rs new file mode 100644 index 000000000..e69de29bb diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_video.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_video.rs new file mode 100644 index 000000000..e69de29bb diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/mod.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/mod.rs index 6eb9e0195..695076375 100644 --- a/server-rs/crates/api-server/src/editor_agent/editor_tools/mod.rs +++ b/server-rs/crates/api-server/src/editor_agent/editor_tools/mod.rs @@ -1,9 +1,3 @@ pub mod common; pub mod edit_image; -mod generate_ui_design; -mod generate_character; -mod generate_image; -mod generate_icon_spritesheet; -mod generate_sound_effect; -mod generate_background_music; -mod generate_video; +pub mod generate_image;