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_agent_harness::prompt::{ToolPromptSpec, build_tools_system_prompt};
use platform_llm::{EDITOR_AGENT_GPT5_MODEL, LlmClient, LlmMessage, LlmRunRequest};

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
            .run(request)
            .await
            .map_err(|e| PromptError::CompletionError(e.to_string()))?;
        Ok(response.text)
    }

    fn build_system_message(&self, text: &str) -> LlmMessage {
        LlmMessage::system(text)
    }

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

    fn build_error_message(&self, error: &PromptError) -> LlmMessage {
        LlmMessage::system(format!(
            "ERROR {}",
            error.display_with_agent_label("美术 Agent")
        ))
    }
}

fn build_editor_agent_llm_request(messages: Vec<LlmMessage>) -> LlmRunRequest {
    LlmRunRequest::new(messages)
        .with_model(EDITOR_AGENT_GPT5_MODEL)
        .with_openai_chat()
        .with_max_output_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
    }
}

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

    #[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.api_kind, LlmApiKind::OpenAiChat);
        assert_eq!(request.max_output_tokens, Some(1024));
        assert_eq!(request.request_timeout_ms, Some(480_000));
        assert_eq!(request.messages.len(), 2);
    }
}