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

move context to tool own. fix memory not update.

Browse Source
This commit is contained in:
王德宇
2026-07-10 14:21:36 +08:00
parent e5eee1dd4e
commit 337d358ede
8 changed files with 71 additions and 169 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server-rs/crates/api-server/src/editor_agent/agent.rs
+7 -11
View File
@@ -1,4 +1,5 @@
use serde_json::Value;
use tracing::log::info;
use module_editor_agent::agent::agent::{Agent, LlmApiAdaptor};
use module_editor_agent::agent::agent_builder::AgentBuilder;
use module_editor_agent::agent::error::PromptError;
@@ -6,16 +7,19 @@ use module_editor_agent::agent::hook::Hook;
use module_editor_agent::agent::memory::AgentMemory;
use module_editor_agent::agent::run::ToolCallFlow;
use module_editor_agent::agent::tool::{Tool, ToolCall, ToolDyn};
use platform_llm::{LlmClient, LlmMessage};
use platform_llm::{LlmClient, LlmMessage, LlmMessageRole};
pub(crate) struct LlmCompletionModel {
client: LlmClient,
}
impl LlmApiAdaptor<LlmMessage> for LlmCompletionModel {
async fn complete(&self, messages: &[LlmMessage]) -> Result<String, PromptError> {
async fn complete<'a>(
&self,
messages: impl Iterator<Item = &'a LlmMessage> + Send,
) -> Result<String, PromptError> {
use platform_llm::LlmTextRequest;
let request = LlmTextRequest::new(messages.to_vec()).with_request_timeout_ms(30_000);
let request = LlmTextRequest::new(messages.cloned().collect()).with_request_timeout_ms(30_000);
let response = self
.client
.request_text(request)
@@ -40,7 +44,6 @@ pub struct LlmChatAgentBuilder {
hooks: Vec<Box<dyn Hook>>,
max_turns: usize,
memory_data: Option<Box<dyn AgentMemory<LlmMessage>>>,
context: Option<serde_json::Value>,
}
impl AgentBuilder<LlmMessage, LlmCompletionModel> for LlmChatAgentBuilder {
@@ -54,7 +57,6 @@ impl AgentBuilder<LlmMessage, LlmCompletionModel> for LlmChatAgentBuilder {
hooks: Vec::new(),
max_turns: 10,
memory_data: None,
context: None,
}
}
@@ -88,11 +90,6 @@ impl AgentBuilder<LlmMessage, LlmCompletionModel> for LlmChatAgentBuilder {
self
}
fn context(mut self, context: Value) -> Self {
self.context = Some(context);
self
}
fn build(self) -> Agent<LlmCompletionModel, LlmMessage> {
let model = LlmCompletionModel {
client: self.client.expect("call .with_client() first"),
@@ -113,7 +110,6 @@ impl AgentBuilder<LlmMessage, LlmCompletionModel> for LlmChatAgentBuilder {
agent.hooks = self.hooks;
agent.default_max_turns = self.max_turns;
agent.memory = self.memory_data;
agent.context = self.context;
agent.system_prompt = Some(LlmMessage::system(&system_prompt));
agent
}
server-rs/crates/api-server/src/editor_agent/api.rs
+20 -10
View File
@@ -30,6 +30,7 @@ use spacetime_client::{
use crate::api_response::json_success_body;
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::utils::{EditorAgentMessageResponse, ImageId, ImageMetadata, 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, write_messages_document, IntoImageId};
use crate::editor_project::{EditorGenerationCaller, current_utc_micros, map_editor_project_error};
@@ -69,12 +70,22 @@ pub async fn editor_agent_message(
let mut document: EditorAgentConversationMessagesDocument =
read_messages_document(&state, &conversation).await?;
let now = now_rfc3339();
if !attachments.is_empty() {
let mut attachment_info = String::new();
attachment_info.push_str("user has just uploaded attachments of the order: ");
for a in &attachments {
attachment_info.push_str(&format!("{} ,", a.clone().into_image_id()))
}
document.messages.push(EditorAgentMessage{
id: 0,
role: EditorAgentMessageRole::System,
text: attachment_info,
attachments:Vec::new(),
tool_call: None,
created_at: now.clone(),
});
}
// Build conversation history as LlmMessage vec
@@ -90,7 +101,6 @@ pub async fn editor_agent_message(
// Save user message to document
let was_empty = document.messages.is_empty();
let user_now = now_rfc3339();
// TODO tell agent info about attachments
let user_message = EditorAgentMessage {
id: document.messages.len(),
@@ -98,7 +108,7 @@ pub async fn editor_agent_message(
text: payload.text.trim().to_string(),
attachments,
tool_call: None,
created_at: user_now,
created_at: now,
};
document.messages.push(user_message.clone());
write_messages_document(&state, &conversation, &document).await?;
@@ -118,7 +128,7 @@ pub async fn editor_agent_message(
}
// Build tool context from document
let tool_context_value = build_tool_context(&document);
let tool_context = build_tool_context(&document);
// Build and run agent
let llm_client = state.llm_client().ok_or_else(|| {
@@ -129,15 +139,14 @@ pub async fn editor_agent_message(
let memory = VecMemory::new(previous_messages);
let agent = LlmChatAgentBuilder::new()
let mut agent = LlmChatAgentBuilder::new()
.with_client(llm_client)
.tool(EditImageTool {})
.tool(EditImageTool { context: tool_context })
.max_turns(3)
.memory(memory)
.context(tool_context_value)
.build();
let agent_result = agent.prompt(LlmMessage::user("")).await;
let agent_result = agent.prompt(LlmMessage::user(user_message.text)).await;
let assistant_now = now_rfc3339();
@@ -205,7 +214,7 @@ fn build_delta_messages(
Ok(messages)
}
fn build_tool_context(document: &EditorAgentConversationMessagesDocument) -> Value {
fn build_tool_context(document: &EditorAgentConversationMessagesDocument) -> EditorToolContext {
let mut images: HashMap<ImageId, ImageMetadata> = HashMap::new();
for msg in document.messages.iter().rev() {
@@ -223,7 +232,7 @@ fn build_tool_context(document: &EditorAgentConversationMessagesDocument) -> Val
}
}
json!({ "images": images })
EditorToolContext { images }
}
#[derive(Debug, Serialize)]
@@ -508,7 +517,8 @@ pub async fn confirm_editor_agent_tool_call(
build_editor_agent_canvas_completion(&project, "edit-image", &title);
// Execute the real generation
let edit_tool = EditImageTool {};
let tool_context = build_tool_context(&document);
let edit_tool = EditImageTool { context: tool_context };
let result = edit_tool
.execute(
&state,
server-rs/crates/api-server/src/editor_agent/editor_tools/edit_image.rs
+7 -17
View File
@@ -12,18 +12,18 @@ use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use shared_contracts::api::ApiSuccessEnvelope;
use shared_contracts::assets::EditorCanvasGenerationCompletionPayload;
use std::collections::HashMap;
use std::error::Error;
use std::fmt::Display;
pub struct EditImageTool {}
pub struct EditImageTool {
pub context: EditorToolContext,
}
#[derive(Debug, Clone)]
pub enum EditImageError {
ObjectImageNotProvided,
PromptNotProvided,
AssetNotFound(ImageId),
ContextParseError(String),
}
impl Display for EditImageError {
@@ -36,9 +36,6 @@ impl Display for EditImageError {
EditImageError::AssetNotFound(image_id) => {
write!(f, "asset {image_id} not found in context")
}
EditImageError::ContextParseError(msg) => {
write!(f, "failed to parse tool context: {msg}")
}
}
}
}
@@ -101,10 +98,9 @@ impl Tool for EditImageTool {
fn call(
&self,
args: Self::Args,
context: Value,
) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send {
async move {
Self::validate_context_images(&args, &context)?;
self.validate_context_images(&args)?;
if let Some(error) = Self::validate_args(&args) {
return Err(error);
}
@@ -116,9 +112,6 @@ impl Tool for EditImageTool {
fn classify_error(&self, error: &Self::Error) -> ToolFailure {
match error {
EditImageError::ContextParseError(_) => {
ToolFailure::new(ToolFailureKind::Internal, error.to_string())
}
EditImageError::AssetNotFound(_) => {
ToolFailure::new(ToolFailureKind::NotFound, error.to_string())
}
@@ -163,18 +156,15 @@ impl EditImageTool {
/// Validate that all referenced images exist in the context.
fn validate_context_images(
&self,
args: &EditImageToolArgs,
context: &Value,
) -> Result<(), EditImageError> {
let tool_context: EditorToolContext = serde_json::from_value(context.clone())
.map_err(|e| EditImageError::ContextParseError(e.to_string()))?;
if !tool_context.contains_image(&args.object_image_id) {
if !self.context.contains_image(&args.object_image_id) {
return Err(EditImageError::AssetNotFound(args.object_image_id.clone()));
}
for ref_id in &args.reference_image_ids {
if !tool_context.contains_image(ref_id) {
if !self.context.contains_image(ref_id) {
return Err(EditImageError::AssetNotFound(ref_id.clone()));
}
}
server-rs/crates/api-server/src/editor_agent/editor_tools/mod.rs
+7
View File
@@ -1,2 +1,9 @@
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;
server-rs/crates/module-editor-agent/src/agent/agent.rs
+6 -91
View File
@@ -11,7 +11,6 @@ pub struct Agent<M: LlmApiAdaptor<Message>, Message> {
pub default_max_turns: usize,
pub system_prompt: Option<Message>,
pub memory: Option<Box<dyn AgentMemory<Message>>>,
pub context: Option<serde_json::Value>,
}
impl<M, Message> Agent<M, Message>
@@ -27,7 +26,6 @@ where
default_max_turns: 10,
system_prompt: None,
memory: None,
context: None,
}
}
@@ -61,7 +59,7 @@ where
}
pub fn prompt(&self, message: impl Into<Message> + Send) -> PromptRequest<'_, M, Message>
pub fn prompt(&mut self, message: impl Into<Message> + Send) -> PromptRequest<'_, M, Message>
where
Message: 'static,
{
@@ -70,96 +68,13 @@ where
}
pub trait LlmApiAdaptor<Message>: Send + Sync {
fn complete(
fn complete<'a>(
&self,
messages: &[Message],
) -> impl Future<Output = Result<String, PromptError>> + Send;
messages: impl Iterator<Item = &'a Message> + Send,
) -> impl Future<Output = Result<String, PromptError>> + Send
where
Message: 'a;
fn tool_result_message(&self, tool_name: &str, output: &str) -> Message;
fn build_assistant_message(&self, text: &str) -> Message;
}
#[cfg(test)]
mod tests {
use super::*;
use crate::agent::tool::{ToolFailure, ToolFailureKind};
use serde::Deserialize;
use serde_json::json;
use std::fmt::{Display, Formatter};
struct TestModel;
impl LlmApiAdaptor<String> for TestModel {
async fn complete(&self, _messages: &[String]) -> Result<String, PromptError> {
Ok(String::new())
}
fn tool_result_message(&self, tool_name: &str, output: &str) -> String {
format!("{tool_name}: {output}")
}
fn build_assistant_message(&self, text: &str) -> String {
text.to_string()
}
}
#[derive(Debug)]
enum TestToolError {
Recoverable,
Fatal,
}
impl Display for TestToolError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::Recoverable => write!(f, "recoverable failure"),
Self::Fatal => write!(f, "fatal failure"),
}
}
}
impl std::error::Error for TestToolError {}
#[derive(Deserialize)]
struct TestArgs {
fatal: bool,
}
struct FailingTool;
impl Tool for FailingTool {
const NAME: &'static str = "fail";
type Error = TestToolError;
type Args = TestArgs;
type Output = ();
fn description(&self) -> String {
"test failing tool".to_string()
}
fn parameters(&self) -> serde_json::Value {
json!({"type": "object"})
}
async fn call(
&self,
args: Self::Args,
_context: serde_json::Value,
) -> Result<Self::Output, Self::Error> {
if args.fatal {
Err(TestToolError::Fatal)
} else {
Err(TestToolError::Recoverable)
}
}
fn classify_error(&self, error: &Self::Error) -> ToolFailure {
match error {
TestToolError::Recoverable => ToolFailure::other(error.to_string()),
TestToolError::Fatal => {
ToolFailure::new(ToolFailureKind::Internal, error.to_string())
}
}
}
}
}
server-rs/crates/module-editor-agent/src/agent/agent_builder.rs
-1
View File
@@ -14,6 +14,5 @@ pub trait AgentBuilder<Message, M: LlmApiAdaptor<Message>> {
fn max_turns(self, n: usize) -> Self;
fn memory(self, memory: impl AgentMemory<Message> + 'static) -> Self;
fn context(self, context: serde_json::Value) -> Self;
fn build(self) -> Agent<M, Message>;
}
server-rs/crates/module-editor-agent/src/agent/run.rs
+21 -33
View File
@@ -2,7 +2,7 @@ use crate::agent::agent::Agent;
use crate::agent::agent::LlmApiAdaptor;
use crate::agent::error::PromptError;
use crate::agent::hook::Hook;
use crate::agent::memory::{AgentMemory, VecMemory};
use crate::agent::memory::VecMemory;
use crate::agent::run::PromptOutput::{Text, Tool};
use crate::agent::tool::{ToolCall, ToolDyn, ToolExecutionResult, ToolFailure, ToolOutcome};
use serde::Deserialize;
@@ -36,12 +36,11 @@ pub enum ToolCallFlow {
}
pub struct PromptRequest<'a, M: LlmApiAdaptor<Message> + 'a, Message: 'a> {
agent: &'a Agent<M, Message>,
agent: &'a mut Agent<M, Message>,
message: Message,
system_prompt: Option<Message>,
hooks: Vec<Box<dyn Hook + 'static>>,
max_turns: usize,
context: Option<serde_json::Value>,
}
impl<'a, M, Message> PromptRequest<'a, M, Message>
@@ -49,14 +48,14 @@ where
M: LlmApiAdaptor<Message> + 'a,
Message: 'a,
{
pub fn new(agent: &'a Agent<M, Message>, message: Message) -> Self {
pub fn new(agent: &'a mut Agent<M, Message>, message: Message) -> Self {
let max_turns = agent.default_max_turns;
Self {
agent,
message,
system_prompt: None,
hooks: Vec::new(),
max_turns: agent.default_max_turns,
context: agent.context.clone(),
max_turns,
}
}
@@ -79,7 +78,7 @@ where
impl<'a, M, Message> IntoFuture for PromptRequest<'a, M, Message>
where
M: LlmApiAdaptor<Message> + Send + Sync + 'a,
Message: Send + Sync + Clone + 'a,
Message: Send + Sync + Clone + 'a + 'static,
{
type Output = Result<Vec<PromptOutput>, PromptError>;
type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>;
@@ -91,32 +90,28 @@ where
let max_turns = self.max_turns;
Box::pin(async move {
let mut memory: Box<dyn AgentMemory<Message>> = match &agent.memory {
Some(m) => {
let msgs: Vec<Message> = m.get_memory().iter().cloned().collect();
Box::new(VecMemory::new(msgs))
}
None => Box::new(VecMemory::new(Vec::new())),
};
let memory = agent
.memory
.get_or_insert_with(|| Box::new(VecMemory::new(Vec::new())));
memory.append_message(message);
let mut prompt_result: Vec<PromptOutput> = Vec::new();
let mut history: Vec<Message> = Vec::new();
if let Some(ref sp) = agent.system_prompt {
history.push(sp.clone());
}
history.extend(memory.get_memory().iter().cloned());
history.push(message);
for _ in 0..max_turns {
// TODO perf issue for copy cost
let text = agent.model.complete(&history).await?;
let text = {
let messages = agent
.system_prompt
.iter()
.chain(memory.get_memory().iter());
agent.model.complete(messages).await?
};
// Try to parse the LLM reply as JSON (handle Markdown fences)
let cleaned = clean_json_response(&text);
match serde_json::from_str::<LlmJsonResponse>(&cleaned) {
Ok(json_resp) => {
// Check for <end/> marker — stop the _turn immediately
let reply_has_end = json_resp.reply_text.contains("<end/>");
// end cond: <end/> and no tool_calls
let reply_has_end = json_resp.reply_text.contains("<end/>")
&& json_resp.tool_calls.is_empty();
let clean_text = json_resp
.reply_text
.replace("<end/>", "")
@@ -152,10 +147,6 @@ where
})
.collect();
if tool_calls.is_empty() {
return Ok(vec![Text(clean_text)]);
}
for (tc_id, tc) in tool_calls.iter().enumerate() {
// inline run_hooks: before_tool_call hook
let mut should_skip = false;
@@ -185,12 +176,9 @@ where
let tools: Vec<&Box<dyn ToolDyn>> = agent.tools.iter().collect();
let name = tc.name.clone();
let args = tc.args.clone();
let context = self.context.clone().unwrap_or_default();
let fut = async move {
match tools.iter().find(|tool| tool.tool_name() == name) {
Some(tool) => {
tool.call_with_context(args, context.clone()).await
}
Some(tool) => tool.call(args).await,
None => ToolExecutionResult::failed(
Value::Null,
ToolFailure::invalid_args(format!(
server-rs/crates/module-editor-agent/src/agent/tool.rs
+3 -6
View File
@@ -138,7 +138,6 @@ pub trait Tool: Sized {
fn call(
&self,
args: Self::Args,
context: serde_json::Value,
) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send;
fn classify_error(&self, error: &Self::Error) -> ToolFailure {
@@ -151,10 +150,9 @@ pub trait ToolDyn: Send + Sync {
fn tool_name(&self) -> &'static str;
fn description(&self) -> String;
fn parameters(&self) -> serde_json::Value;
fn call_with_context(
fn call(
&self,
args: serde_json::Value,
_context: serde_json::Value,
) -> Pin<Box<dyn Future<Output = ToolExecutionResult> + Send + '_>>;
}
@@ -171,10 +169,9 @@ impl<T: Tool + Send + Sync> ToolDyn for T {
self.parameters()
}
fn call_with_context(
fn call(
&self,
args: serde_json::Value,
context: serde_json::Value,
) -> Pin<Box<dyn Future<Output = ToolExecutionResult> + Send + '_>> {
Box::pin(async move {
let parsed: T::Args = match serde_json::from_value(args) {
@@ -187,7 +184,7 @@ impl<T: Tool + Send + Sync> ToolDyn for T {
}
};
let output = match self.call(parsed, context).await {
let output = match self.call(parsed).await {
Ok(output) => output,
Err(error) => {
return ToolExecutionResult::failed(