Files
rustytorch/crates/models/rtx-llm-tools/src/reasoning.rs
T
2026-03-04 00:08:42 +00:00

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()
}
}