impl image gen tool
This commit is contained in:
@@ -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}")))?;
|
||||
|
||||
@@ -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<ImageId>,
|
||||
#[serde(default)]
|
||||
// TODO restrict to a set of possible values
|
||||
pub aspect_ratio: Option<String>,
|
||||
#[serde(default)]
|
||||
pub image_size: Option<String>,
|
||||
}
|
||||
|
||||
#[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<Output = Result<Self::Output, Self::Error>> + 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<String>,
|
||||
pub asset_object_id: Option<String>,
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
pub source_type: String,
|
||||
pub prompt: String,
|
||||
pub actual_prompt: Option<String>,
|
||||
pub model: String,
|
||||
pub provider: String,
|
||||
pub task_id: String,
|
||||
pub resource: Option<Value>,
|
||||
pub asset: Option<Value>,
|
||||
pub project: Option<Value>,
|
||||
}
|
||||
|
||||
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<Value>,
|
||||
asset_label: Option<String>,
|
||||
canvas_completion: EditorCanvasGenerationCompletionPayload,
|
||||
) -> Result<EditorImageGenerationResult, AppError> {
|
||||
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::<ApiSuccessEnvelope<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::<EditorImageGenerationResult>(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(_)));
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user