GenarrativeAI/Genarrative
11
0
Fork 0
Code Issues 14 Pull Requests 3 Actions Packages Projects 3 Releases Wiki Activity

centralize data key query instead of using obj key.

Browse Source
This commit is contained in:
王德宇
2026-07-10 19:06:39 +08:00
parent 532b6434bd
commit c55a90d23e
8 changed files with 216 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server-rs/crates/api-server/src/editor_agent/api.rs
+60 -7
View File
@@ -48,7 +48,7 @@ use crate::editor_agent::editor_tools::generate_sound_effect::{
use crate::editor_agent::editor_tools::generate_ui_design::GenerateUiDesignTool;
use crate::editor_agent::editor_tools::generate_video::{GenerateVideoTool, GenerateVideoToolArgs};
use crate::editor_agent::utils::{
ImageId, ImageMetadata, IntoImageId, build_editor_agent_canvas_completion,
ImageId, ImageMetadata, IntoDataKey, IntoImageId, build_editor_agent_canvas_completion,
conversation_detail_from_record, conversation_summary_from_record, editor_agent_bad_request,
empty_messages_document, ensure_editor_project_access, normalize_editor_agent_attachments,
now_rfc3339, read_messages_document, require_editor_agent_sidebar_enabled,
@@ -270,18 +270,22 @@ fn build_tool_context(document: &EditorAgentConversationMessagesDocument) -> Edi
for msg in document.messages.iter().rev() {
if let Some(tc) = &msg.tool_call {
for img in &tc.images {
let image_id = ImageId {
id: img.object_key.clone().unwrap_or_default(),
let data_key = img.clone().into_data_key();
let image_id = ImageId::from_data_key(&data_key);
let metadata = ImageMetadata {
tag: String::new(),
data_key,
};
let metadata = ImageMetadata { tag: String::new() };
images.insert(image_id, metadata);
}
}
for a in &msg.attachments {
let image_id = ImageId {
id: a.object_key.clone().unwrap_or_default(),
let data_key = a.clone().into_data_key();
let image_id = ImageId::from_data_key(&data_key);
let metadata = ImageMetadata {
tag: String::new(),
data_key,
};
let metadata = ImageMetadata { tag: String::new() };
images.insert(image_id, metadata);
}
}
@@ -289,6 +293,55 @@ fn build_tool_context(document: &EditorAgentConversationMessagesDocument) -> Edi
EditorToolContext { images }
}
#[cfg(test)]
mod tool_context_tests {
use super::*;
use shared_contracts::editor_agent::{
EDITOR_AGENT_MESSAGES_DOCUMENT_VERSION, EditorAgentGeneratedImage,
EditorAgentMessageRole, EditorAgentToolCallStatus,
};
#[test]
fn tool_context_maps_hashed_ids_back_to_real_data_keys() {
let data_key = "generated-images/proj-1/output.png";
let document = EditorAgentConversationMessagesDocument {
version: EDITOR_AGENT_MESSAGES_DOCUMENT_VERSION,
conversation_id: "editor-agent-conversation-1".to_string(),
messages: vec![EditorAgentMessage {
id: 0,
role: EditorAgentMessageRole::System,
text: String::new(),
attachments: Vec::new(),
tool_call: Some(EditorAgentToolCall {
tool_name: "generate-image".to_string(),
summary: String::new(),
status: EditorAgentToolCallStatus::Completed,
args: json!({}),
images: vec![EditorAgentGeneratedImage {
resource_id: None,
object_key: Some(data_key.to_string()),
asset_object_id: None,
image_src: "/generated-images/proj-1/output.png".to_string(),
thumbnail_src: None,
width: None,
height: None,
}],
videos: Vec::new(),
audios: Vec::new(),
error: None,
}),
created_at: String::new(),
}],
};
let context = build_tool_context(&document);
let image_id = ImageId::from_data_key(data_key);
assert_eq!(context.image_data_key(&image_id), Some(data_key));
assert!(!image_id.id.contains(data_key));
}
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct EditorAgentConversationDeleteResponse {
server-rs/crates/api-server/src/editor_agent/editor_tools/common.rs
+6
View File
@@ -13,4 +13,10 @@ impl EditorToolContext {
pub fn contains_image(&self, image_id: &ImageId) -> bool {
self.images.contains_key(image_id)
}
pub fn image_data_key(&self, image_id: &ImageId) -> Option<&str> {
self.images
.get(image_id)
.map(|metadata| metadata.data_key.as_str())
}
}
server-rs/crates/api-server/src/editor_agent/editor_tools/edit_image.rs
+28 -14
View File
@@ -178,16 +178,14 @@ impl EditImageTool {
let name = Self::NAME.to_string();
let arg_json = serde_json::to_string(&args).unwrap_or_default();
// TODO use the trait
let image_id: ImageId = ImageId {
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()),
};
let data_key = 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());
let image_id = ImageId::from_data_key(data_key);
format!(
"[tool_call:{name}] args: {arg_json} output: edited result saved as image: {image_id}"
)
@@ -206,13 +204,22 @@ impl EditImageTool {
source_resource_id: Option<String>,
canvas_completion: Option<EditorCanvasGenerationCompletionPayload>,
) -> Result<EditorImageEditResult, AppError> {
// TODO id should not be assumed as src key
let source_image_src = args.object_image_id.id;
let reference_image_srcs: Vec<String> = args
let source_image_src = self
.context
.image_data_key(&args.object_image_id)
.ok_or_else(|| image_not_found(&args.object_image_id))?
.to_string();
let reference_image_srcs: Result<Vec<String>, AppError> = args
.reference_image_ids
.into_iter()
.map(|id| id.id)
.map(|image_id| {
self.context
.image_data_key(&image_id)
.map(ToOwned::to_owned)
.ok_or_else(|| image_not_found(&image_id))
})
.collect();
let reference_image_srcs = reference_image_srcs?;
let result = edit_editor_image_for_owner(
state,
@@ -256,3 +263,10 @@ impl EditImageTool {
})
}
}
fn image_not_found(image_id: &ImageId) -> AppError {
AppError::from_status(StatusCode::BAD_REQUEST).with_details(json!({
"provider": "editor-agent",
"message": format!("asset {image_id} not found in context"),
}))
}
server-rs/crates/api-server/src/editor_agent/editor_tools/generate_character.rs
+6 -2
View File
@@ -107,10 +107,14 @@ pub(crate) fn format_tool_message(
result: &EditorImageGenerationResult,
) -> String {
let args = serde_json::to_string(args).unwrap_or_default();
let image_id = result
let data_key = result
.object_key
.as_deref()
.unwrap_or(result.image_src.as_str());
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.trim_start_matches('/').to_string())
.unwrap_or_else(|| result.image_src.clone());
let image_id = crate::editor_agent::utils::ImageId::from_data_key(data_key);
format!(
"[tool_call:{tool_name}] args: {args} output: generated result saved as image: {image_id}"
)
server-rs/crates/api-server/src/editor_agent/editor_tools/generate_icon_spritesheet.rs
+30 -4
View File
@@ -170,12 +170,21 @@ impl GenerateIconSpritesheetTool {
request_context,
caller,
EditorIconSpritesheetGenerationRequest {
reference_image_src: args.reference_image_id.id,
reference_image_src: self
.context
.image_data_key(&args.reference_image_id)
.ok_or_else(|| image_not_found(&args.reference_image_id))?
.to_string(),
reference_image_srcs: Some(
args.reference_image_ids
.into_iter()
.map(|id| id.id)
.collect(),
.map(|image_id| {
self.context
.image_data_key(&image_id)
.map(ToOwned::to_owned)
.ok_or_else(|| image_not_found(&image_id))
})
.collect::<Result<Vec<_>, _>>()?,
),
icon_descriptions: args.icon_descriptions,
model: None,
@@ -212,10 +221,27 @@ impl GenerateIconSpritesheetTool {
result: &EditorIconSpritesheetResult,
) -> String {
let args = serde_json::to_string(args).unwrap_or_default();
let data_key = result
.spritesheet_resource
.as_ref()
.and_then(|resource| resource.get("objectKey"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.trim_start_matches('/').to_string())
.unwrap_or_else(|| result.spritesheet_image_src.clone());
let image_id = ImageId::from_data_key(data_key);
format!(
"[tool_call:{}] args: {args} output: generated spritesheet saved as image: {}",
Self::NAME,
result.spritesheet_image_src
image_id
)
}
}
fn image_not_found(image_id: &ImageId) -> AppError {
AppError::from_status(StatusCode::BAD_REQUEST).with_details(json!({
"provider": "editor-agent",
"message": format!("asset {image_id} not found in context"),
}))
}
server-rs/crates/api-server/src/editor_agent/editor_tools/generate_image.rs
+27 -14
View File
@@ -98,7 +98,7 @@ impl Tool for GenerateImageTool {
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(),
message: "this tool call is pending user confirmation.".to_string(),
})
}
}
@@ -154,13 +154,14 @@ impl GenerateImageTool {
result: &EditorImageGenerationResult,
) -> String {
let args = serde_json::to_string(args).unwrap_or_default();
let image_id = result
let data_key = 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());
let image_id = ImageId::from_data_key(data_key);
format!(
"[tool_call:{}] args: {args} output: generated result saved as image: {image_id}",
Self::NAME,
@@ -231,12 +232,7 @@ impl GenerateImageTool {
.then(|| "birefnet".to_string()),
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(),
),
reference_image_srcs: Some(self.image_data_keys(args.reference_image_ids)?),
project_id: Some(project_id),
asset_kind: Some(asset_kind.to_string()),
generation_inputs,
@@ -267,6 +263,23 @@ impl GenerateImageTool {
}))
})
}
fn image_data_keys(&self, image_ids: Vec<ImageId>) -> Result<Vec<String>, AppError> {
image_ids
.into_iter()
.map(|image_id| {
self.context
.image_data_key(&image_id)
.map(ToOwned::to_owned)
.ok_or_else(|| {
AppError::from_status(StatusCode::BAD_REQUEST).with_details(json!({
"provider": "editor-agent",
"message": format!("asset {image_id} not found in context"),
}))
})
})
.collect()
}
}
#[cfg(test)]
@@ -275,13 +288,15 @@ mod tests {
use std::collections::HashMap;
fn tool() -> GenerateImageTool {
let image_id = ImageId::from_data_key("reference.png");
GenerateImageTool {
context: EditorToolContext {
images: HashMap::from([(
ImageId {
id: "reference.png".to_string(),
image_id,
crate::editor_agent::utils::ImageMetadata {
tag: String::new(),
data_key: "reference.png".to_string(),
},
crate::editor_agent::utils::ImageMetadata { tag: String::new() },
)]),
},
}
@@ -292,9 +307,7 @@ mod tests {
let output = tool()
.call(GenerateImageToolArgs {
prompt: "雨夜的霓虹街道".to_string(),
reference_image_ids: vec![ImageId {
id: "reference.png".to_string(),
}],
reference_image_ids: vec![ImageId::from_data_key("reference.png")],
aspect_ratio: Some("16:9".to_string()),
image_size: Some("2K".to_string()),
})
server-rs/crates/api-server/src/editor_agent/editor_tools/generate_video.rs
+17 -3
View File
@@ -125,11 +125,17 @@ impl GenerateVideoTool {
json!({ "provider": "editor-agent", "message": "视频提示词不能为空" }),
));
}
let reference_image_srcs = args
let reference_image_srcs: Result<Vec<String>, AppError> = args
.reference_image_ids
.into_iter()
.map(|id| id.id)
.map(|image_id| {
self.context
.image_data_key(&image_id)
.map(ToOwned::to_owned)
.ok_or_else(|| image_not_found(&image_id))
})
.collect();
let reference_image_srcs = reference_image_srcs?;
let response = generate_editor_video_for_owner(
state.clone(),
request_context.clone(),
@@ -164,11 +170,19 @@ impl GenerateVideoTool {
args: &GenerateVideoToolArgs,
result: &EditorVideoGenerateResponse,
) -> String {
let video_id = ImageId::from_data_key(result.video_src.as_str());
format!(
"[tool_call:{}] args: {} output: generated video saved as: {}",
Self::NAME,
serde_json::to_string(args).unwrap_or_default(),
result.video_src
video_id
)
}
}
fn image_not_found(image_id: &ImageId) -> AppError {
AppError::from_status(StatusCode::BAD_REQUEST).with_details(json!({
"provider": "editor-agent",
"message": format!("asset {image_id} not found in context"),
}))
}
server-rs/crates/api-server/src/editor_agent/utils.rs
+42 -10
View File
@@ -8,12 +8,13 @@ use platform_oss::{
};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use shared_contracts::assets::{
EditorCanvasGenerationCompletionPayload, EditorCanvasGenerationPlaceholderPayload,
};
use shared_contracts::editor_agent::{
EDITOR_AGENT_MESSAGES_DOCUMENT_VERSION, EditorAgentAttachmentRef, EditorAgentAttachmentSource,
EditorAgentConversationDetail, EditorAgentConversationMessagesDocument,
EditorAgentConversationDetail, EditorAgentConversationMessagesDocument, EditorAgentGeneratedImage,
EditorAgentConversationSummary, EditorAgentMessage,
};
use shared_kernel::{normalize_optional_string, normalize_required_string};
@@ -39,6 +40,16 @@ impl IntoDataKey for EditorAgentAttachmentRef {
.unwrap_or_else(|| self.image_src.clone())
}
}
impl IntoDataKey for EditorAgentGeneratedImage {
fn into_data_key(self) -> String {
self.object_key
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.trim_start_matches('/').to_string())
.unwrap_or(self.image_src)
}
}
impl IntoDataKey for EditorProjectResourceRecord {
fn into_data_key(self) -> String {
self.object_key
@@ -66,23 +77,17 @@ pub trait IntoImageId {
}
impl IntoImageId for EditorAgentAttachmentRef {
fn into_image_id(self) -> ImageId {
ImageId {
id: self.into_data_key(),
}
ImageId::from_data_key(self.into_data_key())
}
}
impl IntoImageId for EditorProjectResourceRecord {
fn into_image_id(self) -> ImageId {
ImageId {
id: self.into_data_key(),
}
ImageId::from_data_key(self.into_data_key())
}
}
impl IntoImageId for EditorAssetRecord {
fn into_image_id(self) -> ImageId {
ImageId {
id: self.into_data_key(),
}
ImageId::from_data_key(self.into_data_key())
}
}
@@ -107,6 +112,16 @@ pub struct ImageId {
pub id: String,
}
impl ImageId {
/// Keep data keys server-side while giving the model a stable image reference.
pub fn from_data_key(data_key: impl AsRef<str>) -> Self {
let digest = Sha256::digest(data_key.as_ref().as_bytes());
Self {
id: format!("sha256:{digest:x}"),
}
}
}
impl Serialize for ImageId {
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
self.id.serialize(serializer)
@@ -122,6 +137,7 @@ impl<'de> Deserialize<'de> for ImageId {
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageMetadata {
pub tag: String,
pub data_key: String,
}
impl Display for ImageId {
@@ -130,6 +146,22 @@ impl Display for ImageId {
}
}
#[cfg(test)]
mod image_id_tests {
use super::ImageId;
#[test]
fn image_id_is_a_stable_hash_of_the_data_key() {
let data_key = "editor-projects/proj-1/image.png";
let image_id = ImageId::from_data_key(data_key);
assert_eq!(image_id, ImageId::from_data_key(data_key));
assert_ne!(image_id, ImageId::from_data_key("another-image.png"));
assert!(image_id.id.starts_with("sha256:"));
assert!(!image_id.id.contains(data_key));
}
}
type EditorAgentConversationLockMap = Mutex<BTreeMap<String, Arc<tokio::sync::Mutex<()>>>>;
static EDITOR_AGENT_CONVERSATION_LOCKS: OnceLock<EditorAgentConversationLockMap> = OnceLock::new();