zclaw_openfang/crates/zclaw-growth/src/injector.rs

//! Prompt Injector - Injects retrieved memories into system prompts
//!
//! This module provides the `PromptInjector` which formats and injects
//! retrieved memories into the agent's system prompt for context enhancement.
//!
//! # Formatting Options
//!
//! - `inject()` - Standard markdown format with sections
//! - `inject_compact()` - Compact format for limited token budgets
//! - `inject_json()` - JSON format for structured processing
//! - `inject_custom()` - Custom template with placeholders

use crate::types::{MemoryEntry, RetrievalConfig, RetrievalResult};

/// Output format for memory injection
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum InjectionFormat {
    /// Standard markdown with sections (default)
    Markdown,
    /// Compact inline format
    Compact,
    /// JSON structured format
    Json,
}

/// Prompt Injector - injects memories into system prompts
pub struct PromptInjector {
    /// Retrieval configuration for token budgets
    config: RetrievalConfig,
    /// Output format
    format: InjectionFormat,
    /// Custom template (uses {{preferences}}, {{knowledge}}, {{experience}} placeholders)
    custom_template: Option<String>,
}

impl Default for PromptInjector {
    fn default() -> Self {
        Self::new()
    }
}

impl PromptInjector {
    /// Create a new prompt injector
    pub fn new() -> Self {
        Self {
            config: RetrievalConfig::default(),
            format: InjectionFormat::Markdown,
            custom_template: None,
        }
    }

    /// Create with custom configuration
    pub fn with_config(config: RetrievalConfig) -> Self {
        Self {
            config,
            format: InjectionFormat::Markdown,
            custom_template: None,
        }
    }

    /// Set the output format
    pub fn with_format(mut self, format: InjectionFormat) -> Self {
        self.format = format;
        self
    }

    /// Set a custom template for injection
    ///
    /// Template placeholders:
    /// - `{{preferences}}` - Formatted preferences section
    /// - `{{knowledge}}` - Formatted knowledge section
    /// - `{{experience}}` - Formatted experience section
    /// - `{{all}}` - All memories combined
    pub fn with_custom_template(mut self, template: impl Into<String>) -> Self {
        self.custom_template = Some(template.into());
        self
    }

    /// Inject memories into a base system prompt
    ///
    /// This method constructs an enhanced system prompt by:
    /// 1. Starting with the base prompt
    /// 2. Adding a "用户偏好" section if preferences exist
    /// 3. Adding a "相关知识" section if knowledge exists
    /// 4. Adding an "经验参考" section if experience exists
    ///
    /// Each section respects the token budget configuration.
    pub fn inject(&self, base_prompt: &str, memories: &RetrievalResult) -> String {
        // If no memories, return base prompt unchanged
        if memories.is_empty() {
            return base_prompt.to_string();
        }

        let mut result = base_prompt.to_string();

        // Inject preferences section
        if !memories.preferences.is_empty() {
            let section = self.format_section(
                "## 用户偏好",
                &memories.preferences,
                self.config.preference_budget,
                |entry| format!("- {}", entry.content),
            );
            result.push_str("\n\n");
            result.push_str(&section);
        }

        // Inject knowledge section
        if !memories.knowledge.is_empty() {
            let section = self.format_section(
                "## 相关知识",
                &memories.knowledge,
                self.config.knowledge_budget,
                |entry| format!("- {}", entry.content),
            );
            result.push_str("\n\n");
            result.push_str(&section);
        }

        // Inject experience section
        if !memories.experience.is_empty() {
            let section = self.format_section(
                "## 经验参考",
                &memories.experience,
                self.config.experience_budget,
                |entry| format!("- {}", entry.content),
            );
            result.push_str("\n\n");
            result.push_str(&section);
        }

        // Add memory context footer
        result.push_str("\n\n");
        result.push_str("<!-- 以上内容基于历史对话自动提取的记忆 -->");

        result
    }

    /// Format a section of memories with token budget
    fn format_section<F>(
        &self,
        header: &str,
        entries: &[MemoryEntry],
        token_budget: usize,
        formatter: F,
    ) -> String
    where
        F: Fn(&MemoryEntry) -> String,
    {
        let mut result = String::new();
        result.push_str(header);
        result.push('\n');

        let mut used_tokens = 0;
        let header_tokens = header.len() / 4;
        used_tokens += header_tokens;

        for entry in entries {
            let line = formatter(entry);
            let line_tokens = line.len() / 4;

            if used_tokens + line_tokens > token_budget {
                // Add truncation indicator
                result.push_str("- ... (更多内容已省略)\n");
                break;
            }

            result.push_str(&line);
            result.push('\n');
            used_tokens += line_tokens;
        }

        result
    }

    /// Build a minimal context string for token-limited scenarios
    pub fn build_minimal_context(&self, memories: &RetrievalResult) -> String {
        if memories.is_empty() {
            return String::new();
        }

        let mut context = String::new();

        // Only include top preference
        if let Some(pref) = memories.preferences.first() {
            context.push_str(&format!("[偏好] {}\n", pref.content));
        }

        // Only include top knowledge
        if let Some(knowledge) = memories.knowledge.first() {
            context.push_str(&format!("[知识] {}\n", knowledge.content));
        }

        context
    }

    /// Inject memories in compact format
    ///
    /// Compact format uses inline notation: [P] for preferences, [K] for knowledge, [E] for experience
    pub fn inject_compact(&self, base_prompt: &str, memories: &RetrievalResult) -> String {
        if memories.is_empty() {
            return base_prompt.to_string();
        }

        let mut result = base_prompt.to_string();
        let mut context_parts = Vec::new();

        // Add compact preferences
        for entry in &memories.preferences {
            context_parts.push(format!("[P] {}", entry.content));
        }

        // Add compact knowledge
        for entry in &memories.knowledge {
            context_parts.push(format!("[K] {}", entry.content));
        }

        // Add compact experience
        for entry in &memories.experience {
            context_parts.push(format!("[E] {}", entry.content));
        }

        if !context_parts.is_empty() {
            result.push_str("\n\n[记忆上下文]\n");
            result.push_str(&context_parts.join("\n"));
        }

        result
    }

    /// Inject memories as JSON structure
    ///
    /// Returns a JSON object with preferences, knowledge, and experience arrays
    pub fn inject_json(&self, base_prompt: &str, memories: &RetrievalResult) -> String {
        if memories.is_empty() {
            return base_prompt.to_string();
        }

        let preferences: Vec<_> = memories.preferences.iter()
            .map(|e| serde_json::json!({
                "content": e.content,
                "importance": e.importance,
                "keywords": e.keywords,
            }))
            .collect();

        let knowledge: Vec<_> = memories.knowledge.iter()
            .map(|e| serde_json::json!({
                "content": e.content,
                "importance": e.importance,
                "keywords": e.keywords,
            }))
            .collect();

        let experience: Vec<_> = memories.experience.iter()
            .map(|e| serde_json::json!({
                "content": e.content,
                "importance": e.importance,
                "keywords": e.keywords,
            }))
            .collect();

        let memories_json = serde_json::json!({
            "preferences": preferences,
            "knowledge": knowledge,
            "experience": experience,
        });

        format!("{}\n\n[记忆上下文]\n{}", base_prompt, serde_json::to_string_pretty(&memories_json).unwrap_or_default())
    }

    /// Inject using custom template
    ///
    /// Template placeholders:
    /// - `{{preferences}}` - Formatted preferences section
    /// - `{{knowledge}}` - Formatted knowledge section
    /// - `{{experience}}` - Formatted experience section
    /// - `{{all}}` - All memories combined
    pub fn inject_custom(&self, template: &str, memories: &RetrievalResult) -> String {
        let mut result = template.to_string();

        // Format each section
        let prefs = if !memories.preferences.is_empty() {
            memories.preferences.iter()
                .map(|e| format!("- {}", e.content))
                .collect::<Vec<_>>()
                .join("\n")
        } else {
            String::new()
        };

        let knowledge = if !memories.knowledge.is_empty() {
            memories.knowledge.iter()
                .map(|e| format!("- {}", e.content))
                .collect::<Vec<_>>()
                .join("\n")
        } else {
            String::new()
        };

        let experience = if !memories.experience.is_empty() {
            memories.experience.iter()
                .map(|e| format!("- {}", e.content))
                .collect::<Vec<_>>()
                .join("\n")
        } else {
            String::new()
        };

        // Combine all
        let all = format!(
            "用户偏好:\n{}\n\n相关知识:\n{}\n\n经验参考:\n{}",
            if prefs.is_empty() { "无" } else { &prefs },
            if knowledge.is_empty() { "无" } else { &knowledge },
            if experience.is_empty() { "无" } else { &experience },
        );

        // Replace placeholders
        result = result.replace("{{preferences}}", &prefs);
        result = result.replace("{{knowledge}}", &knowledge);
        result = result.replace("{{experience}}", &experience);
        result = result.replace("{{all}}", &all);

        result
    }

    /// Inject memories using the configured format
    pub fn inject_with_format(&self, base_prompt: &str, memories: &RetrievalResult) -> String {
        match self.format {
            InjectionFormat::Markdown => self.inject(base_prompt, memories),
            InjectionFormat::Compact => self.inject_compact(base_prompt, memories),
            InjectionFormat::Json => self.inject_json(base_prompt, memories),
        }
    }

    /// Estimate total tokens that will be injected
    pub fn estimate_injection_tokens(&self, memories: &RetrievalResult) -> usize {
        let mut total = 0;

        // Count preference tokens
        for entry in &memories.preferences {
            total += entry.estimated_tokens();
            if total > self.config.preference_budget {
                total = self.config.preference_budget;
                break;
            }
        }

        // Count knowledge tokens
        let mut knowledge_tokens = 0;
        for entry in &memories.knowledge {
            knowledge_tokens += entry.estimated_tokens();
            if knowledge_tokens > self.config.knowledge_budget {
                knowledge_tokens = self.config.knowledge_budget;
                break;
            }
        }
        total += knowledge_tokens;

        // Count experience tokens
        let mut experience_tokens = 0;
        for entry in &memories.experience {
            experience_tokens += entry.estimated_tokens();
            if experience_tokens > self.config.experience_budget {
                experience_tokens = self.config.experience_budget;
                break;
            }
        }
        total += experience_tokens;

        total
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::MemoryType;
    use chrono::Utc;

    fn create_test_entry(content: &str) -> MemoryEntry {
        MemoryEntry {
            uri: "test://uri".to_string(),
            memory_type: MemoryType::Preference,
            content: content.to_string(),
            keywords: vec![],
            importance: 5,
            access_count: 0,
            created_at: Utc::now(),
            last_accessed: Utc::now(),
            overview: None,
            abstract_summary: None,
        }
    }

    #[test]
    fn test_injector_empty_memories() {
        let injector = PromptInjector::new();
        let base = "You are a helpful assistant.";
        let memories = RetrievalResult::default();

        let result = injector.inject(base, &memories);
        assert_eq!(result, base);
    }

    #[test]
    fn test_injector_with_preferences() {
        let injector = PromptInjector::new();
        let base = "You are a helpful assistant.";
        let memories = RetrievalResult {
            preferences: vec![create_test_entry("User prefers concise responses")],
            knowledge: vec![],
            experience: vec![],
            total_tokens: 0,
        };

        let result = injector.inject(base, &memories);
        assert!(result.contains("用户偏好"));
        assert!(result.contains("User prefers concise responses"));
    }

    #[test]
    fn test_injector_with_all_types() {
        let injector = PromptInjector::new();
        let base = "You are a helpful assistant.";

        let memories = RetrievalResult {
            preferences: vec![create_test_entry("Prefers concise")],
            knowledge: vec![create_test_entry("Knows Rust")],
            experience: vec![create_test_entry("Browser skill works well")],
            total_tokens: 0,
        };

        let result = injector.inject(base, &memories);
        assert!(result.contains("用户偏好"));
        assert!(result.contains("相关知识"));
        assert!(result.contains("经验参考"));
    }

    #[test]
    fn test_minimal_context() {
        let injector = PromptInjector::new();
        let memories = RetrievalResult {
            preferences: vec![create_test_entry("Prefers concise")],
            knowledge: vec![create_test_entry("Knows Rust")],
            experience: vec![],
            total_tokens: 0,
        };

        let context = injector.build_minimal_context(&memories);
        assert!(context.contains("[偏好]"));
        assert!(context.contains("[知识]"));
    }

    #[test]
    fn test_estimate_tokens() {
        let injector = PromptInjector::new();
        let memories = RetrievalResult {
            preferences: vec![create_test_entry("Short text")],
            knowledge: vec![],
            experience: vec![],
            total_tokens: 0,
        };

        let estimate = injector.estimate_injection_tokens(&memories);
        assert!(estimate > 0);
    }

    #[test]
    fn test_inject_compact() {
        let injector = PromptInjector::new();
        let base = "You are a helpful assistant.";
        let memories = RetrievalResult {
            preferences: vec![create_test_entry("Prefers concise")],
            knowledge: vec![create_test_entry("Knows Rust")],
            experience: vec![],
            total_tokens: 0,
        };

        let result = injector.inject_compact(base, &memories);
        assert!(result.contains("[P]"));
        assert!(result.contains("[K]"));
        assert!(result.contains("[记忆上下文]"));
    }

    #[test]
    fn test_inject_json() {
        let injector = PromptInjector::new();
        let base = "You are a helpful assistant.";
        let memories = RetrievalResult {
            preferences: vec![create_test_entry("Prefers concise")],
            knowledge: vec![],
            experience: vec![],
            total_tokens: 0,
        };

        let result = injector.inject_json(base, &memories);
        assert!(result.contains("\"preferences\""));
        assert!(result.contains("Prefers concise"));
    }

    #[test]
    fn test_inject_custom() {
        let injector = PromptInjector::new();
        let template = "Context:\n{{all}}";
        let memories = RetrievalResult {
            preferences: vec![create_test_entry("Prefers concise")],
            knowledge: vec![create_test_entry("Knows Rust")],
            experience: vec![],
            total_tokens: 0,
        };

        let result = injector.inject_custom(template, &memories);
        assert!(result.contains("用户偏好"));
        assert!(result.contains("相关知识"));
    }

    #[test]
    fn test_format_selection() {
        let base = "Base";

        let memories = RetrievalResult {
            preferences: vec![create_test_entry("Test")],
            knowledge: vec![],
            experience: vec![],
            total_tokens: 0,
        };

        // Test markdown format
        let injector_md = PromptInjector::new().with_format(InjectionFormat::Markdown);
        let result_md = injector_md.inject_with_format(base, &memories);
        assert!(result_md.contains("## 用户偏好"));

        // Test compact format
        let injector_compact = PromptInjector::new().with_format(InjectionFormat::Compact);
        let result_compact = injector_compact.inject_with_format(base, &memories);
        assert!(result_compact.contains("[P]"));
    }
}