PicoBot/src/tools/task/prompt.rs

use super::types::SubagentDef;
use crate::config::LLMProviderConfig;

/// 子代理系统提示词构建器
pub struct SubagentPromptBuilder;

impl SubagentPromptBuilder {
    /// 构建子代理系统提示词（包含系统环境信息和技能索引）
    pub fn build(
        def: &SubagentDef,
        description: &str,
        prompt: &str,
        config: &LLMProviderConfig,
        skills_index: Option<&str>,
    ) -> String {
        let base_prompt = Self::interpolate_template(def, description, prompt);
        let env_info = crate::agent::generate_system_env_prompt(config);

        // 组合提示词：基础 + 环境 + 技能索引（可选）
        match skills_index {
            Some(index) if !index.is_empty() => {
                format!("{}\n\n{}\n\n{}", base_prompt, env_info, index)
            }
            _ => format!("{}\n\n{}", base_prompt, env_info),
        }
    }

    /// 构建恢复任务的提示词
    pub fn build_resume_prompt(session_description: &str, additional_prompt: &str) -> String {
        format!(
            "你正在继续执行一个之前创建的子代理任务。\n\n\
            任务描述: {}\n\n\
            继续执行指令: {}\n\n\
            你应该:\n\
            1. 回顾之前的工作进度（如果已有历史）\n\
            2. 继续完成任务，不要偏离目标\n\
            3. 完成后给出简洁的总结\n\
            4. 不要尝试创建新的子代理任务\n\n\
            注意: 你在一个独立的执行上下文中，没有访问主对话历史的权限。",
            session_description, additional_prompt
        )
    }

    /// 插值提示词模板
    fn interpolate_template(def: &SubagentDef, description: &str, prompt: &str) -> String {
        // 自定义子代理使用通用模板
        let base = if def.prompt_template.is_empty() {
            "你是一个专注的子代理，正在执行一个独立任务。\n\n\
             任务描述: {{description}}\n\n\
             你应该:\n\
             1. 专注于完成任务，不要偏离目标\n\
             2. 使用可用的工具进行必要操作\n\
             3. 完成后给出简洁的总结\n\
             4. 当任务复杂度较高时，可以使用 `task` 工具创建子代理来处理独立子任务\n\n\
             任务追踪:\n\
             你可以使用 `todo_write` 工具追踪子任务进度。规则：同一时间只有一个 in_progress，完成后再标记下一个，3步以上才使用。\n\n\
             注意: 你没有访问主对话历史的权限，这是一个独立的执行上下文。"
        } else {
            &def.prompt_template
        };

        let mut result = base
            .replace("{{description}}", description)
            .replace("{{prompt}}", prompt);

        if let Some(ref body) = def.body {
            result.push_str("\n\n");
            result.push_str(body);
        }

        result
    }
}

/// 从子代理输出提取简洁摘要
pub fn extract_summary(content: &str) -> String {
    // 取第一段或前 500 字符
    let first_paragraph = content
        .lines()
        .take_while(|line| !line.trim().is_empty())
        .collect::<Vec<_>>()
        .join("\n");

    if first_paragraph.len() > 500 {
        first_paragraph.chars().take(500).collect()
    } else if first_paragraph.is_empty() {
        content.chars().take(200).collect()
    } else {
        first_paragraph
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tools::task::types::SubagentSource;

    fn test_def() -> SubagentDef {
        SubagentDef {
            name: "general".to_string(),
            description: "测试".to_string(),
            prompt_template: "任务: {{description}}\n指令: {{prompt}}".to_string(),
            body: None,
            capability: crate::domain::CapabilityPolicy::default(),
            max_execution_secs: None,
            source: SubagentSource::Builtin,
            path: None,
            provider: None,
            model: None,
        }
    }

    #[test]
    fn test_interpolates_template() {
        let def = test_def();
        let result = SubagentPromptBuilder::build(
            &def,
            "审查代码",
            "检查安全漏洞",
            &LLMProviderConfig {
                provider_type: "openai".to_string(),
                name: "test".to_string(),
                base_url: "http://localhost".to_string(),
                api_key: "test".to_string(),
                extra_headers: std::collections::HashMap::new(),
                llm_timeout_secs: 120,
                memory_maintenance_timeout_secs: 600,
                max_retries: 3,
                model_id: "test".to_string(),
                temperature: None,
                max_tokens: None,
                context_window_tokens: None,
                model_extra: std::collections::HashMap::new(),
                max_tool_iterations: 1,
                tool_result_max_chars: 1000,
                context_tool_result_trim_chars: 1000,
                max_run_secs: crate::agent::DEFAULT_MAX_RUN_SECS,
                max_images_in_context: 1,
                max_image_age_rounds: 10,
            },
            None,
        );
        assert!(result.contains("任务: 审查代码"));
        assert!(result.contains("指令: 检查安全漏洞"));
    }

    #[test]
    fn test_appends_body() {
        let mut def = test_def();
        def.body = Some("额外指令".to_string());
        let result = SubagentPromptBuilder::build(
            &def,
            "描述",
            "指令",
            &LLMProviderConfig {
                provider_type: "openai".to_string(),
                name: "test".to_string(),
                base_url: "http://localhost".to_string(),
                api_key: "test".to_string(),
                extra_headers: std::collections::HashMap::new(),
                llm_timeout_secs: 120,
                memory_maintenance_timeout_secs: 600,
                max_retries: 3,
                model_id: "test".to_string(),
                temperature: None,
                max_tokens: None,
                context_window_tokens: None,
                model_extra: std::collections::HashMap::new(),
                max_tool_iterations: 1,
                tool_result_max_chars: 1000,
                context_tool_result_trim_chars: 1000,
                max_run_secs: crate::agent::DEFAULT_MAX_RUN_SECS,
                max_images_in_context: 1,
                max_image_age_rounds: 10,
            },
            None,
        );
        assert!(result.contains("额外指令"));
    }
}