swarm_engine_llm/

batch_processor.rs

//! Batch Processor - ManagerAgent 向け Batch LLM 処理
//!
//! ManagerAgent の `BatchDecisionRequest` を処理するための抽象化レイヤー。
//!
//! # 設計
//!
//! ```text
//! Core Layer
//! ├── ManagerAgent trait (prepare / finalize)
//! ├── DefaultBatchManagerAgent   ← Core層のデフォルト実装
//! ├── ContextStore / ContextView ← 正規化されたコンテキスト
//! └── ContextResolver            ← スコープ解決
//!
//! LLM Layer
//! ├── PromptBuilder              ← ResolvedContext → プロンプト
//! ├── BatchProcessor trait       ← Batch 処理の抽象
//! │   └── LlmBatchProcessor   ← Ollama 実装（仮想バッチ）
//! │
//! └── BatchInvoker 実装          ← LLM Batch 呼び出し
//!     └── OllamaBatchInvoker
//! ```
//!
//! # 仮想バッチ vs 真のバッチ
//!
//! Ollama は真の Batch API を持たないため、`LlmBatchProcessor` は
//! 内部で並列/順次処理を行う「仮想バッチ」として実装されます。
//! 将来 vLLM 等の真の Batch API を持つバックエンドに切り替える際は、
//! `BatchProcessor` trait の別実装を提供するだけで対応可能です。
//!
//! # 型の統一
//!
//! LLM層はCore層の型を直接使用するため、変換ロジックは不要です：
//! - `WorkerDecisionRequest` - リクエスト
//! - `DecisionResponse` - レスポンス

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use std::collections::HashMap;

use swarm_engine_core::actions::ActionDef;
use swarm_engine_core::agent::{BatchDecisionRequest, DecisionResponse, WorkerDecisionRequest};
use swarm_engine_core::exploration::DependencyGraph;
use swarm_engine_core::types::{LoraConfig, WorkerId};

use crate::decider::{LlmDecider, LlmError};

// ============================================================================
// BatchProcessor Trait
// ============================================================================

/// Batch 処理結果
pub type BatchProcessResult = Vec<(WorkerId, Result<DecisionResponse, BatchProcessError>)>;

/// Batch 処理エラー
#[derive(Debug, Clone, thiserror::Error)]
pub enum BatchProcessError {
    /// 一時的エラー（リトライ可能）
    #[error("Batch process error (transient): {0}")]
    Transient(String),

    /// 恒久的エラー（リトライ不可）
    #[error("Batch process error: {0}")]
    Permanent(String),
}

impl BatchProcessError {
    pub fn transient(message: impl Into<String>) -> Self {
        Self::Transient(message.into())
    }

    pub fn permanent(message: impl Into<String>) -> Self {
        Self::Permanent(message.into())
    }

    pub fn is_transient(&self) -> bool {
        matches!(self, Self::Transient(_))
    }

    pub fn message(&self) -> &str {
        match self {
            Self::Transient(msg) => msg,
            Self::Permanent(msg) => msg,
        }
    }
}

impl From<LlmError> for BatchProcessError {
    fn from(e: LlmError) -> Self {
        if e.is_transient() {
            Self::Transient(e.message().to_string())
        } else {
            Self::Permanent(e.message().to_string())
        }
    }
}

impl From<swarm_engine_core::error::SwarmError> for BatchProcessError {
    fn from(err: swarm_engine_core::error::SwarmError) -> Self {
        if err.is_transient() {
            Self::Transient(err.message())
        } else {
            Self::Permanent(err.message())
        }
    }
}

impl From<BatchProcessError> for swarm_engine_core::error::SwarmError {
    fn from(err: BatchProcessError) -> Self {
        match err {
            BatchProcessError::Transient(message) => {
                swarm_engine_core::error::SwarmError::LlmTransient { message }
            }
            BatchProcessError::Permanent(message) => {
                swarm_engine_core::error::SwarmError::LlmPermanent { message }
            }
        }
    }
}

/// Batch Processor trait
///
/// `BatchDecisionRequest` を受け取り、各 Worker への決定を返す。
/// バックエンド（Ollama, vLLM 等）に応じた実装を提供する。
pub trait BatchProcessor: Send + Sync {
    /// Batch リクエストを処理
    ///
    /// # Arguments
    /// * `request` - Core の `BatchDecisionRequest`
    ///
    /// # Returns
    /// 各 Worker への決定結果（WorkerId とペア）
    fn process(
        &self,
        request: BatchDecisionRequest,
    ) -> Pin<Box<dyn Future<Output = BatchProcessResult> + Send + '_>>;

    /// タスクとアクション一覧からアクション依存グラフを生成
    ///
    /// Swarm の Ticks 開始前に呼び出され、アクション間の依存関係を計画する。
    /// LLM を使用して動的に依存グラフを生成する。
    ///
    /// # Default
    /// デフォルトでは None を返す（依存グラフ生成をスキップ）。
    fn plan_dependencies(
        &self,
        _task: &str,
        _actions: &[ActionDef],
    ) -> Pin<Box<dyn Future<Output = Option<DependencyGraph>> + Send + '_>> {
        Box::pin(async { None })
    }

    /// ヘルスチェック
    fn is_healthy(&self) -> Pin<Box<dyn Future<Output = bool> + Send + '_>>;

    /// プロセッサ名
    fn name(&self) -> &str;
}

// ============================================================================
// LlmBatchProcessor
// ============================================================================

/// Ollama Batch Processor 設定
#[derive(Debug, Clone)]
pub struct LlmBatchProcessorConfig {
    /// 並列実行するか（false の場合は順次処理）
    pub parallel: bool,
    /// 並列実行時の最大同時実行数
    pub max_concurrency: usize,
    /// DependencyGraph 生成時の最大リトライ回数
    pub max_retries: Option<usize>,
}

impl Default for LlmBatchProcessorConfig {
    fn default() -> Self {
        Self {
            parallel: true,
            max_concurrency: 4,
            max_retries: Some(5),
        }
    }
}

/// Ollama Batch Processor
///
/// Ollama は真の Batch API を持たないため、仮想バッチとして実装。
/// 内部で `LlmDecider` を使用して並列/順次処理を行う。
pub struct LlmBatchProcessor<D: LlmDecider> {
    decider: Arc<D>,
    config: LlmBatchProcessorConfig,
}

impl<D: LlmDecider> LlmBatchProcessor<D> {
    /// 新しい LlmBatchProcessor を作成
    pub fn new(decider: D) -> Self {
        Self {
            decider: Arc::new(decider),
            config: LlmBatchProcessorConfig::default(),
        }
    }

    /// Arc でラップされた Decider から作成
    pub fn from_arc(decider: Arc<D>) -> Self {
        Self {
            decider,
            config: LlmBatchProcessorConfig::default(),
        }
    }

    /// 設定を指定して作成
    pub fn with_config(mut self, config: LlmBatchProcessorConfig) -> Self {
        self.config = config;
        self
    }
}

impl<D: LlmDecider + 'static> BatchProcessor for LlmBatchProcessor<D> {
    fn process(
        &self,
        request: BatchDecisionRequest,
    ) -> Pin<Box<dyn Future<Output = BatchProcessResult> + Send + '_>> {
        Box::pin(async move {
            if request.requests.is_empty() {
                return vec![];
            }

            // Core の WorkerDecisionRequest をそのまま使用（変換不要）
            let requests: Vec<(WorkerId, WorkerDecisionRequest)> = request
                .requests
                .into_iter()
                .map(|r| (r.worker_id, r))
                .collect();

            if self.config.parallel {
                self.process_parallel(requests).await
            } else {
                self.process_sequential(requests).await
            }
        })
    }

    fn plan_dependencies(
        &self,
        task: &str,
        actions: &[ActionDef],
    ) -> Pin<Box<dyn Future<Output = Option<DependencyGraph>> + Send + '_>> {
        let task = task.to_string();
        let actions: Vec<ActionDef> = actions.to_vec();
        let decider = Arc::clone(&self.decider);

        Box::pin(async move {
            use swarm_engine_core::actions::ActionCategory;
            use swarm_engine_core::exploration::DependencyGraphBuilder;

            let action_names: Vec<String> = actions.iter().map(|a| a.name.clone()).collect();

            // 1. Discover (NodeExpand) と NotDiscover (NodeStateChange) に分離
            let discover: Vec<&ActionDef> = actions
                .iter()
                .filter(|a| a.category == ActionCategory::NodeExpand)
                .collect();
            let not_discover: Vec<&ActionDef> = actions
                .iter()
                .filter(|a| a.category == ActionCategory::NodeStateChange)
                .collect();

            tracing::debug!(
                discover = ?discover.iter().map(|a| &a.name).collect::<Vec<_>>(),
                not_discover = ?not_discover.iter().map(|a| &a.name).collect::<Vec<_>>(),
                "Separated actions by category"
            );

            // 2. Discover も Binary + Vote でソート（順序関係を保持）
            let sorted_discover = if discover.len() <= 1 {
                discover.iter().map(|a| a.name.clone()).collect()
            } else {
                binary_sort_actions(&task, &discover, decider.as_ref()).await
            };

            tracing::debug!(
                sorted = ?sorted_discover,
                "Sorted Discover actions via binary comparison"
            );

            // 3. NotDiscover を Binary + Vote でソート
            let sorted_not_discover = if not_discover.len() <= 1 {
                not_discover.iter().map(|a| a.name.clone()).collect()
            } else {
                binary_sort_actions(&task, &not_discover, decider.as_ref()).await
            };

            tracing::debug!(
                sorted = ?sorted_not_discover,
                "Sorted NotDiscover actions via binary comparison"
            );

            // 4. グラフ構築: Discover（線形）→ NotDiscover（線形）
            let mut builder = DependencyGraphBuilder::new()
                .task(&task)
                .available_actions(action_names.clone());

            // 最初の Discover を Start node として設定
            if !sorted_discover.is_empty() {
                builder = builder.start_node(&sorted_discover[0]);
            } else if !sorted_not_discover.is_empty() {
                // Discover がなければ最初の NotDiscover を Start に
                builder = builder.start_node(&sorted_not_discover[0]);
            }

            // NotDiscover の最後を Terminal に
            if let Some(last) = sorted_not_discover.last() {
                builder = builder.terminal_node(last);
            } else if !sorted_discover.is_empty() {
                // NotDiscover がなければ最後の Discover を Terminal に
                builder = builder.terminal_node(sorted_discover.last().unwrap());
            }

            // Discover 間のエッジ（線形）
            for window in sorted_discover.windows(2) {
                builder = builder.edge(&window[0], &window[1], 0.9);
            }

            // 最後の Discover → 最初の NotDiscover へのエッジ
            if !sorted_discover.is_empty() && !sorted_not_discover.is_empty() {
                builder = builder.edge(
                    sorted_discover.last().unwrap(),
                    &sorted_not_discover[0],
                    0.9,
                );
            }

            // NotDiscover 間のエッジ（線形）
            for window in sorted_not_discover.windows(2) {
                builder = builder.edge(&window[0], &window[1], 0.9);
            }

            let graph = builder.build();

            tracing::info!(
                discover_order = ?sorted_discover,
                not_discover_order = ?sorted_not_discover,
                edges = graph.edges().len(),
                "DependencyGraph generated via binary sort"
            );

            Some(graph)
        })
    }

    fn is_healthy(&self) -> Pin<Box<dyn Future<Output = bool> + Send + '_>> {
        let decider = Arc::clone(&self.decider);
        Box::pin(async move { decider.is_healthy().await })
    }

    fn name(&self) -> &str {
        self.decider.model_name()
    }
}

impl<D: LlmDecider + 'static> LlmBatchProcessor<D> {
    /// 並列実行（LoRA グルーピング + Semaphore で同時実行数を制限）
    ///
    /// # LoRA グルーピング
    ///
    /// llama.cpp の continuous batching では、同じ LoRA 設定のリクエストは
    /// 効率的にバッチ処理される。異なる LoRA を混ぜると効率が落ちるため、
    /// リクエストを LoRA 設定でグルーピングして処理する。
    ///
    /// ```text
    /// リクエスト群
    /// ├── LoRA A のリクエスト群 → 並列実行（グループ内）
    /// ├── LoRA B のリクエスト群 → 並列実行（グループ内）
    /// └── LoRA なしのリクエスト群 → 並列実行（グループ内）
    /// ```
    ///
    /// グループ間は順次処理（同じ LoRA を連続して処理することで効率化）
    async fn process_parallel(
        &self,
        requests: Vec<(WorkerId, WorkerDecisionRequest)>,
    ) -> BatchProcessResult {
        // リクエストを LoRA 設定でグルーピング
        let grouped = group_by_lora(requests);

        let group_count = grouped.len();
        if group_count > 1 {
            tracing::debug!(
                groups = group_count,
                "Processing requests in {} LoRA groups",
                group_count
            );
        }

        // 各グループを順次処理（グループ内は並列）
        let mut all_results = Vec::new();
        for (lora_config, group_requests) in grouped {
            if group_count > 1 {
                tracing::trace!(
                    lora = ?lora_config,
                    count = group_requests.len(),
                    "Processing LoRA group"
                );
            }
            let results = self.process_group(group_requests).await;
            all_results.extend(results);
        }

        all_results
    }

    /// 単一グループの並列処理（Semaphore で同時実行数を制限）
    async fn process_group(
        &self,
        requests: Vec<(WorkerId, WorkerDecisionRequest)>,
    ) -> BatchProcessResult {
        use futures::future::join_all;
        use tokio::sync::Semaphore;

        // サーバーからスロット数を取得、取得できなければconfig値を使用
        let max_concurrency = self
            .decider
            .max_concurrency()
            .await
            .unwrap_or(self.config.max_concurrency);

        let semaphore = Arc::new(Semaphore::new(max_concurrency));

        let futures: Vec<_> = requests
            .into_iter()
            .map(|(worker_id, req)| {
                let decider = Arc::clone(&self.decider);
                let sem = Arc::clone(&semaphore);
                async move {
                    // スロットを取得してから実行
                    let _permit = sem.acquire().await.expect("Semaphore closed");
                    let result = decider.decide(req).await;
                    (worker_id, result)
                }
            })
            .collect();

        let results = join_all(futures).await;

        results
            .into_iter()
            .map(|(worker_id, result)| {
                let mapped = result.map_err(BatchProcessError::from);
                (worker_id, mapped)
            })
            .collect()
    }

    /// 順次実行
    async fn process_sequential(
        &self,
        requests: Vec<(WorkerId, WorkerDecisionRequest)>,
    ) -> BatchProcessResult {
        let mut results = Vec::with_capacity(requests.len());

        for (worker_id, req) in requests {
            let result = self.decider.decide(req).await;
            let mapped = result.map_err(BatchProcessError::from);
            results.push((worker_id, mapped));
        }

        results
    }
}

/// リクエストを LoRA 設定でグルーピング
///
/// 同じ LoRA 設定（または LoRA なし）のリクエストをまとめる。
/// HashMap の順序は不定だが、グループ内の順序は保持される。
fn group_by_lora(
    requests: Vec<(WorkerId, WorkerDecisionRequest)>,
) -> HashMap<Option<LoraConfig>, Vec<(WorkerId, WorkerDecisionRequest)>> {
    let mut groups: HashMap<Option<LoraConfig>, Vec<(WorkerId, WorkerDecisionRequest)>> =
        HashMap::new();

    for (worker_id, req) in requests {
        let lora_key = req.lora.clone();
        groups.entry(lora_key).or_default().push((worker_id, req));
    }

    groups
}

// ============================================================================
// Helper Functions
// ============================================================================

/// Binary + Vote でアクションをソート（バッチ版）
///
/// 全ペア × 3回分のプロンプトを一括でバッチ送信し、結果を集計。
/// 勝ち数でソート（勝ち数が少ない = 先に来る）。
async fn binary_sort_actions<D: LlmDecider>(
    task: &str,
    actions: &[&ActionDef],
    decider: &D,
) -> Vec<String> {
    use futures::future::join_all;
    use std::collections::HashMap;

    if actions.len() <= 1 {
        return actions.iter().map(|a| a.name.clone()).collect();
    }

    // 全ペア × 3回分のリクエストを作成
    // (pair_index, vote_index, prompt, a_name, b_name)
    let mut requests: Vec<(usize, usize, String, String, String)> = Vec::new();
    let mut pair_index = 0;

    for i in 0..actions.len() {
        for j in (i + 1)..actions.len() {
            let a = actions[i];
            let b = actions[j];
            let prompt = format!(
                "Goal: {}\n- {}: {}\n- {}: {}\nWhich comes first: {} or {}?\nAnswer (one word):",
                task, a.name, a.description, b.name, b.description, a.name, b.name
            );

            // 同じペアを3回投げる
            for vote_idx in 0..3 {
                requests.push((
                    pair_index,
                    vote_idx,
                    prompt.clone(),
                    a.name.clone(),
                    b.name.clone(),
                ));
            }
            pair_index += 1;
        }
    }

    let total_requests = requests.len();
    tracing::debug!(
        pairs = pair_index,
        total_requests = total_requests,
        "Binary sort: sending batch requests"
    );

    // 全リクエストを並列で送信
    // Note: Binary sort does not use LoRA (base model only)
    let futures: Vec<_> = requests
        .into_iter()
        .map(|(pair_idx, vote_idx, prompt, a_name, b_name)| {
            let decider_ref = decider;
            async move {
                let result = decider_ref.call_raw(&prompt, None).await;
                (pair_idx, vote_idx, result, a_name, b_name)
            }
        })
        .collect();

    let results = join_all(futures).await;

    // ペアごとに投票結果を集計
    // pair_index -> (a_count, b_count, a_name, b_name)
    let mut pair_votes: HashMap<usize, (usize, usize, String, String)> = HashMap::new();

    for (pair_idx, _vote_idx, result, a_name, b_name) in results {
        let entry = pair_votes
            .entry(pair_idx)
            .or_insert((0, 0, a_name.clone(), b_name.clone()));

        if let Ok(response) = result {
            let response_upper = response.to_uppercase();
            let a_upper = a_name.to_uppercase();
            let b_upper = b_name.to_uppercase();

            if response_upper.contains(&a_upper) {
                entry.0 += 1;
            } else if response_upper.contains(&b_upper) {
                entry.1 += 1;
            }
        }
    }

    // 各アクションの「勝ち数」をカウント
    let mut wins: HashMap<String, usize> = HashMap::new();
    for a in actions {
        wins.insert(a.name.clone(), 0);
    }

    for (_pair_idx, (a_count, b_count, a_name, b_name)) in pair_votes {
        // winner = 「先に来る方」なので、もう一方が「後」= 勝ち
        if a_count >= b_count {
            // a が先 → b に勝ち+1
            *wins.get_mut(&b_name).unwrap() += 1;
        } else {
            // b が先 → a に勝ち+1
            *wins.get_mut(&a_name).unwrap() += 1;
        }
    }

    // 勝ち数が少ない順にソート（先に来るものが少ない）
    let mut sorted: Vec<_> = wins.into_iter().collect();
    sorted.sort_by_key(|(_, count)| *count);

    tracing::debug!(
        sorted = ?sorted.iter().map(|(n, c)| format!("{}:{}", n, c)).collect::<Vec<_>>(),
        "Binary sort completed"
    );

    sorted.into_iter().map(|(name, _)| name).collect()
}

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

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

    #[test]
    fn test_batch_process_error_transient() {
        let err = BatchProcessError::transient("connection timeout");
        assert!(err.is_transient());
        assert_eq!(err.message(), "connection timeout");
    }

    #[test]
    fn test_batch_process_error_permanent() {
        let err = BatchProcessError::permanent("invalid model");
        assert!(!err.is_transient());
        assert_eq!(err.message(), "invalid model");
    }

    #[test]
    fn test_batch_process_error_from_llm_error() {
        let llm_err = LlmError::transient("timeout");
        let batch_err: BatchProcessError = llm_err.into();
        assert!(batch_err.is_transient());
        assert_eq!(batch_err.message(), "timeout");
    }

    #[test]
    fn test_ollama_batch_processor_config_default() {
        let config = LlmBatchProcessorConfig::default();
        assert!(config.parallel);
        assert_eq!(config.max_concurrency, 4);
    }

    // =========================================================================
    // Binary Sort Tests
    // =========================================================================

    use std::collections::HashMap;

    /// 同期版の binary_sort (テスト用)
    /// wins の計算ロジックをテスト
    fn binary_sort_sync(
        actions: &[&str],
        // (a, b) -> winner (先に来る方)
        comparator: impl Fn(&str, &str) -> String,
    ) -> Vec<String> {
        if actions.len() <= 1 {
            return actions.iter().map(|s| s.to_string()).collect();
        }

        let mut wins: HashMap<String, usize> = HashMap::new();
        for &a in actions {
            wins.insert(a.to_string(), 0);
        }

        for i in 0..actions.len() {
            for j in (i + 1)..actions.len() {
                let a = actions[i];
                let b = actions[j];
                let winner = comparator(a, b);

                // winner = 先に来る方 → もう一方が後 = 勝ち
                if winner == a {
                    *wins.get_mut(b).unwrap() += 1;
                } else {
                    *wins.get_mut(a).unwrap() += 1;
                }
            }
        }

        let mut sorted: Vec<_> = wins.into_iter().collect();
        sorted.sort_by_key(|(_, count)| *count);
        sorted.into_iter().map(|(name, _)| name).collect()
    }

    #[test]
    fn test_binary_sort_two_actions() {
        // Fetch が先、Summarize が後
        let result = binary_sort_sync(
            &["Fetch", "Summarize"],
            |a, _b| a.to_string(), // 常に a が先
        );
        assert_eq!(result, vec!["Fetch", "Summarize"]);

        // Summarize が先、Fetch が後
        let result = binary_sort_sync(
            &["Fetch", "Summarize"],
            |_a, b| b.to_string(), // 常に b が先
        );
        assert_eq!(result, vec!["Summarize", "Fetch"]);
    }

    #[test]
    fn test_binary_sort_three_actions() {
        // Test -> Deploy の順
        // comparator: 常に正しい順序を返す
        let result = binary_sort_sync(&["Test", "Deploy", "Build"], |a, b| {
            let order = ["Build", "Test", "Deploy"];
            let a_idx = order.iter().position(|&x| x == a).unwrap();
            let b_idx = order.iter().position(|&x| x == b).unwrap();
            if a_idx < b_idx {
                a.to_string()
            } else {
                b.to_string()
            }
        });
        assert_eq!(result, vec!["Build", "Test", "Deploy"]);
    }

    #[test]
    fn test_binary_sort_wins_calculation() {
        // 3つのアクション: A, B, C
        // 正しい順序: A -> B -> C
        // 比較結果:
        //   A vs B -> A が先 -> B に+1
        //   A vs C -> A が先 -> C に+1
        //   B vs C -> B が先 -> C に+1
        // wins = {A: 0, B: 1, C: 2}
        // ソート後: A(0), B(1), C(2)

        let mut wins: HashMap<String, usize> = HashMap::new();
        wins.insert("A".to_string(), 0);
        wins.insert("B".to_string(), 0);
        wins.insert("C".to_string(), 0);

        // A vs B: A が先 → B に+1
        *wins.get_mut("B").unwrap() += 1;
        // A vs C: A が先 → C に+1
        *wins.get_mut("C").unwrap() += 1;
        // B vs C: B が先 → C に+1
        *wins.get_mut("C").unwrap() += 1;

        assert_eq!(wins["A"], 0);
        assert_eq!(wins["B"], 1);
        assert_eq!(wins["C"], 2);

        let mut sorted: Vec<_> = wins.into_iter().collect();
        sorted.sort_by_key(|(_, count)| *count);
        let result: Vec<_> = sorted.into_iter().map(|(name, _)| name).collect();

        assert_eq!(result, vec!["A", "B", "C"]);
    }

    /// response から winner を抽出するロジックのテスト
    fn extract_winner(response: &str, a: &str, b: &str) -> Option<String> {
        let response_upper = response.to_uppercase();
        let a_upper = a.to_uppercase();
        let b_upper = b.to_uppercase();

        if response_upper.contains(&a_upper) {
            Some(a.to_string())
        } else if response_upper.contains(&b_upper) {
            Some(b.to_string())
        } else {
            None
        }
    }

    #[test]
    fn test_extract_winner() {
        // 正常ケース
        assert_eq!(
            extract_winner("Fetch", "Fetch", "Summarize"),
            Some("Fetch".to_string())
        );
        assert_eq!(
            extract_winner("Summarize", "Fetch", "Summarize"),
            Some("Summarize".to_string())
        );

        // 先頭スペース
        assert_eq!(
            extract_winner(" Fetch", "Fetch", "Summarize"),
            Some("Fetch".to_string())
        );

        // 大文字小文字
        assert_eq!(
            extract_winner("fetch", "Fetch", "Summarize"),
            Some("Fetch".to_string())
        );
        assert_eq!(
            extract_winner("FETCH", "Fetch", "Summarize"),
            Some("Fetch".to_string())
        );

        // 文中に含まれる
        assert_eq!(
            extract_winner("The answer is Fetch.", "Fetch", "Summarize"),
            Some("Fetch".to_string())
        );

        // どちらも含まれない
        assert_eq!(extract_winner("Unknown", "Fetch", "Summarize"), None);

        // 両方含まれる場合は先にマッチした方
        assert_eq!(
            extract_winner("Fetch then Summarize", "Fetch", "Summarize"),
            Some("Fetch".to_string())
        );
    }

    #[test]
    fn test_vote_majority() {
        // 3回の投票で多数決
        fn vote_majority(responses: &[&str], a: &str, b: &str) -> String {
            let mut a_count = 0;
            let mut b_count = 0;

            for response in responses {
                if let Some(winner) = extract_winner(response, a, b) {
                    if winner == a {
                        a_count += 1;
                    } else {
                        b_count += 1;
                    }
                }
            }

            if a_count >= b_count {
                a.to_string()
            } else {
                b.to_string()
            }
        }

        // 3回とも Fetch
        assert_eq!(
            vote_majority(&["Fetch", "Fetch", "Fetch"], "Fetch", "Summarize"),
            "Fetch"
        );

        // 2回 Fetch, 1回 Summarize
        assert_eq!(
            vote_majority(&["Fetch", "Summarize", "Fetch"], "Fetch", "Summarize"),
            "Fetch"
        );

        // 2回 Summarize, 1回 Fetch
        assert_eq!(
            vote_majority(&["Summarize", "Summarize", "Fetch"], "Fetch", "Summarize"),
            "Summarize"
        );

        // 同数の場合は a (Fetch) を返す
        assert_eq!(
            vote_majority(&["Fetch", "Summarize", "Unknown"], "Fetch", "Summarize"),
            "Fetch"
        );
    }

    // =========================================================================
    // LoRA Grouping Tests
    // =========================================================================

    use swarm_engine_core::context::{ContextTarget, GlobalContext, ResolvedContext};

    fn create_test_request(
        worker_id: usize,
        lora: Option<LoraConfig>,
    ) -> (WorkerId, WorkerDecisionRequest) {
        let global = GlobalContext {
            tick: 0,
            max_ticks: 100,
            progress: 0.0,
            success_rate: 0.0,
            task_description: Some("test".to_string()),
            hint: None,
        };
        let context = ResolvedContext::new(global, ContextTarget::Worker(WorkerId(worker_id)));

        (
            WorkerId(worker_id),
            WorkerDecisionRequest {
                worker_id: WorkerId(worker_id),
                query: format!("query_{}", worker_id),
                context,
                lora,
            },
        )
    }

    #[test]
    fn test_group_by_lora_single_group_no_lora() {
        let requests = vec![
            create_test_request(0, None),
            create_test_request(1, None),
            create_test_request(2, None),
        ];

        let groups = group_by_lora(requests);

        assert_eq!(groups.len(), 1);
        assert!(groups.contains_key(&None));
        assert_eq!(groups[&None].len(), 3);
    }

    #[test]
    fn test_group_by_lora_single_group_with_lora() {
        let lora = LoraConfig::with_id(0);
        let requests = vec![
            create_test_request(0, Some(lora.clone())),
            create_test_request(1, Some(lora.clone())),
        ];

        let groups = group_by_lora(requests);

        assert_eq!(groups.len(), 1);
        assert!(groups.contains_key(&Some(lora)));
    }

    #[test]
    fn test_group_by_lora_multiple_groups() {
        let lora_a = LoraConfig::with_id(0);
        let lora_b = LoraConfig::with_id(1);

        let requests = vec![
            create_test_request(0, Some(lora_a.clone())),
            create_test_request(1, Some(lora_b.clone())),
            create_test_request(2, Some(lora_a.clone())),
            create_test_request(3, None),
            create_test_request(4, Some(lora_b.clone())),
        ];

        let groups = group_by_lora(requests);

        assert_eq!(groups.len(), 3);
        assert_eq!(groups[&Some(lora_a)].len(), 2);
        assert_eq!(groups[&Some(lora_b)].len(), 2);
        assert_eq!(groups[&None].len(), 1);
    }

    #[test]
    fn test_group_by_lora_preserves_order_within_group() {
        let lora = LoraConfig::with_id(0);
        let requests = vec![
            create_test_request(5, Some(lora.clone())),
            create_test_request(3, Some(lora.clone())),
            create_test_request(7, Some(lora.clone())),
        ];

        let groups = group_by_lora(requests);
        let group = &groups[&Some(lora)];

        // グループ内の順序は保持される
        assert_eq!(group[0].0, WorkerId(5));
        assert_eq!(group[1].0, WorkerId(3));
        assert_eq!(group[2].0, WorkerId(7));
    }

    #[test]
    fn test_group_by_lora_different_scales() {
        // 同じ ID でも scale が違えば別グループ
        let lora_full = LoraConfig::new(0, 1.0);
        let lora_half = LoraConfig::new(0, 0.5);

        let requests = vec![
            create_test_request(0, Some(lora_full.clone())),
            create_test_request(1, Some(lora_half.clone())),
            create_test_request(2, Some(lora_full.clone())),
        ];

        let groups = group_by_lora(requests);

        assert_eq!(groups.len(), 2);
        assert_eq!(groups[&Some(lora_full)].len(), 2);
        assert_eq!(groups[&Some(lora_half)].len(), 1);
    }

    #[test]
    fn test_group_by_lora_empty() {
        let requests: Vec<(WorkerId, WorkerDecisionRequest)> = vec![];
        let groups = group_by_lora(requests);
        assert!(groups.is_empty());
    }
}