swarm-engine-llm 0.1.6

//! JSON出力専用プロンプトテンプレート
//!
//! LLMにJSON形式で応答させるための共通テンプレート。
//!
//! # 設計
//!
//! Few-shot example + 「Your JSON:」で終わるパターンを統一的に提供。
//!
//! ```text
//! ┌─────────────────────────────────────────┐
//! │ You are a JSON-only response AI [role]. │  ← Header
//! │                                         │
//! │ Example interaction:                    │  ← Few-shot
//! │ User: [example_input]                   │
//! │ Assistant: [example_json]               │
//! │                                         │
//! │ Now respond with ONLY a JSON object:    │  ← Transition
//! │ [content]                               │  ← Content
//! │ Your JSON:                              │  ← Footer
//! └─────────────────────────────────────────┘
//! ```
//!
//! # 使用例
//!
//! ```ignore
//! let template = JsonPromptTemplate::new()
//!     .with_role("for action selection")
//!     .with_example(
//!         "Task: read a file. Available: Read, Write.",
//!         r#"{"tool": "Read", "target": "file.txt"}"#,
//!     );
//!
//! let prompt = template.build("Task: find errors. Available: Grep, Read.");
//! ```

/// JSON出力専用プロンプトテンプレート
#[derive(Debug, Clone)]
pub struct JsonPromptTemplate {
    /// 役割説明（例: "for exploration strategy selection"）
    role: Option<String>,
    /// Few-shot example の入力側
    example_input: String,
    /// Few-shot example の出力側（JSON）
    example_json: String,
}

impl JsonPromptTemplate {
    /// 新しいテンプレートを作成
    pub fn new() -> Self {
        Self {
            role: None,
            example_input: String::new(),
            example_json: String::new(),
        }
    }

    /// 役割説明を設定
    pub fn with_role(mut self, role: impl Into<String>) -> Self {
        self.role = Some(role.into());
        self
    }

    /// Few-shot example を設定
    pub fn with_example(mut self, input: impl Into<String>, json: impl Into<String>) -> Self {
        self.example_input = input.into();
        self.example_json = json.into();
        self
    }

    /// プロンプトを生成
    ///
    /// # 引数
    ///
    /// - `content`: コンテキスト/タスク情報（複数行可）
    pub fn build(&self, content: &str) -> String {
        let role_part = self
            .role
            .as_ref()
            .map(|r| format!(" {}", r))
            .unwrap_or_default();

        format!(
            "You are a JSON-only response AI{role}.

Example interaction:
User: {input}
Assistant: {json}

Now respond with ONLY a JSON object:
{content}
Your JSON:",
            role = role_part,
            input = self.example_input,
            json = self.example_json,
            content = content,
        )
    }

    /// プロンプトを生成（User: プレフィックス付き）
    ///
    /// contentの各行に「User:」を付けたい場合に使用。
    pub fn build_with_user_prefix(&self, content: &str) -> String {
        let role_part = self
            .role
            .as_ref()
            .map(|r| format!(" {}", r))
            .unwrap_or_default();

        format!(
            "You are a JSON-only response AI{role}.

Example interaction:
User: {input}
Assistant: {json}

Now respond with ONLY a JSON object:
User: {content}
Your JSON:",
            role = role_part,
            input = self.example_input,
            json = self.example_json,
            content = content,
        )
    }
}

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

// ============================================================================
// Pre-defined Templates
// ============================================================================

/// Action選択用テンプレート
pub fn action_selection_template() -> JsonPromptTemplate {
    JsonPromptTemplate::new().with_example(
        "Task is to read a file. Available: Read, Write.",
        r#"{"tool": "Read", "target": "file.txt", "args": {}, "confidence": 0.9}"#,
    )
}

/// Strategy選択用テンプレート
pub fn strategy_selection_template() -> JsonPromptTemplate {
    JsonPromptTemplate::new()
        .with_role("for exploration strategy selection")
        .with_example(
            "frontier=5, visits=10, failure=5%, current=FIFO",
            r#"{"strategy":"UCB1","change":true,"confidence":0.8,"reason":"early stage needs exploration"}"#,
        )
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_basic_template() {
        let template =
            JsonPromptTemplate::new().with_example("input example", r#"{"result": "output"}"#);

        let prompt = template.build("actual content");

        assert!(prompt.contains("JSON-only response AI"));
        assert!(prompt.contains("Example interaction:"));
        assert!(prompt.contains("User: input example"));
        assert!(prompt.contains(r#"Assistant: {"result": "output"}"#));
        assert!(prompt.contains("actual content"));
        assert!(prompt.ends_with("Your JSON:"));
    }

    #[test]
    fn test_template_with_role() {
        let template = JsonPromptTemplate::new()
            .with_role("for testing")
            .with_example("test input", "{}");

        let prompt = template.build("content");

        assert!(prompt.contains("JSON-only response AI for testing"));
    }

    #[test]
    fn test_build_with_user_prefix() {
        let template = JsonPromptTemplate::new().with_example("ex", "{}");

        let prompt = template.build_with_user_prefix("my content");

        assert!(prompt.contains("User: my content"));
        assert!(prompt.ends_with("Your JSON:"));
    }

    #[test]
    fn test_action_selection_template() {
        let template = action_selection_template();
        let prompt = template.build("Find the bug");

        assert!(prompt.contains("Read, Write"));
        assert!(prompt.contains("tool"));
        assert!(prompt.contains("Find the bug"));
    }

    #[test]
    fn test_strategy_selection_template() {
        let template = strategy_selection_template();
        let prompt = template.build_with_user_prefix("frontier=10, visits=50");

        assert!(prompt.contains("exploration strategy selection"));
        assert!(prompt.contains("UCB1"));
        assert!(prompt.contains("User: frontier=10"));
    }
}