swarm-engine-core 0.1.6

//! DPO (Direct Preference Optimization) LearnModel
//!
//! group_id でグループ化された Episode を比較し、DPO 学習用データを生成する。

use std::collections::HashMap;

use super::super::episode::{Episode, EpisodeContext, Outcome};
use super::super::record::Record;
use super::super::training::TrainingData;
use super::{LearnError, LearnModel};
use crate::types::GroupId;

/// DPO 学習用の比較ペア
///
/// 同じ group_id 内の成功/失敗 Episode から生成される。
#[derive(Debug, Clone)]
pub struct DpoPair {
    /// 成功した Episode
    pub chosen: Episode,
    /// 失敗した Episode
    pub rejected: Episode,
    /// 共通の group_id
    pub group_id: GroupId,
    /// 品質差（chosen.score - rejected.score）
    pub quality_gap: f64,
}

impl DpoPair {
    /// 新しい DpoPair を作成
    pub fn new(chosen: Episode, rejected: Episode, group_id: GroupId) -> Self {
        let chosen_score = chosen.outcome.score();
        let rejected_score = rejected.outcome.score();
        let quality_gap = chosen_score - rejected_score;

        Self {
            chosen,
            rejected,
            group_id,
            quality_gap,
        }
    }
}

/// DPO LearnModel の設定
#[derive(Debug, Clone)]
pub struct DpoConfig {
    /// 最小品質差（この差未満のペアは除外）
    pub min_quality_gap: f64,
    /// 最大ペア数（None なら無制限）
    pub max_pairs: Option<usize>,
    /// 同じエピソードの重複使用を許可
    pub allow_reuse: bool,
}

impl Default for DpoConfig {
    fn default() -> Self {
        Self {
            min_quality_gap: 0.1, // 10% 以上の差
            max_pairs: None,
            allow_reuse: true,
        }
    }
}

/// 汎用 DPO LearnModel
///
/// group_id でグループ化された Episode を比較し、DPO 学習用データを生成する。
///
/// ## 設計思想
///
/// DPO 学習では「同じ条件で複数回実行した結果を比較」する。
/// - group_id: 同じ条件での実行グループ（Eval -n 5 で 5 回実行など）
/// - 成功 Episode と失敗 Episode をペアにして比較
///
/// ## 使用方法
///
/// ```ignore
/// // Eval で group_id 付きの Episode を収集
/// let episodes: Vec<Episode> = ...;
///
/// // DPO ペアを生成
/// let dpo_learn = DpoLearnModel::new();
/// let pairs = dpo_learn.build_pairs(&episodes);
///
/// // TrainingData に変換
/// let training_data: Vec<TrainingData> = pairs
///     .iter()
///     .filter_map(|pair| dpo_learn.convert_pair(pair).ok())
///     .collect();
/// ```
pub struct DpoLearnModel<F>
where
    F: Fn(&Episode) -> Option<(String, String)> + Send + Sync,
{
    /// システムプロンプト
    system_prompt: String,
    /// 設定
    config: DpoConfig,
    /// Episode から (prompt, response) を抽出する関数
    extractor: F,
}

impl<F> DpoLearnModel<F>
where
    F: Fn(&Episode) -> Option<(String, String)> + Send + Sync,
{
    /// 新しい DpoLearnModel を作成
    pub fn new(extractor: F) -> Self {
        Self {
            system_prompt: String::new(),
            config: DpoConfig::default(),
            extractor,
        }
    }

    /// システムプロンプトを設定
    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
        self.system_prompt = prompt.into();
        self
    }

    /// 設定を適用
    pub fn with_config(mut self, config: DpoConfig) -> Self {
        self.config = config;
        self
    }

    /// 最小品質差を設定
    pub fn with_min_quality_gap(mut self, gap: f64) -> Self {
        self.config.min_quality_gap = gap;
        self
    }

    /// 最大ペア数を設定
    pub fn with_max_pairs(mut self, max: usize) -> Self {
        self.config.max_pairs = Some(max);
        self
    }

    /// group_id でグループ化された Episode から DPO ペアを生成
    pub fn build_pairs(&self, episodes: &[Episode]) -> Vec<DpoPair> {
        // group_id でグループ化
        let mut by_group: HashMap<GroupId, Vec<&Episode>> = HashMap::new();
        for ep in episodes {
            if let Some(gid) = ep.group_id {
                by_group.entry(gid).or_default().push(ep);
            }
        }

        let mut pairs = Vec::new();

        for (group_id, group_episodes) in by_group {
            // 成功/失敗で分類
            let (successes, failures): (Vec<_>, Vec<_>) = group_episodes
                .into_iter()
                .partition(|ep| ep.outcome.is_success());

            if successes.is_empty() || failures.is_empty() {
                continue;
            }

            // スコアでソート（高い順）
            let mut sorted_successes: Vec<_> = successes;
            sorted_successes.sort_by(|a, b| {
                let a_score = a.outcome.score();
                let b_score = b.outcome.score();
                b_score
                    .partial_cmp(&a_score)
                    .unwrap_or(std::cmp::Ordering::Equal)
            });

            // スコアでソート（低い順）
            let mut sorted_failures: Vec<_> = failures;
            sorted_failures.sort_by(|a, b| {
                let a_score = a.outcome.score();
                let b_score = b.outcome.score();
                a_score
                    .partial_cmp(&b_score)
                    .unwrap_or(std::cmp::Ordering::Equal)
            });

            // ペア作成
            for success_ep in &sorted_successes {
                for failure_ep in &sorted_failures {
                    let chosen_score = success_ep.outcome.score();
                    let rejected_score = failure_ep.outcome.score();
                    let gap = chosen_score - rejected_score;

                    if gap < self.config.min_quality_gap {
                        continue;
                    }

                    let pair = DpoPair::new((*success_ep).clone(), (*failure_ep).clone(), group_id);
                    pairs.push(pair);

                    if !self.config.allow_reuse {
                        break;
                    }
                }

                if !self.config.allow_reuse {
                    break;
                }
            }
        }

        // 品質差でソート（大きい順）
        pairs.sort_by(|a, b| {
            b.quality_gap
                .partial_cmp(&a.quality_gap)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // 最大数で制限
        if let Some(max) = self.config.max_pairs {
            pairs.truncate(max);
        }

        pairs
    }

    /// DPO ペアを TrainingData に変換
    pub fn convert_pair(&self, pair: &DpoPair) -> Result<TrainingData, LearnError> {
        let (chosen_prompt, chosen_response) = (self.extractor)(&pair.chosen)
            .ok_or_else(|| LearnError::MissingData("chosen prompt/response".into()))?;

        let (rejected_prompt, rejected_response) = (self.extractor)(&pair.rejected)
            .ok_or_else(|| LearnError::MissingData("rejected prompt/response".into()))?;

        // prompt が一致することを確認（正規化後）
        if chosen_prompt != rejected_prompt {
            return Err(LearnError::InvalidEpisode(format!(
                "Prompt mismatch: '{}' vs '{}'",
                chosen_prompt, rejected_prompt
            )));
        }

        let training = if self.system_prompt.is_empty() {
            TrainingData::dpo(&chosen_prompt, &chosen_response, &rejected_response)
        } else {
            TrainingData::dpo_with_system(
                &self.system_prompt,
                &chosen_prompt,
                &chosen_response,
                &rejected_response,
            )
        };

        Ok(training
            .with_episode_id(pair.chosen.id.to_string())
            .with_custom("rejected_episode_id", pair.rejected.id.to_string())
            .with_custom("quality_gap", pair.quality_gap.to_string())
            .with_custom("group_id", pair.group_id.0.to_string()))
    }

    /// 複数のペアを一括変換
    pub fn convert_pairs(&self, pairs: &[DpoPair]) -> Vec<TrainingData> {
        pairs
            .iter()
            .filter_map(|pair| self.convert_pair(pair).ok())
            .collect()
    }
}

/// LearnModel trait の実装（Record ベースの Episode 構築用）
///
/// DPO は通常、既存の Episode を比較するため、build_episodes は空を返す。
/// 実際の DPO ペア生成は build_pairs メソッドを使用。
impl<F> LearnModel for DpoLearnModel<F>
where
    F: Fn(&Episode) -> Option<(String, String)> + Send + Sync,
{
    fn name(&self) -> &str {
        "dpo"
    }

    fn objective(&self) -> &str {
        "Learn preferences from success/failure Episode pairs within the same group"
    }

    fn build_episodes(&self, _records: &[Record]) -> Vec<Episode> {
        // DPO は既存の Episode を比較するため、Record から Episode は生成しない
        vec![]
    }

    fn evaluate(&self, _context: &EpisodeContext) -> Outcome {
        // DpoLearnModel は複数 Episode を group_id でグルーピングし、
        // 成功/失敗のペアを比較して学習する。
        // 個々の Episode を evaluate() するのは設計として不適切。
        //
        // DPO のフロー:
        //   1. Eval 実行時に Episode が生成される（Outcome は Eval 側で設定）
        //   2. build_pairs() で group_id ごとにグルーピング
        //   3. 成功/失敗 Episode のペアから TrainingData を生成
        //
        // この evaluate() が呼ばれるのは実装ミス。
        panic!(
            "DpoLearnModel::evaluate() should not be called.\n\
             DPO learning compares multiple Episodes by group_id, not individual Episode evaluation.\n\
             Use build_pairs() to generate training pairs from Episodes."
        );
    }

    fn convert(&self, _episode: &Episode) -> Result<TrainingData, LearnError> {
        // 単一の Episode からは DPO TrainingData は生成できない
        // convert_pair を使用すること
        Err(LearnError::InvalidEpisode(
            "DPO requires pairs, use convert_pair instead".into(),
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::learn::episode::EpisodeBuilder;
    use crate::learn::record::ActionRecord;
    use crate::types::TaskId;

    fn create_test_episode(
        task_id: TaskId,
        group_id: GroupId,
        success: bool,
        score: f64,
    ) -> Episode {
        let outcome = if success {
            Outcome::success(score)
        } else {
            Outcome::failure("test failure")
        };

        EpisodeBuilder::default()
            .learn_model("test")
            .task_id(task_id)
            .group_id(group_id)
            .record(ActionRecord::new(1, 0, "TestAction").success(success))
            .outcome(outcome)
            .build()
    }

    fn test_extractor(ep: &Episode) -> Option<(String, String)> {
        // テスト用: 固定の prompt/response を返す
        Some((
            "test prompt".to_string(),
            format!("response for {:?}", ep.id),
        ))
    }

    #[test]
    fn test_build_pairs_basic() {
        let group_id = GroupId::new();
        let task1 = TaskId::new();
        let task2 = TaskId::new();

        let episodes = vec![
            create_test_episode(task1, group_id, true, 0.9),
            create_test_episode(task2, group_id, false, 0.0),
        ];

        let dpo = DpoLearnModel::new(test_extractor);
        let pairs = dpo.build_pairs(&episodes);

        assert_eq!(pairs.len(), 1);
        assert!(pairs[0].quality_gap > 0.0);
    }

    #[test]
    fn test_build_pairs_different_groups() {
        let group1 = GroupId::new();
        let group2 = GroupId::new();

        let episodes = vec![
            create_test_episode(TaskId::new(), group1, true, 0.9),
            create_test_episode(TaskId::new(), group2, false, 0.0),
        ];

        let dpo = DpoLearnModel::new(test_extractor);
        let pairs = dpo.build_pairs(&episodes);

        // 異なる group_id なのでペアにならない
        assert!(pairs.is_empty());
    }

    #[test]
    fn test_min_quality_gap() {
        let group_id = GroupId::new();

        let episodes = vec![
            create_test_episode(TaskId::new(), group_id, true, 0.6),
            create_test_episode(TaskId::new(), group_id, false, 0.0),
        ];

        // 0.5 以上の差を要求
        let dpo = DpoLearnModel::new(test_extractor).with_min_quality_gap(0.5);
        let pairs = dpo.build_pairs(&episodes);

        // 0.6 - 0.0 = 0.6 なのでペアになる
        assert_eq!(pairs.len(), 1);

        // 0.7 以上の差を要求
        let dpo = DpoLearnModel::new(test_extractor).with_min_quality_gap(0.7);
        let pairs = dpo.build_pairs(&episodes);

        // 差が足りないのでペアにならない
        assert!(pairs.is_empty());
    }
}