598 lines
17 KiB
Rust
598 lines
17 KiB
Rust
//! Chain-of-thought reasoning framework with step validation and error recovery
|
|
|
|
use crate::error::{ReasoningError, Result};
|
|
use async_trait::async_trait;
|
|
use indexmap::IndexMap;
|
|
use std::collections::{HashMap, HashSet};
|
|
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
use uuid::Uuid;
|
|
|
|
/// Evidence supporting a reasoning step
|
|
#[derive(Debug, Clone)]
|
|
pub struct Evidence {
|
|
pub id: String,
|
|
pub source_type: String,
|
|
pub description: String,
|
|
pub reliability: f64, // 0.0 to 1.0
|
|
pub timestamp: Instant,
|
|
}
|
|
|
|
impl Evidence {
|
|
pub fn new(source_type: impl Into<String>, description: impl Into<String>) -> Self {
|
|
Self {
|
|
id: Uuid::new_v4().to_string(),
|
|
source_type: source_type.into(),
|
|
description: description.into(),
|
|
reliability: 1.0,
|
|
timestamp: Instant::now(),
|
|
}
|
|
}
|
|
|
|
pub fn with_reliability(mut self, reliability: f64) -> Self {
|
|
self.reliability = reliability.clamp(0.0, 1.0);
|
|
self
|
|
}
|
|
}
|
|
|
|
/// Chain-of-thought reasoning step
|
|
#[derive(Debug, Clone)]
|
|
pub struct ReasoningStep {
|
|
pub id: String,
|
|
pub description: String,
|
|
pub evidence: Vec<Evidence>,
|
|
pub conclusion: String,
|
|
pub confidence: f64, // 0.0 to 1.0
|
|
pub step_type: StepType,
|
|
pub dependencies: HashSet<String>, // IDs of steps this depends on
|
|
pub created_at: Instant,
|
|
pub validation_status: ValidationStatus,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
|
pub enum StepType {
|
|
Premise,
|
|
Deduction,
|
|
Induction,
|
|
Abduction,
|
|
Hypothesis,
|
|
Observation,
|
|
Conclusion,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub enum ValidationStatus {
|
|
Pending,
|
|
Valid,
|
|
Invalid { reason: String },
|
|
Warning { message: String },
|
|
}
|
|
|
|
impl ReasoningStep {
|
|
pub fn new(
|
|
id: impl Into<String>,
|
|
description: impl Into<String>,
|
|
evidence: Vec<Evidence>,
|
|
conclusion: impl Into<String>,
|
|
) -> Self {
|
|
let id = id.into();
|
|
let evidence_deps: HashSet<String> = evidence
|
|
.iter()
|
|
.filter(|e| e.source_type.starts_with("from_step"))
|
|
.map(|e| e.description.clone())
|
|
.collect();
|
|
|
|
Self {
|
|
id,
|
|
description: description.into(),
|
|
evidence,
|
|
conclusion: conclusion.into(),
|
|
confidence: 1.0,
|
|
step_type: StepType::Deduction,
|
|
dependencies: evidence_deps,
|
|
created_at: Instant::now(),
|
|
validation_status: ValidationStatus::Pending,
|
|
}
|
|
}
|
|
|
|
pub fn set_confidence(&mut self, confidence: f64) {
|
|
self.confidence = confidence.clamp(0.0, 1.0);
|
|
}
|
|
|
|
pub fn set_step_type(&mut self, step_type: StepType) {
|
|
self.step_type = step_type;
|
|
}
|
|
|
|
pub fn add_evidence(&mut self, evidence: Evidence) {
|
|
if evidence.source_type.starts_with("from_step") {
|
|
self.dependencies.insert(evidence.description.clone());
|
|
}
|
|
self.evidence.push(evidence);
|
|
}
|
|
|
|
pub fn get_evidence_strength(&self) -> f64 {
|
|
if self.evidence.is_empty() {
|
|
return 0.0;
|
|
}
|
|
|
|
let total_reliability: f64 = self.evidence.iter().map(|e| e.reliability).sum();
|
|
total_reliability / self.evidence.len() as f64
|
|
}
|
|
}
|
|
|
|
/// Step validation trait
|
|
#[async_trait]
|
|
pub trait StepValidation {
|
|
async fn validate_step(
|
|
&self,
|
|
step: &ReasoningStep,
|
|
context: &ReasoningContext,
|
|
) -> Result<ValidationStatus>;
|
|
}
|
|
|
|
/// Rule-based step validator
|
|
pub struct StepValidator {
|
|
contradiction_rules: HashMap<String, Box<dyn Fn(&ReasoningStep) -> bool + Send + Sync>>,
|
|
evidence_requirements: HashMap<StepType, usize>,
|
|
confidence_thresholds: HashMap<StepType, f64>,
|
|
}
|
|
|
|
impl Default for StepValidator {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl StepValidator {
|
|
pub fn new() -> Self {
|
|
let mut evidence_requirements = HashMap::new();
|
|
evidence_requirements.insert(StepType::Premise, 1);
|
|
evidence_requirements.insert(StepType::Deduction, 2);
|
|
evidence_requirements.insert(StepType::Induction, 3);
|
|
evidence_requirements.insert(StepType::Hypothesis, 1);
|
|
evidence_requirements.insert(StepType::Observation, 1);
|
|
evidence_requirements.insert(StepType::Conclusion, 1);
|
|
|
|
let mut confidence_thresholds = HashMap::new();
|
|
confidence_thresholds.insert(StepType::Premise, 0.8);
|
|
confidence_thresholds.insert(StepType::Deduction, 0.7);
|
|
confidence_thresholds.insert(StepType::Conclusion, 0.6);
|
|
|
|
Self {
|
|
contradiction_rules: HashMap::new(),
|
|
evidence_requirements,
|
|
confidence_thresholds,
|
|
}
|
|
}
|
|
|
|
pub fn add_contradiction_rule<F>(&mut self, name: impl Into<String>, rule: F)
|
|
where
|
|
F: Fn(&ReasoningStep) -> bool + Send + Sync + 'static,
|
|
{
|
|
self.contradiction_rules.insert(name.into(), Box::new(rule));
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl StepValidation for StepValidator {
|
|
async fn validate_step(
|
|
&self,
|
|
step: &ReasoningStep,
|
|
_context: &ReasoningContext,
|
|
) -> Result<ValidationStatus> {
|
|
// Check evidence requirements
|
|
let required_evidence = self
|
|
.evidence_requirements
|
|
.get(&step.step_type)
|
|
.unwrap_or(&1);
|
|
|
|
if step.evidence.len() < *required_evidence {
|
|
return Ok(ValidationStatus::Invalid {
|
|
reason: format!(
|
|
"Step requires at least {} evidence items, but only {} provided",
|
|
required_evidence,
|
|
step.evidence.len()
|
|
),
|
|
});
|
|
}
|
|
|
|
// Check confidence threshold
|
|
if let Some(threshold) = self.confidence_thresholds.get(&step.step_type)
|
|
&& step.confidence < *threshold
|
|
{
|
|
return Ok(ValidationStatus::Warning {
|
|
message: format!(
|
|
"Step confidence {} is below recommended threshold {}",
|
|
step.confidence, threshold
|
|
),
|
|
});
|
|
}
|
|
|
|
// Check contradiction rules
|
|
for (rule_name, rule_fn) in &self.contradiction_rules {
|
|
if rule_fn(step) {
|
|
return Ok(ValidationStatus::Invalid {
|
|
reason: format!("Step violates contradiction rule: {rule_name}"),
|
|
});
|
|
}
|
|
}
|
|
|
|
// Check evidence strength
|
|
let evidence_strength = step.get_evidence_strength();
|
|
if evidence_strength < 0.5 {
|
|
return Ok(ValidationStatus::Warning {
|
|
message: format!("Low evidence strength: {evidence_strength:.2}"),
|
|
});
|
|
}
|
|
|
|
Ok(ValidationStatus::Valid)
|
|
}
|
|
}
|
|
|
|
/// Reasoning context for chain-of-thought processing
|
|
#[derive(Debug, Clone)]
|
|
pub struct ReasoningContext {
|
|
pub steps: IndexMap<String, ReasoningStep>,
|
|
pub variables: HashMap<String, String>,
|
|
pub metadata: HashMap<String, serde_json::Value>,
|
|
pub created_at: Instant,
|
|
}
|
|
|
|
impl ReasoningContext {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
steps: IndexMap::new(),
|
|
variables: HashMap::new(),
|
|
metadata: HashMap::new(),
|
|
created_at: Instant::now(),
|
|
}
|
|
}
|
|
|
|
pub fn set_variable(&mut self, key: impl Into<String>, value: impl Into<String>) {
|
|
self.variables.insert(key.into(), value.into());
|
|
}
|
|
|
|
pub fn get_variable(&self, key: &str) -> Option<&str> {
|
|
self.variables.get(key).map(|s| s.as_str())
|
|
}
|
|
|
|
pub fn add_step(&mut self, step: ReasoningStep) {
|
|
self.steps.insert(step.id.clone(), step);
|
|
}
|
|
|
|
pub fn get_step(&self, id: &str) -> Option<&ReasoningStep> {
|
|
self.steps.get(id)
|
|
}
|
|
|
|
pub fn validate_dependencies(&self) -> Result<()> {
|
|
for step in self.steps.values() {
|
|
for dep_id in &step.dependencies {
|
|
if !self.steps.contains_key(dep_id) {
|
|
return Err(ReasoningError::InsufficientEvidence {
|
|
step_id: step.id.clone(),
|
|
}
|
|
.into());
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub fn detect_circular_dependencies(&self) -> Result<()> {
|
|
let mut visited = HashSet::new();
|
|
let mut in_progress = HashSet::new();
|
|
|
|
for step_id in self.steps.keys() {
|
|
if !visited.contains(step_id) {
|
|
self.dfs_cycle_detection(step_id, &mut visited, &mut in_progress)?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn dfs_cycle_detection(
|
|
&self,
|
|
step_id: &str,
|
|
visited: &mut HashSet<String>,
|
|
in_progress: &mut HashSet<String>,
|
|
) -> Result<()> {
|
|
if in_progress.contains(step_id) {
|
|
return Err(ReasoningError::CircularReasoning.into());
|
|
}
|
|
|
|
if visited.contains(step_id) {
|
|
return Ok(());
|
|
}
|
|
|
|
in_progress.insert(step_id.to_string());
|
|
|
|
if let Some(step) = self.steps.get(step_id) {
|
|
for dep_id in &step.dependencies {
|
|
self.dfs_cycle_detection(dep_id, visited, in_progress)?;
|
|
}
|
|
}
|
|
|
|
in_progress.remove(step_id);
|
|
visited.insert(step_id.to_string());
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl Default for ReasoningContext {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
/// Reasoning execution result
|
|
#[derive(Debug)]
|
|
pub struct ReasoningResult {
|
|
pub is_valid: bool,
|
|
pub conclusion: String,
|
|
pub steps_executed: usize,
|
|
pub overall_confidence: f64,
|
|
pub execution_time: Duration,
|
|
pub errors: Vec<String>,
|
|
pub warnings: Vec<String>,
|
|
pub recovered_from_errors: bool,
|
|
pub recovery_count: usize,
|
|
}
|
|
|
|
/// Chain-of-thought reasoning engine
|
|
pub struct ChainOfThought {
|
|
name: String,
|
|
context: ReasoningContext,
|
|
validator: Option<Arc<dyn StepValidation + Send + Sync>>,
|
|
timeout: Option<Duration>,
|
|
error_recovery_enabled: bool,
|
|
recovery_steps: HashMap<String, ReasoningStep>,
|
|
expected_template_steps: Option<usize>,
|
|
template_name: Option<String>,
|
|
}
|
|
|
|
impl std::fmt::Debug for ChainOfThought {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("ChainOfThought")
|
|
.field("name", &self.name)
|
|
.field("context", &self.context)
|
|
.field("validator", &self.validator.is_some())
|
|
.field("timeout", &self.timeout)
|
|
.field("error_recovery_enabled", &self.error_recovery_enabled)
|
|
.field("recovery_steps", &self.recovery_steps.len())
|
|
.field("expected_template_steps", &self.expected_template_steps)
|
|
.field("template_name", &self.template_name)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
impl ChainOfThought {
|
|
pub fn new(name: impl Into<String>) -> Self {
|
|
Self {
|
|
name: name.into(),
|
|
context: ReasoningContext::new(),
|
|
validator: None,
|
|
timeout: None,
|
|
error_recovery_enabled: false,
|
|
recovery_steps: HashMap::new(),
|
|
expected_template_steps: None,
|
|
template_name: None,
|
|
}
|
|
}
|
|
|
|
pub fn with_context(name: impl Into<String>, context: ReasoningContext) -> Self {
|
|
Self {
|
|
name: name.into(),
|
|
context,
|
|
validator: None,
|
|
timeout: None,
|
|
error_recovery_enabled: false,
|
|
recovery_steps: HashMap::new(),
|
|
expected_template_steps: None,
|
|
template_name: None,
|
|
}
|
|
}
|
|
|
|
pub fn set_validator(&mut self, validator: impl StepValidation + Send + Sync + 'static) {
|
|
self.validator = Some(Arc::new(validator));
|
|
}
|
|
|
|
pub fn set_timeout(&mut self, timeout: Duration) {
|
|
self.timeout = Some(timeout);
|
|
}
|
|
|
|
pub fn enable_error_recovery(&mut self, enabled: bool) {
|
|
self.error_recovery_enabled = enabled;
|
|
}
|
|
|
|
pub fn get_context(&self) -> &ReasoningContext {
|
|
&self.context
|
|
}
|
|
|
|
pub fn expected_steps(&self) -> usize {
|
|
self.expected_template_steps
|
|
.unwrap_or(self.context.steps.len())
|
|
}
|
|
|
|
pub fn get_template_name(&self) -> Option<&str> {
|
|
self.template_name.as_deref()
|
|
}
|
|
|
|
pub async fn add_step(&mut self, step: ReasoningStep) -> Result<()> {
|
|
// Validate dependencies exist
|
|
for dep_id in &step.dependencies {
|
|
if !self.context.steps.contains_key(dep_id) {
|
|
return Err(ReasoningError::InsufficientEvidence {
|
|
step_id: step.id.clone(),
|
|
}
|
|
.into());
|
|
}
|
|
}
|
|
|
|
self.context.add_step(step);
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn add_step_as_branch(&mut self, step: ReasoningStep, parent_id: &str) -> Result<()> {
|
|
if !self.context.steps.contains_key(parent_id) {
|
|
return Err(ReasoningError::InsufficientEvidence {
|
|
step_id: step.id.clone(),
|
|
}
|
|
.into());
|
|
}
|
|
|
|
self.context.add_step(step);
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn add_recovery_step(
|
|
&mut self,
|
|
failed_step_id: impl Into<String>,
|
|
recovery_step: ReasoningStep,
|
|
) -> Result<()> {
|
|
self.recovery_steps
|
|
.insert(failed_step_id.into(), recovery_step);
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn execute(&mut self) -> Result<ReasoningResult> {
|
|
let start_time = Instant::now();
|
|
let errors = Vec::new();
|
|
let mut warnings = Vec::new();
|
|
let mut recovery_count = 0;
|
|
let mut recovered_from_errors = false;
|
|
|
|
// Validate overall structure
|
|
if let Err(e) = self.context.validate_dependencies() {
|
|
if self.error_recovery_enabled {
|
|
// Try to recover using recovery steps
|
|
recovery_count += 1;
|
|
recovered_from_errors = true;
|
|
} else {
|
|
return Err(e);
|
|
}
|
|
}
|
|
|
|
self.context.detect_circular_dependencies()?;
|
|
|
|
// Execute validation if validator is present
|
|
if let Some(validator) = &self.validator {
|
|
for step in self.context.steps.values() {
|
|
match validator.validate_step(step, &self.context).await? {
|
|
ValidationStatus::Invalid { reason: _ } => {
|
|
if self.error_recovery_enabled && self.recovery_steps.contains_key(&step.id)
|
|
{
|
|
recovery_count += 1;
|
|
recovered_from_errors = true;
|
|
} else {
|
|
return Err(ReasoningError::StepValidationFailed {
|
|
step_id: step.id.clone(),
|
|
}
|
|
.into());
|
|
}
|
|
}
|
|
ValidationStatus::Warning { message } => {
|
|
warnings.push(message);
|
|
}
|
|
ValidationStatus::Valid => {}
|
|
ValidationStatus::Pending => {}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Check for evidence insufficiency
|
|
for step in self.context.steps.values() {
|
|
if step.evidence.is_empty() && !matches!(step.step_type, StepType::Premise) {
|
|
return Err(ReasoningError::InsufficientEvidence {
|
|
step_id: step.id.clone(),
|
|
}
|
|
.into());
|
|
}
|
|
}
|
|
|
|
// Calculate overall confidence
|
|
let total_confidence: f64 = self
|
|
.context
|
|
.steps
|
|
.values()
|
|
.map(|s| s.confidence * s.get_evidence_strength())
|
|
.sum();
|
|
let overall_confidence = if self.context.steps.is_empty() {
|
|
0.0
|
|
} else {
|
|
total_confidence / self.context.steps.len() as f64
|
|
};
|
|
|
|
// Get final conclusion
|
|
let conclusion = self
|
|
.context
|
|
.steps
|
|
.values()
|
|
.last()
|
|
.map(|s| s.conclusion.clone())
|
|
.unwrap_or_default();
|
|
|
|
let execution_time = start_time.elapsed();
|
|
|
|
Ok(ReasoningResult {
|
|
is_valid: errors.is_empty(),
|
|
conclusion,
|
|
steps_executed: self.context.steps.len(),
|
|
overall_confidence,
|
|
execution_time,
|
|
errors,
|
|
warnings,
|
|
recovered_from_errors,
|
|
recovery_count,
|
|
})
|
|
}
|
|
}
|
|
|
|
/// Template-based reasoning engine
|
|
pub struct ReasoningEngine {
|
|
templates: HashMap<String, Vec<String>>,
|
|
default_validator: Arc<StepValidator>,
|
|
}
|
|
|
|
impl ReasoningEngine {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
templates: HashMap::new(),
|
|
default_validator: Arc::new(StepValidator::new()),
|
|
}
|
|
}
|
|
|
|
pub async fn register_template(
|
|
&mut self,
|
|
name: impl Into<String>,
|
|
steps: Vec<String>,
|
|
) -> Result<()> {
|
|
self.templates.insert(name.into(), steps);
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn create_chain_from_template(
|
|
&self,
|
|
template_name: &str,
|
|
context: ReasoningContext,
|
|
) -> Result<ChainOfThought> {
|
|
let template_steps = self.templates.get(template_name).ok_or_else(|| {
|
|
ReasoningError::InvalidStep(format!("Template not found: {template_name}"))
|
|
})?;
|
|
|
|
let mut chain = ChainOfThought::with_context(template_name, context);
|
|
chain.expected_template_steps = Some(template_steps.len());
|
|
chain.template_name = Some(template_name.to_string());
|
|
chain.set_validator(StepValidator::new());
|
|
|
|
Ok(chain)
|
|
}
|
|
}
|
|
|
|
impl Default for ReasoningEngine {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|