swarm_engine_eval/environments/

deep_search.rs

//! DeepSearchEnvironment - 疑似Deep Search環境
//!
//! 情報検索・評価・統合をシミュレートする環境。
//! 実際のWeb検索ではなく、疑似ドキュメント集合を使用。
//!
//! # SearchEnvironment との違い
//!
//! - SearchEnvironment: 単純なファイル検索（1つの正解を探す）
//! - DeepSearchEnvironment: 情報の信頼性評価、複数ソースからの事実統合
//!
//! # アクション
//!
//! - `Search`: クエリで検索、関連ドキュメントのリストを返す
//! - `ReadDocument`: ドキュメントの内容を読む
//! - `EvaluateSource`: ソースの信頼度を評価
//! - `ExtractFact`: ドキュメントから事実を抽出
//! - `SubmitAnswer`: 回答を提出（完了判定）
//!
//! # 目標
//!
//! 信頼できる情報源から正確な回答を構築する。
//!
//! # DependencyGraph
//!
//! ```text
//! Search → ReadDocument
//! ReadDocument → EvaluateSource | ExtractFact
//! EvaluateSource → ExtractFact
//! ExtractFact → SubmitAnswer (terminal)
//! ```

use std::collections::HashSet;
use std::sync::RwLock;

use swarm_engine_core::agent::WorkResult;
use swarm_engine_core::environment::Environment;
use swarm_engine_core::types::{Action, WorkerId};

// ============================================================================
// Document & Source Definitions
// ============================================================================

/// ドキュメントの信頼度
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Reliability {
    /// 高信頼（公式ソース、学術論文等）
    High,
    /// 中信頼（ニュースサイト、専門ブログ等）
    Medium,
    /// 低信頼（個人ブログ、SNS等）
    Low,
    /// 誤情報を含む可能性
    Unreliable,
}

impl Reliability {
    fn score(&self) -> f64 {
        match self {
            Reliability::High => 1.0,
            Reliability::Medium => 0.7,
            Reliability::Low => 0.4,
            Reliability::Unreliable => 0.1,
        }
    }

    fn description(&self) -> &str {
        match self {
            Reliability::High => "High reliability - Official or academic source",
            Reliability::Medium => "Medium reliability - Reputable news or expert blog",
            Reliability::Low => "Low reliability - Personal blog or forum",
            Reliability::Unreliable => "Unreliable - May contain misinformation",
        }
    }
}

/// 疑似ドキュメント
#[derive(Debug, Clone)]
pub struct Document {
    pub id: String,
    pub title: String,
    pub source: String,
    pub content: String,
    pub keywords: Vec<String>,
    pub reliability: Reliability,
    /// このドキュメントが含む事実（正解かどうかのフラグ付き）
    pub facts: Vec<(String, bool)>,
}

impl Document {
    fn new(id: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            title: String::new(),
            source: String::new(),
            content: String::new(),
            keywords: Vec::new(),
            reliability: Reliability::Medium,
            facts: Vec::new(),
        }
    }

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

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

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

    fn keywords(mut self, keywords: Vec<&str>) -> Self {
        self.keywords = keywords.into_iter().map(String::from).collect();
        self
    }

    fn reliability(mut self, reliability: Reliability) -> Self {
        self.reliability = reliability;
        self
    }

    fn fact(mut self, fact: impl Into<String>, is_correct: bool) -> Self {
        self.facts.push((fact.into(), is_correct));
        self
    }

    /// クエリとの関連度を計算（0.0〜1.0）
    fn relevance(&self, query: &str) -> f64 {
        let query_lower = query.to_lowercase();
        let query_words: HashSet<_> = query_lower.split_whitespace().collect();

        let mut score: f64 = 0.0;

        // タイトルマッチ
        if self.title.to_lowercase().contains(&query_lower) {
            score += 0.4;
        }

        // キーワードマッチ
        for kw in &self.keywords {
            if query_words.contains(kw.to_lowercase().as_str()) {
                score += 0.2;
            }
        }

        // コンテンツマッチ（部分一致）
        let content_lower = self.content.to_lowercase();
        for word in &query_words {
            if content_lower.contains(*word) {
                score += 0.1;
            }
        }

        score.min(1.0)
    }
}

// ============================================================================
// DeepSearchEnvironment
// ============================================================================

/// 検索タスクの目標
#[derive(Debug, Clone)]
pub struct SearchGoal {
    /// 検索クエリ（ユーザーの質問）
    pub query: String,
    /// 正解となる事実のキーワード
    pub expected_facts: Vec<String>,
}

/// 環境の内部状態
struct SearchState {
    /// 疑似ドキュメント集合
    documents: Vec<Document>,
    /// 目標
    goal: SearchGoal,
    /// 検索履歴
    search_history: Vec<String>,
    /// 読んだドキュメント
    read_documents: HashSet<String>,
    /// 評価済みドキュメント
    evaluated_documents: HashSet<String>,
    /// 抽出した事実
    extracted_facts: Vec<(String, String, bool)>, // (doc_id, fact, is_correct)
    /// 完了フラグ
    done: bool,
}

impl SearchState {
    fn reset(&mut self, documents: Vec<Document>, goal: SearchGoal) {
        self.documents = documents;
        self.goal = goal;
        self.search_history.clear();
        self.read_documents.clear();
        self.evaluated_documents.clear();
        self.extracted_facts.clear();
        self.done = false;
    }
}

/// 疑似Deep Search環境
pub struct DeepSearchEnvironment {
    state: RwLock<SearchState>,
    initial_documents: Vec<Document>,
    initial_goal: SearchGoal,
}

impl DeepSearchEnvironment {
    /// 新しい検索環境を作成
    pub fn new(documents: Vec<Document>, goal: SearchGoal) -> Self {
        let state = SearchState {
            documents: documents.clone(),
            goal: goal.clone(),
            search_history: Vec::new(),
            read_documents: HashSet::new(),
            evaluated_documents: HashSet::new(),
            extracted_facts: Vec::new(),
            done: false,
        };

        Self {
            state: RwLock::new(state),
            initial_documents: documents,
            initial_goal: goal,
        }
    }

    /// プリセット: 技術質問シナリオ
    pub fn tech_question_scenario() -> Self {
        let documents = vec![
            Document::new("doc1")
                .title("Rust Memory Safety - Official Documentation")
                .source("doc.rust-lang.org")
                .content("Rust guarantees memory safety without garbage collection through its ownership system. The borrow checker enforces rules at compile time.")
                .keywords(vec!["rust", "memory", "safety", "ownership", "borrow"])
                .reliability(Reliability::High)
                .fact("Rust uses ownership system for memory safety", true)
                .fact("Rust has no garbage collector", true),

            Document::new("doc2")
                .title("Why Rust is Memory Safe - Tech Blog")
                .source("techblog.example.com")
                .content("Rust achieves memory safety through compile-time checks. The ownership model prevents data races and null pointer dereferences.")
                .keywords(vec!["rust", "memory", "safe", "compile"])
                .reliability(Reliability::Medium)
                .fact("Rust prevents data races at compile time", true)
                .fact("Rust has no runtime overhead for safety", true),

            Document::new("doc3")
                .title("Rust vs Go Memory Management")
                .source("forum.dev")
                .content("Some say Rust is slower because of all its safety checks. Go is better for most projects.")
                .keywords(vec!["rust", "go", "memory", "performance"])
                .reliability(Reliability::Low)
                .fact("Rust safety checks cause runtime overhead", false),

            Document::new("doc4")
                .title("Memory Safety Myths Debunked")
                .source("random-blog.net")
                .content("Rust is just hype. All languages can be memory safe if you're careful. Rust's borrow checker is annoying and unnecessary.")
                .keywords(vec!["rust", "memory", "hype"])
                .reliability(Reliability::Unreliable)
                .fact("Rust borrow checker is unnecessary", false),

            Document::new("doc5")
                .title("Understanding Ownership in Rust")
                .source("rust-book.example.org")
                .content("Each value in Rust has a variable that's called its owner. There can only be one owner at a time. When the owner goes out of scope, the value will be dropped.")
                .keywords(vec!["rust", "ownership", "scope", "drop"])
                .reliability(Reliability::High)
                .fact("Each value has exactly one owner", true)
                .fact("Values are dropped when owner goes out of scope", true),
        ];

        let goal = SearchGoal {
            query: "How does Rust achieve memory safety?".to_string(),
            expected_facts: vec![
                "ownership".to_string(),
                "borrow".to_string(),
                "compile".to_string(),
            ],
        };

        Self::new(documents, goal)
    }

    // ------------------------------------------------------------------------
    // Helpers
    // ------------------------------------------------------------------------

    /// Resolve doc_id from target.
    /// - If target is a valid doc_id (exists in documents), use it directly
    /// - Otherwise, auto-select first unread document
    fn resolve_doc_id_for_read(&self, target: &str) -> Option<String> {
        let state = self.state.read().unwrap();

        // Check if target is a valid doc_id
        if state.documents.iter().any(|d| d.id == target) {
            return Some(target.to_string());
        }

        // Auto-select first unread document
        state
            .documents
            .iter()
            .find(|d| !state.read_documents.contains(&d.id))
            .map(|d| d.id.clone())
    }

    /// Resolve doc_id from target.
    /// - If target is a valid doc_id (exists and has been read), use it directly
    /// - Otherwise, auto-select first read but unevaluated document
    fn resolve_doc_id_for_eval(&self, target: &str) -> Option<String> {
        let state = self.state.read().unwrap();

        // Check if target is a valid doc_id that has been read
        if state.read_documents.contains(target) {
            return Some(target.to_string());
        }

        // Auto-select first read but unevaluated document
        state
            .read_documents
            .iter()
            .find(|id| !state.evaluated_documents.contains(*id))
            .cloned()
    }

    /// Resolve doc_id from target.
    /// - If target is a valid doc_id (exists and has been read), use it directly
    /// - Otherwise, auto-select first read document
    fn resolve_doc_id_for_extract(&self, target: &str) -> Option<String> {
        let state = self.state.read().unwrap();

        // Check if target is a valid doc_id that has been read
        if state.read_documents.contains(target) {
            return Some(target.to_string());
        }

        // Auto-select first read document
        state.read_documents.iter().next().cloned()
    }

    // ------------------------------------------------------------------------
    // Action Handlers
    // ------------------------------------------------------------------------

    fn handle_search(&self, action: &Action) -> WorkResult {
        let target = action.params.target.as_deref().unwrap_or("");

        // If target is "node:X" format, use the goal query as the search query
        let query = if target.starts_with("node:") || target.is_empty() {
            let state = self.state.read().unwrap();
            state.goal.query.clone()
        } else {
            target.to_string()
        };

        if query.is_empty() {
            return WorkResult::env_failure("Search query is required");
        }

        let mut state = self.state.write().unwrap();
        state.search_history.push(query.clone());

        // 関連度でソートしてトップ5を返す
        let mut results: Vec<_> = state
            .documents
            .iter()
            .map(|doc| (doc, doc.relevance(&query)))
            .filter(|(_, rel)| *rel > 0.1)
            .collect();

        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
        results.truncate(5);

        if results.is_empty() {
            return WorkResult::env_success_with_data(
                "No relevant documents found",
                "Try different search terms".to_string(),
            );
        }

        // 表示用出力
        let output: Vec<String> = results
            .iter()
            .map(|(doc, rel)| {
                format!(
                    "- [{}] {} (source: {}, relevance: {:.0}%)",
                    doc.id,
                    doc.title,
                    doc.source,
                    rel * 100.0
                )
            })
            .collect();

        // 発見したドキュメントIDのリスト（ExploMap で展開される）
        let discovered_targets: Vec<String> =
            results.iter().map(|(doc, _)| doc.id.clone()).collect();

        WorkResult::env_success_with_discoveries(
            format!("Found {} relevant documents", results.len()),
            output.join("\n"),
            discovered_targets,
        )
    }

    fn handle_read_document(&self, action: &Action) -> WorkResult {
        let target = action.params.target.as_deref().unwrap_or("");

        let doc_id = match self.resolve_doc_id_for_read(target) {
            Some(id) => id,
            None => return WorkResult::env_failure("No unread documents available"),
        };

        let mut state = self.state.write().unwrap();

        let doc = match state.documents.iter().find(|d| d.id == doc_id) {
            Some(d) => d.clone(),
            None => return WorkResult::env_failure(format!("Document '{}' not found", doc_id)),
        };

        state.read_documents.insert(doc_id.to_string());

        let output = format!(
            "Title: {}\nSource: {}\n\nContent:\n{}",
            doc.title, doc.source, doc.content
        );

        WorkResult::env_success_with_data(format!("Read document '{}'", doc_id), output)
    }

    fn handle_evaluate_source(&self, action: &Action) -> WorkResult {
        let target = action.params.target.as_deref().unwrap_or("");
        let doc_id = match self.resolve_doc_id_for_eval(target) {
            Some(id) => id,
            None => {
                return WorkResult::env_failure(
                    "No documents available for evaluation (read documents first)",
                )
            }
        };

        let mut state = self.state.write().unwrap();

        // 先に読んでいないと評価できない
        if !state.read_documents.contains(&doc_id) {
            return WorkResult::env_failure(format!(
                "Must read document '{}' before evaluating",
                doc_id
            ));
        }

        let doc = match state.documents.iter().find(|d| d.id == doc_id) {
            Some(d) => d.clone(),
            None => return WorkResult::env_failure(format!("Document '{}' not found", doc_id)),
        };

        state.evaluated_documents.insert(doc_id.to_string());

        let output = format!(
            "Source: {}\nReliability: {:?} ({:.0}%)\nAssessment: {}",
            doc.source,
            doc.reliability,
            doc.reliability.score() * 100.0,
            doc.reliability.description()
        );

        WorkResult::env_success_with_data(format!("Evaluated source '{}'", doc_id), output)
    }

    fn handle_extract_fact(&self, action: &Action) -> WorkResult {
        let target = action.params.target.as_deref().unwrap_or("");
        let doc_id = match self.resolve_doc_id_for_extract(target) {
            Some(id) => id,
            None => return WorkResult::env_failure("No read documents available for extraction"),
        };

        let claim = action
            .params
            .args
            .get("claim")
            .map(|s| s.as_str())
            .unwrap_or("ownership"); // Default claim for auto-mode

        let mut state = self.state.write().unwrap();

        // 先に読んでいないと抽出できない
        if !state.read_documents.contains(&doc_id) {
            return WorkResult::env_failure(format!(
                "Must read document '{}' before extracting facts",
                doc_id
            ));
        }

        let doc = match state.documents.iter().find(|d| d.id == doc_id) {
            Some(d) => d.clone(),
            None => return WorkResult::env_failure(format!("Document '{}' not found", doc_id)),
        };

        // claim がドキュメントの facts に含まれるかチェック
        let claim_lower = claim.to_lowercase();
        let matched_fact = doc.facts.iter().find(|(fact, _)| {
            let fact_lower = fact.to_lowercase();
            claim_lower
                .split_whitespace()
                .any(|w| fact_lower.contains(w))
        });

        let (extracted, is_correct) = match matched_fact {
            Some((fact, correct)) => (fact.clone(), *correct),
            None => {
                return WorkResult::env_failure(format!(
                    "Claim '{}' not found in document '{}'",
                    claim, doc_id
                ));
            }
        };

        state
            .extracted_facts
            .push((doc_id.to_string(), extracted.clone(), is_correct));

        let confidence = if is_correct {
            doc.reliability.score()
        } else {
            0.0
        };

        let output = format!(
            "Extracted: \"{}\"\nSource reliability: {:?}\nConfidence: {:.0}%",
            extracted,
            doc.reliability,
            confidence * 100.0
        );

        WorkResult::env_success_with_data(format!("Extracted fact from '{}'", doc_id), output)
    }

    fn handle_submit_answer(&self, action: &Action) -> WorkResult {
        let answer = action.params.target.as_deref().unwrap_or("");
        if answer.is_empty() {
            return WorkResult::env_failure("Answer is required");
        }

        let mut state = self.state.write().unwrap();

        // 統計を先に計算
        let correct_count = state
            .extracted_facts
            .iter()
            .filter(|(_, _, correct)| *correct)
            .count();

        let incorrect_count = state
            .extracted_facts
            .iter()
            .filter(|(_, _, correct)| !*correct)
            .count();

        // 期待されるキーワードがいくつ含まれているか
        let answer_lower = answer.to_lowercase();
        let matched_count = state
            .goal
            .expected_facts
            .iter()
            .filter(|kw| answer_lower.contains(&kw.to_lowercase()))
            .count();

        let total_expected = state.goal.expected_facts.len();
        let keyword_coverage = matched_count as f64 / total_expected as f64;

        // スコア計算
        let correct_score = correct_count as f64 * 0.3;
        let incorrect_penalty = incorrect_count as f64 * 0.2;
        let coverage_score = keyword_coverage * 0.4;
        let total_score = (correct_score + coverage_score - incorrect_penalty)
            .max(0.0)
            .min(1.0);

        let success = total_score >= 0.6 && incorrect_count == 0;

        state.done = true;

        let summary = format!(
            "Answer evaluation:\n\
             - Correct facts extracted: {}\n\
             - Incorrect facts extracted: {}\n\
             - Keyword coverage: {:.0}%\n\
             - Total score: {:.0}%\n\
             - Result: {}",
            correct_count,
            incorrect_count,
            keyword_coverage * 100.0,
            total_score * 100.0,
            if success {
                "SUCCESS"
            } else {
                "NEEDS IMPROVEMENT"
            }
        );

        if success {
            WorkResult::done_success(summary)
        } else {
            WorkResult::done_failure(summary)
        }
    }
}

impl Environment for DeepSearchEnvironment {
    fn step(&self, _worker_id: WorkerId, action: &Action) -> WorkResult {
        match action.name.as_str() {
            "Search" => self.handle_search(action),
            "ReadDocument" => self.handle_read_document(action),
            "EvaluateSource" => self.handle_evaluate_source(action),
            "ExtractFact" => self.handle_extract_fact(action),
            "SubmitAnswer" => self.handle_submit_answer(action),
            _ => WorkResult::unsupported(&action.name),
        }
    }

    fn reset(&self) {
        let mut state = self.state.write().unwrap();
        state.reset(self.initial_documents.clone(), self.initial_goal.clone());
    }

    fn name(&self) -> &str {
        "DeepSearchEnvironment"
    }
}

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

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

    fn is_success(result: &WorkResult) -> bool {
        match result {
            WorkResult::Acted { action_result, .. } => action_result.success,
            WorkResult::Done { success, .. } => *success,
            _ => false,
        }
    }

    fn is_done(result: &WorkResult) -> bool {
        matches!(result, WorkResult::Done { .. })
    }

    fn get_output(result: &WorkResult) -> Option<String> {
        match result {
            WorkResult::Acted { action_result, .. } => action_result
                .output
                .as_ref()
                .map(|o| o.as_text().to_string()),
            WorkResult::Done { message, .. } => message.clone(),
            _ => None,
        }
    }

    fn get_error(result: &WorkResult) -> Option<String> {
        match result {
            WorkResult::Acted { action_result, .. } => action_result.error.clone(),
            _ => None,
        }
    }

    fn action(name: &str, target: Option<&str>) -> Action {
        Action {
            name: name.into(),
            params: swarm_engine_core::types::ActionParams {
                target: target.map(|s| s.into()),
                args: HashMap::new(),
                data: vec![],
            },
        }
    }

    fn action_with_args(name: &str, target: Option<&str>, args: Vec<(&str, &str)>) -> Action {
        let mut a = action(name, target);
        for (k, v) in args {
            a.params.args.insert(k.to_string(), v.to_string());
        }
        a
    }

    #[test]
    fn test_search_returns_relevant_documents() {
        let env = DeepSearchEnvironment::tech_question_scenario();

        let act = action("Search", Some("rust memory safety"));
        let result = env.step(WorkerId(0), &act);

        assert!(is_success(&result));
        let output = get_output(&result).unwrap();
        assert!(output.contains("doc1"));
    }

    #[test]
    fn test_read_document() {
        let env = DeepSearchEnvironment::tech_question_scenario();

        let act = action("ReadDocument", Some("doc1"));
        let result = env.step(WorkerId(0), &act);

        assert!(is_success(&result));
        let output = get_output(&result).unwrap();
        assert!(output.contains("ownership"));
    }

    #[test]
    fn test_evaluate_requires_read_first() {
        let env = DeepSearchEnvironment::tech_question_scenario();

        // Try to evaluate without reading first
        let act = action("EvaluateSource", Some("doc1"));
        let result = env.step(WorkerId(0), &act);

        assert!(!is_success(&result));
        // Error message: either "Must read" (explicit doc_id) or "No documents available" (auto-resolve)
        let err = get_error(&result).unwrap();
        assert!(
            err.contains("Must read") || err.contains("No documents available"),
            "Unexpected error: {}",
            err
        );
    }

    #[test]
    fn test_extract_fact() {
        let env = DeepSearchEnvironment::tech_question_scenario();

        // Read first
        let read = action("ReadDocument", Some("doc1"));
        env.step(WorkerId(0), &read);

        // Extract fact
        let extract = action_with_args("ExtractFact", Some("doc1"), vec![("claim", "ownership")]);
        let result = env.step(WorkerId(0), &extract);

        assert!(is_success(&result));
    }

    #[test]
    fn test_full_workflow() {
        let env = DeepSearchEnvironment::tech_question_scenario();

        // 1. Search
        let search = action("Search", Some("rust memory"));
        let result = env.step(WorkerId(0), &search);
        assert!(is_success(&result));

        // 2. Read high-reliability document
        let read = action("ReadDocument", Some("doc1"));
        let result = env.step(WorkerId(0), &read);
        assert!(is_success(&result));

        // 3. Evaluate source
        let eval = action("EvaluateSource", Some("doc1"));
        let result = env.step(WorkerId(0), &eval);
        assert!(is_success(&result));

        // 4. Extract fact
        let extract = action_with_args("ExtractFact", Some("doc1"), vec![("claim", "ownership")]);
        let result = env.step(WorkerId(0), &extract);
        assert!(is_success(&result));

        // 5. Submit answer
        let submit = action(
            "SubmitAnswer",
            Some(
                "Rust achieves memory safety through ownership and borrow checking at compile time",
            ),
        );
        let result = env.step(WorkerId(0), &submit);
        assert!(is_done(&result));
    }

    #[test]
    fn test_node_target_format() {
        let env = DeepSearchEnvironment::tech_question_scenario();

        // Search with "node:1" target should use goal.query
        let search = action("Search", Some("node:1"));
        let result = env.step(WorkerId(0), &search);
        assert!(
            is_success(&result),
            "Search with node:1 failed: {:?}",
            get_error(&result)
        );
        let output = get_output(&result).unwrap();
        assert!(output.contains("doc"), "Search output: {}", output);

        // ReadDocument with "node:2" target should auto-select first unread doc
        let read = action("ReadDocument", Some("node:2"));
        let result = env.step(WorkerId(0), &read);
        assert!(
            is_success(&result),
            "ReadDocument with node:2 failed: {:?}",
            get_error(&result)
        );
    }
}