Genarrative/server-rs/crates/platform-editor-agent/src/agent/agent.rs

use crate::framework::agent::{Agent, LlmApiAdaptor};
use crate::framework::agent_builder::AgentBuilder;
use crate::framework::error::PromptError;
use crate::framework::hook::Hook;
use crate::framework::memory::AgentMemory;
use crate::framework::tool::{Tool, ToolDyn};
use platform_llm::{EDITOR_AGENT_GPT5_MODEL, LlmClient, LlmMessage, LlmTextRequest};
use serde_json::Value;

const EDITOR_AGENT_LLM_MAX_OUTPUT_TOKENS: u32 = 1024;
const EDITOR_AGENT_LLM_HARD_REQUEST_TIMEOUT_MS: u64 = 480_000;

pub struct LlmCompletionModel {
    client: LlmClient,
}

impl LlmApiAdaptor<LlmMessage> for LlmCompletionModel {
    async fn complete<'a>(
        &self,
        messages: impl Iterator<Item = &'a LlmMessage> + Send,
    ) -> Result<String, PromptError> {
        let request = build_editor_agent_llm_request(messages.cloned().collect());
        let response = self
            .client
            .request_text(request)
            .await
            .map_err(|e| PromptError::CompletionError(e.to_string()))?;
        Ok(response.content)
    }

    fn tool_result_message(&self, tool_name: &str, output: &str) -> LlmMessage {
        LlmMessage::system(format!("Tool '{tool_name}' returned: {output}"))
    }

    fn build_assistant_message(&self, text: &str) -> LlmMessage {
        LlmMessage::assistant(text)
    }
}

fn build_editor_agent_llm_request(messages: Vec<LlmMessage>) -> LlmTextRequest {
    LlmTextRequest::new(messages)
        .with_model(EDITOR_AGENT_GPT5_MODEL)
        .with_max_tokens(EDITOR_AGENT_LLM_MAX_OUTPUT_TOKENS)
        .with_request_timeout_ms(EDITOR_AGENT_LLM_HARD_REQUEST_TIMEOUT_MS)
}

pub struct LlmChatAgentBuilder {
    client: Option<LlmClient>,
    system_prompt_parts: Vec<String>,
    tools: Vec<Box<dyn ToolDyn>>,
    hooks: Vec<Box<dyn Hook>>,
    max_turns: usize,
    memory_data: Option<Box<dyn AgentMemory<LlmMessage>>>,
}

impl AgentBuilder<LlmMessage, LlmCompletionModel> for LlmChatAgentBuilder {
    type Client = LlmClient;

    fn new() -> Self {
        Self {
            client: None,
            system_prompt_parts: Vec::new(),
            tools: Vec::new(),
            hooks: Vec::new(),
            max_turns: 10,
            memory_data: None,
        }
    }

    fn with_client(mut self, client: LlmClient) -> Self {
        self.client = Some(client);
        self
    }

    fn system_prompt(mut self, system_prompt: impl Into<String>) -> Self {
        self.system_prompt_parts.push(system_prompt.into());
        self
    }

    fn tool(mut self, tool: impl Tool + Send + Sync + 'static) -> Self {
        self.tools.push(Box::new(tool));
        self
    }

    fn add_hook(mut self, hook: impl Hook + 'static) -> Self {
        self.hooks.push(Box::new(hook));
        self
    }

    fn max_turns(mut self, n: usize) -> Self {
        self.max_turns = n;
        self
    }

    fn memory(mut self, memory: impl AgentMemory<LlmMessage> + 'static) -> Self {
        self.memory_data = Some(Box::new(memory));
        self
    }

    fn build(self) -> Agent<LlmCompletionModel, LlmMessage> {
        let model = LlmCompletionModel {
            client: self.client.expect("call .with_client() first"),
        };
        let mut agent = Agent::new(model);
        let tool_specs = self
            .tools
            .iter()
            .map(|tool| ToolPromptSpec {
                name: tool.tool_name().to_string(),
                description: tool.description(),
                parameters: tool.parameters(),
            })
            .collect::<Vec<_>>();
        let base_prompt = self.system_prompt_parts.join("\n\n");
        let system_prompt = build_tools_system_prompt(&base_prompt, &tool_specs);
        agent.tools = self.tools;
        agent.hooks = self.hooks;
        agent.default_max_turns = self.max_turns;
        agent.memory = self.memory_data;
        agent.system_prompt = Some(LlmMessage::system(&system_prompt));
        agent
    }
}

struct ToolPromptSpec {
    name: String,
    description: String,
    parameters: Value,
}

fn build_tools_system_prompt(base_prompt: &str, tool_specs: &[ToolPromptSpec]) -> String {
    let mut prompt = String::new();
    prompt.push_str(base_prompt);
    prompt.push_str("\n\nYou have access to the following tools.\n\n");

    if tool_specs.is_empty() {
        prompt.push_str("(No tools available.)\n");
    } else {
        prompt.push_str("## JSON Response Format\n");
        prompt.push_str("respond with valid JSON only (no markdown fences):\n");
        prompt.push_str("{\n");
        prompt.push_str("  \"reply_text\": \"your message to the user\",\n");
        prompt.push_str("  \"tool_calls\": [\n    {\n");
        prompt.push_str("      \"tool_name\": \"tool_name_here\",\n");
        prompt.push_str("      \"args\": { \"argument_name\": \"argument_value\" }\n");
        prompt.push_str("    }\n  ]\n");
        prompt.push_str("}\n\n");
        prompt.push_str("If you don't need to use a tool, respond with:\n");
        prompt.push_str("{\n");
        prompt.push_str("  \"reply_text\": \"your message\",\n");
        prompt.push_str("  \"tool_calls\": []\n");
        prompt.push_str("}\n\n");
        prompt.push_str("## Available Tools\n\n");

        for tool in tool_specs {
            prompt.push_str(&format!("- {}\n", tool.name));
            prompt.push_str(&format!("  Description: {}\n", tool.description));
            prompt.push_str("  Arguments JSON Schema:\n");
            let parameters = serde_json::to_string_pretty(&tool.parameters)
                .unwrap_or_else(|_| tool.parameters.to_string());
            prompt.push_str(&parameters);
            prompt.push_str("\n");
        }

        prompt.push_str(
      "as you see, tool_calls is an array, several tools calls can be executed in one turn concurrently. ",
    );
        prompt.push_str("your valid tool call will be recorded as system message");
        prompt.push_str(
          "some tools calls needs user's confirmation, you should not retry the same tool call in this case.\
           And if all tool calls are pending confirmation, you should just end the turn, as you cant do more before user's action"
        );
    }

    prompt
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn editor_agent_llm_request_keeps_gpt5_contract() {
        let request = build_editor_agent_llm_request(vec![
            LlmMessage::system("系统提示"),
            LlmMessage::user("用户请求"),
        ]);

        assert_eq!(request.model.as_deref(), Some(EDITOR_AGENT_GPT5_MODEL));
        assert_eq!(request.max_tokens, Some(1024));
        assert_eq!(request.request_timeout_ms, Some(480_000));
        assert_eq!(request.messages.len(), 2);
    }
}