rustytorch/crates/models/rtx-llm-tools/tests/reasoning_tests.rs

use rtx_llm_tools::{
    error::{LlmToolsError, ReasoningError},
    reasoning::{
        ChainOfThought, Evidence, ReasoningContext, ReasoningEngine, ReasoningStep, StepValidator,
    },
};

#[tokio::test]
async fn test_basic_reasoning_chain() {
    let mut chain = ChainOfThought::new("basic_logic");

    let step1 = ReasoningStep::new(
        "step1",
        "All humans are mortal",
        vec![Evidence::new(
            "premise",
            "Universal truth about human mortality",
        )],
        "Humans have finite lifespans",
    );

    let step2 = ReasoningStep::new(
        "step2",
        "Socrates is human",
        vec![Evidence::new(
            "fact",
            "Historical record confirms Socrates was human",
        )],
        "Socrates belongs to the category of humans",
    );

    let step3 = ReasoningStep::new(
        "step3",
        "Therefore, Socrates is mortal",
        vec![
            Evidence::new("from_step1", "step1"),
            Evidence::new("from_step2", "step2"),
        ],
        "By logical deduction, Socrates must be mortal",
    );

    chain.add_step(step1).await.unwrap();
    chain.add_step(step2).await.unwrap();
    chain.add_step(step3).await.unwrap();

    let result = chain.execute().await.unwrap();
    assert!(result.is_valid);
    assert_eq!(
        result.conclusion,
        "By logical deduction, Socrates must be mortal"
    );
    assert_eq!(result.steps_executed, 3);
}

#[tokio::test]
async fn test_step_validation_failure() {
    let mut chain = ChainOfThought::new("invalid_logic");

    let step1 = ReasoningStep::new(
        "step1",
        "All birds can fly",
        vec![Evidence::new(
            "assumption",
            "Common but incorrect assumption",
        )],
        "Every bird has flight capability",
    );

    let step2 = ReasoningStep::new(
        "step2",
        "Penguins are birds",
        vec![Evidence::new("fact", "Penguins are classified as birds")],
        "Penguins belong to the bird category",
    );

    let step3 = ReasoningStep::new(
        "step3",
        "Therefore, penguins can fly",
        vec![
            Evidence::new("from_step1", "step1"),
            Evidence::new("from_step2", "step2"),
        ],
        "By deduction, penguins should be able to fly",
    );

    // Add a validator that knows penguins can't fly
    let mut validator = StepValidator::new();
    validator.add_contradiction_rule("penguin_flight", |step: &ReasoningStep| {
        step.conclusion.contains("penguins") && step.conclusion.contains("fly")
    });

    chain.set_validator(validator);
    chain.add_step(step1).await.unwrap();
    chain.add_step(step2).await.unwrap();
    chain.add_step(step3).await.unwrap();

    let result = chain.execute().await;
    assert!(result.is_err());

    match result.unwrap_err() {
        LlmToolsError::Reasoning(ReasoningError::Contradiction { step1, step2 }) => {
            assert!(step1.contains("step3") || step2.contains("step3"));
        }
        _ => panic!("Expected contradiction error"),
    }
}

#[tokio::test]
async fn test_evidence_insufficiency() {
    let mut chain = ChainOfThought::new("insufficient_evidence");

    let step = ReasoningStep::new(
        "step1",
        "Quantum consciousness exists",
        vec![], // No evidence provided
        "Human consciousness operates on quantum principles",
    );

    chain.add_step(step).await.unwrap();

    let result = chain.execute().await;
    assert!(result.is_err());

    match result.unwrap_err() {
        LlmToolsError::Reasoning(ReasoningError::InsufficientEvidence { step_id }) => {
            assert_eq!(step_id, "step1");
        }
        _ => panic!("Expected insufficient evidence error"),
    }
}

#[tokio::test]
async fn test_circular_reasoning_detection() {
    let mut chain = ChainOfThought::new("circular");

    let step1 = ReasoningStep::new(
        "step1",
        "A is true because B is true",
        vec![Evidence::new("from_step2", "step2")],
        "A is established",
    );

    let step2 = ReasoningStep::new(
        "step2",
        "B is true because A is true",
        vec![Evidence::new("from_step1", "step1")],
        "B is established",
    );

    chain.add_step(step1).await.unwrap();
    chain.add_step(step2).await.unwrap();

    let result = chain.execute().await;
    assert!(result.is_err());

    match result.unwrap_err() {
        LlmToolsError::Reasoning(ReasoningError::CircularReasoning) => {
            // Expected
        }
        _ => panic!("Expected circular reasoning error"),
    }
}

#[tokio::test]
async fn test_reasoning_context_variables() {
    let mut context = ReasoningContext::new();
    context.set_variable("domain", "mathematics");
    context.set_variable("rigor_level", "formal");

    let mut chain = ChainOfThought::with_context("math_proof", context);

    let step = ReasoningStep::new(
        "step1",
        "For all real numbers x, x² ≥ 0",
        vec![Evidence::new(
            "axiom",
            "Property of real numbers and multiplication",
        )],
        "Squares of real numbers are non-negative",
    );

    chain.add_step(step).await.unwrap();

    let result = chain.execute().await.unwrap();
    assert!(result.is_valid);

    let context = chain.get_context();
    assert_eq!(context.get_variable("domain"), Some("mathematics"));
    assert_eq!(context.get_variable("rigor_level"), Some("formal"));
}

#[tokio::test]
async fn test_step_branching() {
    let mut chain = ChainOfThought::new("branching");

    // Main premise
    let main_step = ReasoningStep::new(
        "main",
        "We need to solve the climate crisis",
        vec![Evidence::new(
            "consensus",
            "Scientific consensus on climate change",
        )],
        "Action is required to address climate change",
    );

    // Branch 1: Technology solution
    let tech_step = ReasoningStep::new(
        "tech_branch",
        "Technology can provide solutions",
        vec![Evidence::new("from_main", "main")],
        "Renewable energy and innovation can help",
    );

    // Branch 2: Policy solution
    let policy_step = ReasoningStep::new(
        "policy_branch",
        "Policy changes are necessary",
        vec![Evidence::new("from_main", "main")],
        "Regulation and incentives can drive change",
    );

    // Convergence step
    let convergence_step = ReasoningStep::new(
        "convergence",
        "Multiple approaches needed",
        vec![
            Evidence::new("from_tech", "tech_branch"),
            Evidence::new("from_policy", "policy_branch"),
        ],
        "Both technological and policy solutions are required",
    );

    chain.add_step(main_step).await.unwrap();
    chain.add_step_as_branch(tech_step, "main").await.unwrap();
    chain.add_step_as_branch(policy_step, "main").await.unwrap();
    chain.add_step(convergence_step).await.unwrap();

    let result = chain.execute().await.unwrap();
    assert!(result.is_valid);
    assert_eq!(result.steps_executed, 4);
}

#[tokio::test]
async fn test_reasoning_engine_with_templates() {
    let mut engine = ReasoningEngine::new();

    // Register a reasoning template for scientific method
    engine
        .register_template(
            "scientific_method",
            vec![
                "observation".to_string(),
                "hypothesis".to_string(),
                "prediction".to_string(),
                "experiment".to_string(),
                "analysis".to_string(),
                "conclusion".to_string(),
            ],
        )
        .await
        .unwrap();

    let mut context = ReasoningContext::new();
    context.set_variable("subject", "plant growth");
    context.set_variable("hypothesis", "Plants grow faster with more sunlight");

    let chain = engine
        .create_chain_from_template("scientific_method", context)
        .await
        .unwrap();
    assert_eq!(chain.expected_steps(), 6);
    assert_eq!(chain.get_template_name(), Some("scientific_method"));
}

#[tokio::test]
async fn test_reasoning_with_uncertainty() {
    let mut chain = ChainOfThought::new("uncertainty");

    let mut step1 = ReasoningStep::new(
        "step1",
        "There's a 70% chance of rain tomorrow",
        vec![Evidence::new("weather_model", "Meteorological prediction")],
        "Rain is likely but not certain",
    );
    step1.set_confidence(0.7);

    let mut step2 = ReasoningStep::new(
        "step2",
        "If it rains, the picnic will be cancelled",
        vec![Evidence::new("policy", "Event cancellation policy")],
        "Rain would force picnic cancellation",
    );
    step2.set_confidence(0.9);

    let step3 = ReasoningStep::new(
        "step3",
        "Therefore, there's about a 63% chance the picnic will be cancelled",
        vec![
            Evidence::new("from_step1", "step1"),
            Evidence::new("from_step2", "step2"),
        ],
        "Probability of cancellation is 0.7 × 0.9 = 0.63",
    );

    chain.add_step(step1).await.unwrap();
    chain.add_step(step2).await.unwrap();
    chain.add_step(step3).await.unwrap();

    let result = chain.execute().await.unwrap();
    assert!(result.is_valid);
    assert!((result.overall_confidence - 0.63).abs() < 0.01);
}

#[tokio::test]
async fn test_reasoning_timeout() {
    let mut chain = ChainOfThought::new("timeout_test");
    chain.set_timeout(std::time::Duration::from_millis(100));

    // Add a step that simulates long computation
    let slow_step = ReasoningStep::new(
        "slow",
        "Complex calculation requiring time",
        vec![Evidence::new("computation", "Mathematical proof")],
        "Result of lengthy computation",
    );

    chain.add_step(slow_step).await.unwrap();

    // Mock long execution time
    let result = chain.execute().await;
    // This should timeout (depending on implementation details)
    // For now, just verify it completes or times out appropriately
    assert!(
        result.is_ok()
            || matches!(
                result.unwrap_err(),
                LlmToolsError::Reasoning(ReasoningError::Timeout { .. })
            )
    );
}

#[tokio::test]
async fn test_step_dependency_validation() {
    let mut chain = ChainOfThought::new("dependency_validation");

    // Try to add a step that depends on a non-existent step
    let invalid_step = ReasoningStep::new(
        "invalid",
        "This depends on missing step",
        vec![Evidence::new("from_missing", "nonexistent_step")],
        "Should fail validation",
    );

    chain.add_step(invalid_step).await.unwrap();

    let result = chain.execute().await;
    assert!(result.is_err());

    match result.unwrap_err() {
        LlmToolsError::Reasoning(ReasoningError::InsufficientEvidence { .. }) => {
            // Expected - evidence references non-existent step
        }
        _ => panic!("Expected insufficient evidence error due to missing dependency"),
    }
}

#[tokio::test]
async fn test_reasoning_step_recovery() {
    let mut chain = ChainOfThought::new("recovery");
    chain.enable_error_recovery(true);

    let faulty_step = ReasoningStep::new(
        "faulty",
        "This step has issues",
        vec![Evidence::new("weak", "Questionable source")],
        "Dubious conclusion",
    );

    let recovery_step = ReasoningStep::new(
        "recovery",
        "Alternative approach",
        vec![Evidence::new("strong", "Reliable source")],
        "Better conclusion",
    );

    chain.add_step(faulty_step).await.unwrap();
    chain
        .add_recovery_step("faulty", recovery_step)
        .await
        .unwrap();

    let result = chain.execute().await.unwrap();
    assert!(result.is_valid);
    assert!(result.recovered_from_errors);
    assert_eq!(result.recovery_count, 1);
}