swarm_engine_core/learn/

learn_model.rs

//! LearnModel - 学習の統合モデル
//!
//! ## 設計思想
//!
//! LearnModel は「何を学習するか」を統合的に定義する。
//!
//! - **何を目的とするか** (objective)
//! - **何を Episode として切り出すか** (build_episodes)
//! - **何を Success/Failure とするか** (evaluate)
//! - **どう TrainingData に変換するか** (convert)
//!
//! ## Learn の価値
//!
//! Core（Swarm本体）は性能制約で 3-gram までしか取れない。
//! しかし Learn は非同期/オフラインなので、5-gram や 10-gram など
//! 自由に分析できる。これが Learn モジュールの価値。
//!
//! ## Record による抽象化
//!
//! ActionEvent と LlmDebugEvent を `Record` enum で統一的に扱う。
//! LearnModel は Record のストリームから Episode を構築する。
//!
//! ```text
//! ActionEvent ──┐
//!               ├──▶ Vec<Record> ──▶ LearnModel.build_episodes()
//! LlmDebugEvent ┘                         ↓
//!                                    Vec<Episode>
//!                                         ↓
//!                                    LearnModel.convert()
//!                                         ↓
//!                                    TrainingData
//! ```

use crate::events::ActionEvent;

use super::episode::{Episode, EpisodeContext, Outcome};
use super::record::{ActionRecord, Record, RecordStream};
use super::training::TrainingData;

// ============================================================================
// System Event Constants
// ============================================================================

/// システムイベント定数
pub mod system_events {
    /// Tick 開始イベント
    pub const TICK_START: &str = "tick_start";
    /// Tick 終了イベント
    pub const TICK_END: &str = "tick_end";
    /// タスク完了イベント
    pub const DONE: &str = "done";

    /// デフォルトのシステムイベント一覧
    pub const DEFAULT_SYSTEM_EVENTS: &[&str] = &[TICK_START, TICK_END, DONE];
}

// ============================================================================
// LearnModel Trait
// ============================================================================

/// 学習の統合モデル
///
/// 何を学習対象とし、何を成功とするかを統合的に定義する。
/// Record[] から Episode を構築し、TrainingData に変換するまでの全責務を担う。
///
/// ## Record による統一インターフェース
///
/// ActionEvent も LlmDebugEvent も `Record` として統一的に扱う。
/// これにより：
/// - ActionEvent ベースの Learn
/// - LlmDebugEvent ベースの Learn
/// - 両方を混ぜた Learn
/// 全て同じインターフェースで実装可能。
pub trait LearnModel: Send + Sync {
    /// 名前
    fn name(&self) -> &str;

    /// 目的を表す説明
    fn objective(&self) -> &str;

    /// Record のストリームから Episode を構築
    ///
    /// N-gram、Worker単位、任意のグルーピングが可能。
    /// Core が 3-gram までしか取れなくても、Learn は 5-gram や 10-gram を
    /// 自由に構築できる。
    fn build_episodes(&self, records: &[Record]) -> Vec<Episode>;

    /// Records から Success/Failure を判定
    ///
    /// 純粋なロジック: EpisodeContext (Records) → Outcome
    /// build_episodes() 内でこれを呼んで Episode.outcome を設定する。
    fn evaluate(&self, context: &EpisodeContext) -> Outcome;

    /// Episode を TrainingData に変換
    fn convert(&self, episode: &Episode) -> Result<TrainingData, LearnError>;

    /// 複数 Episode を一括変換（デフォルト実装）
    fn convert_batch(&self, episodes: &[Episode]) -> Vec<TrainingData> {
        episodes
            .iter()
            .filter_map(|ep| self.convert(ep).ok())
            .collect()
    }

    /// 便利メソッド: ActionEvent[] から直接変換
    fn build_episodes_from_actions(&self, actions: &[ActionEvent]) -> Vec<Episode> {
        let records: Vec<Record> = actions.iter().map(Record::from).collect();
        self.build_episodes(&records)
    }
}

// ============================================================================
// LearnError
// ============================================================================

/// LearnModel のエラー型
#[derive(Debug, thiserror::Error)]
pub enum LearnError {
    #[error("Build error: {0}")]
    Build(String),

    #[error("Conversion error: {0}")]
    Conversion(String),

    #[error("Missing data: {0}")]
    MissingData(String),

    #[error("Invalid episode: {0}")]
    InvalidEpisode(String),

    #[error("{0}")]
    Other(String),
}

// ============================================================================
// 組み込み LearnModel 実装
// ============================================================================

/// Worker タスク完了ベースの LearnModel
///
/// Worker が開始から done までに実行したアクション列を Episode として構築。
pub struct WorkerTaskLearn {
    /// システムプロンプト
    system_prompt: String,
    /// 最小アクション数
    min_actions: usize,
}

impl WorkerTaskLearn {
    pub fn new() -> Self {
        Self {
            system_prompt:
                "You are an intelligent agent that diagnoses and resolves system issues."
                    .to_string(),
            min_actions: 2,
        }
    }

    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
        self.system_prompt = prompt.into();
        self
    }

    pub fn with_min_actions(mut self, min: usize) -> Self {
        self.min_actions = min;
        self
    }
}

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

impl LearnModel for WorkerTaskLearn {
    fn name(&self) -> &str {
        "worker_task"
    }

    fn objective(&self) -> &str {
        "Learn complete worker task sequences from start to done"
    }

    fn evaluate(&self, context: &EpisodeContext) -> Outcome {
        // 空のコンテキストは評価不能
        if context.is_empty() {
            return Outcome::failure("Empty context: no actions to evaluate");
        }

        // 最後のアクションが done かつ success なら成功
        let last_action = context.iter::<ActionRecord>().last();

        match last_action {
            Some(action) if action.is_terminal() => {
                if action.success {
                    Outcome::success_binary()
                } else {
                    Outcome::failure("Task failed")
                }
            }
            _ => Outcome::Unknown,
        }
    }

    fn build_episodes(&self, records: &[Record]) -> Vec<Episode> {
        use std::collections::HashMap;

        let stream = RecordStream::new(records);

        // Worker ID ごとにグルーピング（Action Records のみ）
        let mut worker_actions: HashMap<usize, Vec<&ActionRecord>> = HashMap::new();
        for record in stream.actions() {
            worker_actions
                .entry(record.worker_id)
                .or_default()
                .push(record);
        }

        let mut episodes = Vec::new();

        for (_worker_id, worker_records) in worker_actions {
            // done で終わるシーケンスを探す
            let mut current_sequence: Vec<&ActionRecord> = Vec::new();

            for record in worker_records {
                current_sequence.push(record);

                if record.is_terminal() {
                    // シーケンス完了
                    if current_sequence.len() >= self.min_actions {
                        // context を構築
                        let mut context = EpisodeContext::new();
                        for r in &current_sequence {
                            context.push((*r).clone());
                        }

                        // evaluate で outcome を判定
                        let outcome = self.evaluate(&context);

                        let episode = Episode::builder()
                            .learn_model("worker_task")
                            .context(context)
                            .outcome(outcome)
                            .build();

                        episodes.push(episode);
                    }

                    current_sequence.clear();
                }
            }
        }

        episodes
    }

    fn convert(&self, episode: &Episode) -> Result<TrainingData, LearnError> {
        if !episode.outcome.is_success() {
            return Err(LearnError::InvalidEpisode(
                "Episode is not successful".into(),
            ));
        }

        let action_count = episode.context.iter::<ActionRecord>().count();
        if action_count < self.min_actions {
            return Err(LearnError::InvalidEpisode(format!(
                "Too few actions: {} < {}",
                action_count, self.min_actions
            )));
        }

        let actions: Vec<&str> = episode
            .context
            .iter::<ActionRecord>()
            .filter(|a| !a.is_terminal())
            .map(|a| a.action.as_str())
            .collect();

        let prompt = format!(
            "Diagnose and resolve the issue.\nAvailable actions: {}",
            actions.join(", ")
        );

        let response = format!("Execute the following sequence: {}", actions.join(" -> "));

        Ok(TrainingData::sft(&self.system_prompt, &prompt, &response)
            .with_episode_id(episode.id.to_string())
            .with_outcome_score(episode.outcome.score()))
    }
}

/// Worker Decision 学習モデル（シーケンスベース）
///
/// **ターゲット**: Worker の Decision（アクション選択）を改善
/// **手法**: 成功シーケンスを使った学習
///
/// ## 学習目的
///
/// Worker の Decider（アクション決定LLM）を fine-tuning するためのデータ生成。
/// 成功したアクションシーケンスから「どの順序でアクションを実行すべきか」を学習。
///
/// ## プロンプト形式
///
/// - **入力**: コンテキスト + 利用可能なアクション一覧
/// - **出力**: 最適なアクションシーケンス
///
/// ## 用途
///
/// - LoRA fine-tuning 用のデータ生成（Decider向け）
/// - 成功パターンの抽出と再利用
pub struct WorkerDecisionSequenceLearn {
    /// システムプロンプト
    system_prompt: String,
    /// 最小アクション数（これ以下は無視）
    min_actions: usize,
    /// 利用可能なアクション一覧（プロンプト生成用）
    available_actions: Vec<String>,
    /// システムイベント（フィルタ対象）
    system_events: Vec<String>,
}

impl WorkerDecisionSequenceLearn {
    pub fn new() -> Self {
        Self {
            system_prompt: "You are an intelligent agent that diagnoses and resolves system issues. \
                           Given a context and available actions, determine the optimal action sequence.".to_string(),
            min_actions: 3,
            available_actions: vec![
                "CheckStatus".to_string(),
                "ReadLogs".to_string(),
                "AnalyzeMetrics".to_string(),
                "Diagnose".to_string(),
                "Restart".to_string(),
            ],
            system_events: system_events::DEFAULT_SYSTEM_EVENTS
                .iter()
                .map(|s| s.to_string())
                .collect(),
        }
    }

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

    /// 最小アクション数を設定
    pub fn with_min_actions(mut self, min: usize) -> Self {
        self.min_actions = min;
        self
    }

    /// 利用可能なアクション一覧を設定
    pub fn with_available_actions(mut self, actions: Vec<String>) -> Self {
        self.available_actions = actions;
        self
    }

    /// システムイベント（フィルタ対象）を追加
    pub fn with_system_event(mut self, event: impl Into<String>) -> Self {
        self.system_events.push(event.into());
        self
    }

    /// アクションがシステムイベントかどうか判定
    fn is_system_event(&self, action: &str) -> bool {
        self.system_events.iter().any(|e| e == action)
    }
}

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

impl LearnModel for WorkerDecisionSequenceLearn {
    fn name(&self) -> &str {
        "worker_decision_sequence"
    }

    fn objective(&self) -> &str {
        "Learn successful action sequences for problem resolution"
    }

    fn evaluate(&self, context: &EpisodeContext) -> Outcome {
        // 空のコンテキストは評価不能
        if context.is_empty() {
            return Outcome::failure("Empty context: no actions to evaluate");
        }

        // 成功したアクション（システムイベント除く）をカウント
        let successful_actions: Vec<_> = context
            .iter::<ActionRecord>()
            .filter(|a| a.success && !self.is_system_event(&a.action))
            .collect();

        if successful_actions.len() >= self.min_actions {
            Outcome::success(1.0)
        } else {
            Outcome::failure(format!(
                "Insufficient successful actions: {} < {}",
                successful_actions.len(),
                self.min_actions
            ))
        }
    }

    fn build_episodes(&self, records: &[Record]) -> Vec<Episode> {
        // 成功したアクション（システムイベント除く）のみ抽出
        let successful_actions: Vec<&ActionRecord> = records
            .iter()
            .filter_map(Record::as_action)
            .filter(|a| a.success && !self.is_system_event(&a.action))
            .collect();

        if successful_actions.len() < self.min_actions {
            return vec![];
        }

        // 全成功アクションを1つのEpisodeとして構築
        let mut context = EpisodeContext::new();
        for action in &successful_actions {
            context.push((*action).clone());
        }

        let outcome = self.evaluate(&context);

        let episode = Episode::builder()
            .learn_model("worker_decision_sequence")
            .context(context)
            .outcome(outcome)
            .build();

        vec![episode]
    }

    fn convert(&self, episode: &Episode) -> Result<TrainingData, LearnError> {
        if !episode.outcome.is_success() {
            return Err(LearnError::InvalidEpisode(
                "Episode is not successful".into(),
            ));
        }

        let actions: Vec<&str> = episode
            .context
            .iter::<ActionRecord>()
            .map(|a| a.action.as_str())
            .collect();

        if actions.len() < self.min_actions {
            return Err(LearnError::InvalidEpisode(format!(
                "Too few actions: {} < {}",
                actions.len(),
                self.min_actions
            )));
        }

        // プロンプト生成
        let available = self.available_actions.join(", ");
        let prompt = format!(
            "Current context: default\n\
             Available actions: {}\n\n\
             What is the best sequence of actions to resolve this issue?",
            available
        );

        // レスポンス生成
        let action_sequence = actions.join(" -> ");
        let response = format!(
            "Based on the context, the optimal action sequence is: {}",
            action_sequence
        );

        Ok(TrainingData::sft(&self.system_prompt, &prompt, &response)
            .with_episode_id(episode.id.to_string())
            .with_outcome_score(episode.outcome.score()))
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::events::{ActionContext, ActionEventBuilder, ActionEventResult};
    use crate::types::WorkerId;
    use std::time::Duration;

    fn make_action(tick: u64, worker_id: usize, action: &str, success: bool) -> ActionEvent {
        let result = if success {
            ActionEventResult::success()
        } else {
            ActionEventResult::failure("error")
        };

        ActionEventBuilder::new(tick, WorkerId(worker_id), action)
            .result(result)
            .duration(Duration::from_millis(10))
            .context(ActionContext::new())
            .build()
    }

    fn make_records(actions: &[ActionEvent]) -> Vec<Record> {
        actions.iter().map(Record::from).collect()
    }

    #[test]
    fn test_record_accessors() {
        let action = make_action(1, 5, "CheckStatus", true);
        let record = Record::from(&action);

        assert!(record.is_action());
        assert!(!record.is_llm());
        assert_eq!(record.worker_id(), Some(5));
        assert!(record.as_action().is_some());
        assert!(record.as_llm().is_none());
    }

    #[test]
    fn test_record_stream_group_by_worker() {
        let actions = vec![
            make_action(1, 0, "A", true),
            make_action(2, 1, "B", true),
            make_action(3, 0, "C", true),
            make_action(4, 1, "D", true),
        ];
        let records = make_records(&actions);
        let stream = RecordStream::new(&records);

        let groups = stream.group_by_worker();
        assert_eq!(groups.len(), 2);
        assert_eq!(groups.get(&0).map(|v| v.len()), Some(2));
        assert_eq!(groups.get(&1).map(|v| v.len()), Some(2));
    }

    #[test]
    fn test_worker_task_learn_build_episodes() {
        let learn = WorkerTaskLearn::new().with_min_actions(2);

        let actions = vec![
            make_action(1, 0, "CheckStatus", true),
            make_action(2, 0, "ReadLogs", true),
            make_action(3, 0, "done", true),
            make_action(4, 1, "Grep", true),
            make_action(5, 1, "done", false),
        ];
        let records = make_records(&actions);

        let episodes = learn.build_episodes(&records);

        // Worker 0: success, Worker 1: failure
        assert_eq!(episodes.len(), 2);

        let worker0_ep = episodes.iter().find(|ep| ep.worker_id() == Some(0));
        assert!(worker0_ep.is_some());
        assert!(worker0_ep.unwrap().outcome.is_success());

        let worker1_ep = episodes.iter().find(|ep| ep.worker_id() == Some(1));
        assert!(worker1_ep.is_some());
        assert!(worker1_ep.unwrap().outcome.is_failure());
    }

    #[test]
    fn test_worker_task_learn_convert_success_only() {
        let learn = WorkerTaskLearn::new();

        // 失敗 Episode は変換エラー
        let failed_ep = Episode::builder()
            .learn_model("worker_task")
            .outcome(Outcome::failure("test"))
            .build();

        assert!(learn.convert(&failed_ep).is_err());

        // 成功 Episode は変換可能
        let success_ep = Episode::builder()
            .learn_model("worker_task")
            .record(ActionRecord::new(1, 0, "Check").success(true))
            .record(ActionRecord::new(2, 0, "Fix").success(true))
            .record(ActionRecord::new(3, 0, "done").success(true))
            .outcome(Outcome::success_binary())
            .build();

        assert!(learn.convert(&success_ep).is_ok());
    }

    // ========================================================================
    // WorkerDecisionSequenceLearn Tests
    // ========================================================================

    #[test]
    fn test_worker_decision_sequence_learn_build_episodes() {
        let learn = WorkerDecisionSequenceLearn::new().with_min_actions(3);

        let actions = vec![
            make_action(1, 0, "CheckStatus", true),
            make_action(2, 0, "ReadLogs", true),
            make_action(3, 0, "Restart", true),
            make_action(4, 0, "tick_end", true), // system event, should be filtered
            make_action(5, 0, "done", true),     // system event, should be filtered
        ];
        let records = make_records(&actions);

        let episodes = learn.build_episodes(&records);

        // 1 episode（成功アクション3つ: CheckStatus, ReadLogs, Restart）
        assert_eq!(episodes.len(), 1);
        assert!(episodes[0].outcome.is_success());

        // システムイベントは除外されている
        let action_count = episodes[0].context.iter::<ActionRecord>().count();
        assert_eq!(action_count, 3);
    }

    #[test]
    fn test_worker_decision_sequence_learn_filters_failed_actions() {
        let learn = WorkerDecisionSequenceLearn::new().with_min_actions(3);

        let actions = vec![
            make_action(1, 0, "CheckStatus", true),
            make_action(2, 0, "ReadLogs", false), // failed, should be filtered
            make_action(3, 0, "Restart", true),
        ];
        let records = make_records(&actions);

        let episodes = learn.build_episodes(&records);

        // 成功アクションが2つしかないので Episode は生成されない
        assert_eq!(episodes.len(), 0);
    }

    #[test]
    fn test_worker_decision_sequence_learn_convert() {
        let learn = WorkerDecisionSequenceLearn::new().with_available_actions(vec![
            "A".to_string(),
            "B".to_string(),
            "C".to_string(),
        ]);

        let episode = Episode::builder()
            .learn_model("worker_decision_sequence")
            .record(ActionRecord::new(1, 0, "A").success(true))
            .record(ActionRecord::new(2, 0, "B").success(true))
            .record(ActionRecord::new(3, 0, "C").success(true))
            .outcome(Outcome::success(1.0))
            .build();

        let result = learn.convert(&episode);
        assert!(result.is_ok());

        let training_data = result.unwrap();
        assert!(training_data.prompt.contains("Available actions: A, B, C"));
        assert!(training_data.chosen.contains("A -> B -> C"));
    }
}