From 5818ff341e68fe51080581ab6e28e8a3a506c8b4 Mon Sep 17 00:00:00 2001
From: CI-DEV <154627941+IlumCI@users.noreply.github.com>
Date: Wed, 19 Nov 2025 22:06:38 +0200
Subject: [PATCH] Update tree_of_thought_agent.py

---
 swarms/agents/tree_of_thought_agent.py | 815 +------------------------
 1 file changed, 29 insertions(+), 786 deletions(-)

diff --git a/swarms/agents/tree_of_thought_agent.py b/swarms/agents/tree_of_thought_agent.py
index 71dd43f8..8168536a 100644
--- a/swarms/agents/tree_of_thought_agent.py
+++ b/swarms/agents/tree_of_thought_agent.py
@@ -1,147 +1,3 @@
-"""
-Tree-of-Thought (ToT) Reasoning Framework Implementation.
-
-This module implements a comprehensive Tree-of-Thought reasoning system based on
-the formal framework where we introduce a tree-structured latent variable R
-representing multiple candidate reasoning paths.
-
-Mathematical Foundation:
-    Core Probabilistic Model:
-        p_θ(y | x) = Σ_{R ∈ T} p_θ(R | x) · p_θ(y | R, x)
-    
-    Where:
-    - x = input (question, task description) ∈ X
-    - y = final answer ∈ Y
-    - R = {r^(1), r^(2), ..., r^(k)} = set of candidate reasoning paths
-    - T = set of reasoning trees
-    - θ = model parameters
-    
-    Tree Structure:
-        T = (V, E) where:
-        - V = {v₁, v₂, ..., v_n} = nodes (thoughts)
-        - E = {(v_i, v_j) | v_i → v_j} = edges (reasoning transitions)
-        - Root: v_root = initial problem state
-        - Leaves: L = {v | children(v) = ∅}
-    
-    Path Probability:
-        P(path = (v₀, v₁, ..., v_k)) = Π_{i=0}^{k-1} P(v_{i+1} | v_i, x)
-        
-        Where P(v_{i+1} | v_i, x) is the transition probability.
-    
-    Marginalization over Tree:
-        p_θ(y | x) = Σ_{path ∈ paths(T)} P(path) · p_θ(y | path, x)
-        
-        Where paths(T) is the set of all root-to-leaf paths.
-    
-    Information-Theoretic Tree Search:
-        Information gain at node v:
-            I(v; Y | x) = H(Y | x) - H(Y | v, x)
-        
-        Expected information gain:
-            E[I(v; Y | x)] = Σ_{child} P(child | v) · I(child; Y | x)
-    
-    Quantum Tree Superposition:
-        |ψ_tree⟩ = Σ_{path} α_path |path⟩ ⊗ |y_path⟩
-        
-        Where:
-        - α_path = √(P(path)) = amplitude for path
-        - |path⟩ = quantum state representing reasoning path
-        - |y_path⟩ = answer state for path
-        
-        Measurement probability:
-            P(y | x) = |⟨y | ψ_tree⟩|² = |Σ_{path: y_path=y} α_path|²
-    
-    Monte Carlo Tree Search (MCTS):
-        UCB1 formula:
-            UCB1(v) = Q(v) + c · √(ln(N(v_parent)) / N(v))
-        
-        Where:
-        - Q(v) = average value: Q(v) = (1/N(v)) Σ_{i=1}^{N(v)} V_i
-        - N(v) = visit count
-        - c = exploration constant (typically √2)
-        - V_i = evaluation value from simulation i
-        
-        Value backpropagation:
-            Q(v) ← (N(v) · Q(v) + V_new) / (N(v) + 1)
-            N(v) ← N(v) + 1
-        
-        Selection policy:
-            v* = argmax_{v ∈ children(v_parent)} UCB1(v)
-    
-    Beam Search (Pruned Tree Search):
-        Beam width B, keep top-B nodes at each depth:
-            Beam_d = {v | v ∈ Top_B(score(v), v ∈ candidates_d)}
-        
-        Score function:
-            score(v) = α · heuristic(v) + β · depth_penalty(v) + γ · path_prob(v)
-        
-        Where:
-        - heuristic(v) = evaluator score
-        - depth_penalty(v) = -λ · depth(v)
-        - path_prob(v) = log P(path_to_v)
-    
-    Statistical Mechanics (Tree Energy):
-        Energy of path:
-            E(path, x) = -log P(path | x) = -Σ_{i} log P(v_{i+1} | v_i, x)
-        
-        Boltzmann distribution over paths:
-            P(path | x) = (1/Z(x)) exp(-E(path, x) / T)
-        
-        Partition function:
-            Z(x) = Σ_{path ∈ paths(T)} exp(-E(path, x) / T)
-        
-        Free energy:
-            F(x) = -T log Z(x)
-    
-    Graph-Theoretic Properties:
-        Tree depth: D = max_{path} |path|
-        Branching factor: b = avg_{v} |children(v)|
-        Tree size: |T| = Σ_{d=0}^D b^d (for balanced tree)
-        
-        Path diversity:
-            Diversity(T) = (1/|L|) Σ_{l₁, l₂ ∈ L} distance(l₁, l₂)
-        
-        Where distance is edit distance or semantic distance.
-    
-    Optimization Objective:
-        Best path selection:
-            path* = argmax_{path} [log p_θ(y | path, x) + λ · log P(path | x)]
-        
-        Multi-objective:
-            path* = argmax_{path} [w₁ · correctness + w₂ · efficiency + w₃ · diversity]
-    
-    Computational Complexity:
-        Time: O(b^D · (expand_cost + eval_cost))
-        - b = branching factor
-        - D = max depth
-        - expand_cost = cost to generate children
-        - eval_cost = cost to evaluate node
-        
-        With beam search (width B):
-            Time: O(B · D · (expand_cost + eval_cost))
-        
-        With MCTS (N simulations):
-            Time: O(N · (selection_cost + expand_cost + eval_cost + backprop_cost))
-    
-    Variational Tree Inference:
-        ELBO for tree search:
-            log p_θ(y | x) ≥ E_{q_φ(path|x,y)}[log p_θ(y | path, x)] - KL(q_φ(path|x,y) || p_θ(path|x))
-        
-        Where q_φ is the search policy (beam search, MCTS, etc.)
-
-At inference time:
-    1. Build a tree of reasoning paths by expanding nodes
-    2. Evaluate partial reasoning paths using heuristic
-    3. Search the tree using beam search or MCTS
-    4. Extract final answer from best leaf node
-    
-    Search strategies:
-    - Beam search: Top-B paths at each depth
-    - MCTS: UCB1-based exploration-exploitation
-    - BFS/DFS: Exhaustive or depth-limited search
-    - Quantum: Superposition of all paths with measurement
-"""
-
 from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional, Tuple, Union, Callable
 from enum import Enum
@@ -154,8 +10,6 @@ from loguru import logger
 
 
 class SearchStrategy(str, Enum):
-    """Search strategies for exploring the reasoning tree."""
-    
     BEAM = "beam"
     BFS = "bfs"
     DFS = "dfs"
@@ -164,23 +18,8 @@ class SearchStrategy(str, Enum):
 
 
 class TreeInformationTheory:
-    """
-    Information-theoretic utilities for tree reasoning analysis.
-    
-    Implements entropy, information gain, and tree diversity measures.
-    """
-    
     @staticmethod
     def path_entropy(path_probs: List[float]) -> float:
-        """
-        Calculate entropy of path distribution: H(Path | X) = -Σ P(path) log P(path).
-        
-        Args:
-            path_probs: List of path probabilities
-            
-        Returns:
-            Entropy value in bits
-        """
         if not path_probs:
             return 0.0
         
@@ -204,16 +43,6 @@ class TreeInformationTheory:
         prior_entropy: float,
         conditional_entropy: float
     ) -> float:
-        """
-        Calculate information gain: I(V; Y | X) = H(Y | X) - H(Y | V, X).
-        
-        Args:
-            prior_entropy: H(Y | X) - entropy before observing node
-            conditional_entropy: H(Y | V, X) - entropy after observing node
-            
-        Returns:
-            Information gain value
-        """
         return prior_entropy - conditional_entropy
     
     @staticmethod
@@ -221,16 +50,6 @@ class TreeInformationTheory:
         node_probs: List[float],
         child_entropies: List[float]
     ) -> float:
-        """
-        Calculate expected information gain: E[I(v; Y | x)] = Σ P(child) · I(child; Y | x).
-        
-        Args:
-            node_probs: Probabilities of child nodes P(child | v)
-            child_entropies: Information gains for each child
-            
-        Returns:
-            Expected information gain
-        """
         if not node_probs or not child_entropies:
             return 0.0
         
@@ -250,15 +69,6 @@ class TreeInformationTheory:
     
     @staticmethod
     def tree_diversity(leaf_paths: List[List[str]]) -> float:
-        """
-        Calculate tree diversity: Diversity(T) = (1/|L|) Σ distance(l₁, l₂).
-        
-        Args:
-            leaf_paths: List of paths to leaves (each path is list of node texts)
-            
-        Returns:
-            Diversity score (higher = more diverse)
-        """
         if len(leaf_paths) < 2:
             return 0.0
         
@@ -281,16 +91,6 @@ class TreeInformationTheory:
     
     @staticmethod
     def _edit_distance(s1: str, s2: str) -> float:
-        """
-        Calculate normalized edit distance between two strings.
-        
-        Args:
-            s1: First string
-            s2: Second string
-            
-        Returns:
-            Normalized edit distance [0, 1]
-        """
         if not s1 and not s2:
             return 0.0
         if not s1 or not s2:
@@ -317,23 +117,8 @@ class TreeInformationTheory:
 
 
 class QuantumTreeSearch:
-    """
-    Quantum-inspired tree search with superposition of paths.
-    
-    Implements: |ψ_tree⟩ = Σ_path α_path |path⟩
-    """
-    
     @staticmethod
     def calculate_path_amplitudes(path_probs: List[float]) -> List[float]:
-        """
-        Calculate quantum amplitudes: α_path = √(P(path)).
-        
-        Args:
-            path_probs: List of path probabilities
-            
-        Returns:
-            List of amplitudes
-        """
         return [math.sqrt(max(0.0, p)) for p in path_probs]
     
     @staticmethod
@@ -342,17 +127,6 @@ class QuantumTreeSearch:
         answers: List[str],
         path_probs: Optional[List[float]] = None
     ) -> Tuple[str, float]:
-        """
-        Quantum measurement: P(y | x) = |⟨y | ψ_tree⟩|² = |Σ_{path: y_path=y} α_path|².
-        
-        Args:
-            paths: List of reasoning paths (each path is list of node texts)
-            answers: List of answers corresponding to paths
-            path_probs: Optional path probabilities (uniform if None)
-            
-        Returns:
-            Tuple of (most likely answer, probability)
-        """
         if not paths or not answers:
             return "", 0.0
         
@@ -389,17 +163,6 @@ class QuantumTreeSearch:
         path_probs: List[float],
         num_samples: int = 1
     ) -> List[Any]:
-        """
-        Sample nodes using quantum-inspired superposition.
-        
-        Args:
-            nodes: List of nodes to sample from
-            path_probs: Probabilities for each node's path
-            num_samples: Number of samples to generate
-            
-        Returns:
-            List of sampled nodes
-        """
         if not nodes:
             return []
         
@@ -430,37 +193,12 @@ class QuantumTreeSearch:
 
 
 class TreeEnergyFunction:
-    """
-    Energy-based functions for tree reasoning (statistical mechanics).
-    
-    Implements: E(path, x) = -log P(path | x)
-    """
-    
     @staticmethod
     def calculate_path_energy(path_logprob: float) -> float:
-        """
-        Calculate energy of path: E(path, x) = -log P(path | x).
-        
-        Args:
-            path_logprob: Log probability of path
-            
-        Returns:
-            Energy value
-        """
         return -path_logprob
     
     @staticmethod
     def boltzmann_path_weight(energy: float, temperature: float) -> float:
-        """
-        Calculate Boltzmann weight: w(path) = exp(-E(path, x) / T).
-        
-        Args:
-            energy: Energy value E(path, x)
-            temperature: Temperature parameter T
-            
-        Returns:
-            Boltzmann weight
-        """
         if temperature <= 0:
             return 0.0 if energy > 0 else 1.0
         
@@ -471,16 +209,6 @@ class TreeEnergyFunction:
         path_energies: List[float],
         temperature: float
     ) -> float:
-        """
-        Calculate partition function: Z(x) = Σ_{path} exp(-E(path, x) / T).
-        
-        Args:
-            path_energies: List of energy values for paths
-            temperature: Temperature parameter T
-            
-        Returns:
-            Partition function value
-        """
         if temperature <= 0:
             return 1.0
         
@@ -489,16 +217,6 @@ class TreeEnergyFunction:
     
     @staticmethod
     def tree_free_energy(partition_function: float, temperature: float) -> float:
-        """
-        Calculate free energy: F(x) = -T log Z(x).
-        
-        Args:
-            partition_function: Partition function Z(x)
-            temperature: Temperature parameter T
-            
-        Returns:
-            Free energy value
-        """
         if partition_function <= 0:
             return float('inf')
         
@@ -514,18 +232,6 @@ class TreeEnergyFunction:
         temperature: float,
         num_samples: int = 1
     ) -> List[Any]:
-        """
-        Sample paths using Boltzmann distribution.
-        
-        Args:
-            paths: List of paths to sample from
-            path_logprobs: Log probabilities for each path
-            temperature: Temperature parameter T
-            num_samples: Number of samples to generate
-            
-        Returns:
-            List of sampled paths
-        """
         if not paths:
             return []
         
@@ -555,27 +261,11 @@ class TreeEnergyFunction:
 
 
 class TreeGraphTheory:
-    """
-    Graph-theoretic utilities for tree reasoning.
-    
-    Implements tree properties, path finding, and optimization.
-    """
-    
     @staticmethod
     def calculate_path_probability(
         node_scores: List[float],
         normalize: bool = True
     ) -> float:
-        """
-        Calculate path probability: P(path) = Π P(v_{i+1} | v_i, x).
-        
-        Args:
-            node_scores: Scores/probabilities for each node in path
-            normalize: Whether to normalize scores to probabilities
-            
-        Returns:
-            Path probability
-        """
         if not node_scores:
             return 0.0
         
@@ -605,18 +295,6 @@ class TreeGraphTheory:
         path_lengths: List[int],
         lambda_reg: float = 0.1
     ) -> Optional[List[Any]]:
-        """
-        Find optimal path: path* = argmax [score(path) - λ · length(path)].
-        
-        Args:
-            paths: List of paths (each path is list of nodes)
-            path_scores: Scores for each path
-            path_lengths: Lengths of each path
-            lambda_reg: Regularization parameter λ
-            
-        Returns:
-            Optimal path, or None if empty
-        """
         if not paths:
             return None
         
@@ -639,17 +317,6 @@ class TreeGraphTheory:
         get_depth: Callable[[Any], int],
         get_children: Callable[[Any], List[Any]]
     ) -> Dict[str, float]:
-        """
-        Calculate tree metrics: depth, branching factor, size.
-        
-        Args:
-            nodes: List of all nodes in tree
-            get_depth: Function to get depth of node
-            get_children: Function to get children of node
-            
-        Returns:
-            Dictionary with tree metrics
-        """
         if not nodes:
             return {
                 "max_depth": 0.0,
@@ -677,12 +344,6 @@ class TreeGraphTheory:
 
 
 class EnhancedMCTS:
-    """
-    Enhanced MCTS with mathematical foundations.
-    
-    Implements UCB1: UCB1(v) = Q(v) + c · √(ln(N(v_parent)) / N(v))
-    """
-    
     @staticmethod
     def calculate_ucb1(
         node_value: float,
@@ -690,18 +351,6 @@ class EnhancedMCTS:
         parent_visits: int,
         exploration_constant: float = math.sqrt(2)
     ) -> float:
-        """
-        Calculate UCB1 value: UCB1(v) = Q(v) + c · √(ln(N(v_parent)) / N(v)).
-        
-        Args:
-            node_value: Average value Q(v)
-            node_visits: Visit count N(v)
-            parent_visits: Parent visit count N(v_parent)
-            exploration_constant: Exploration constant c
-            
-        Returns:
-            UCB1 value
-        """
         if node_visits == 0:
             return float('inf')
         
@@ -721,17 +370,6 @@ class EnhancedMCTS:
         current_visits: int,
         new_value: float
     ) -> Tuple[float, int]:
-        """
-        Update node value: Q(v) ← (N(v) · Q(v) + V_new) / (N(v) + 1).
-        
-        Args:
-            current_value: Current average value Q(v)
-            current_visits: Current visit count N(v)
-            new_value: New evaluation value V_new
-            
-        Returns:
-            Tuple of (updated_value, updated_visits)
-        """
         new_visits = current_visits + 1
         updated_value = (current_value * current_visits + new_value) / new_visits
         
@@ -740,20 +378,6 @@ class EnhancedMCTS:
 
 @dataclass
 class ThoughtNode:
-    """
-    Represents a node in the Tree of Thought reasoning structure.
-    
-    Attributes:
-        id: Unique identifier for the node
-        depth: Depth of the node in the tree (0 = root)
-        text: Reasoning text accumulated so far
-        score: Evaluator score (heuristic value)
-        children: List of child ThoughtNode instances
-        parent: Optional parent ThoughtNode
-        is_leaf: Whether this is a leaf node (no children)
-        visit_count: Number of times visited (for MCTS)
-        value_sum: Sum of values from evaluations (for MCTS)
-    """
     
     id: UUID
     depth: int
@@ -766,23 +390,11 @@ class ThoughtNode:
     value_sum: float = 0.0
     
     def get_path(self) -> str:
-        """
-        Get the full reasoning path from root to this node.
-        
-        Returns:
-            Complete reasoning text from root to this node
-        """
         if self.parent is None:
             return self.text
         return self.parent.get_path() + self.text
     
     def get_average_value(self) -> float:
-        """
-        Get the average value for MCTS.
-        
-        Returns:
-            Average value (value_sum / visit_count) or 0.0 if not visited
-        """
         if self.visit_count == 0:
             return 0.0
         return self.value_sum / self.visit_count
@@ -790,23 +402,6 @@ class ThoughtNode:
 
 @dataclass
 class ToTConfig:
-    """
-    Configuration for Tree-of-Thought reasoning.
-    
-    Attributes:
-        max_depth: Maximum depth of the reasoning tree
-        branch_factor: Number of candidate thoughts to generate per node
-        beam_width: Number of nodes to keep in beam search
-        temperature: Sampling temperature for thought generation
-        top_p: Nucleus sampling parameter
-        max_thought_length: Maximum length of a single thought in tokens
-        search_strategy: Search strategy (beam, bfs, dfs, mcts)
-        mcts_simulations: Number of MCTS simulations (if using MCTS)
-        mcts_exploration: Exploration constant for UCB1 (if using MCTS)
-        evaluation_temperature: Temperature for evaluation prompts
-        answer_prefix: Prefix for final answer extraction
-        return_tree: Whether to return the full tree structure
-    """
     
     max_depth: int = 5
     branch_factor: int = 3
@@ -823,13 +418,6 @@ class ToTConfig:
 
 
 class LLMBackend:
-    """
-    Abstract interface for LLM backend.
-    
-    This defines the contract that any LLM implementation must follow
-    to work with the ToT framework.
-    """
-    
     def generate(
         self,
         prompt: str,
@@ -838,36 +426,10 @@ class LLMBackend:
         top_p: float = 0.9,
         stop: Optional[List[str]] = None,
     ) -> str:
-        """
-        Generate text from the LLM.
-        
-        Args:
-            prompt: Input prompt
-            max_tokens: Maximum tokens to generate
-            temperature: Sampling temperature
-            top_p: Nucleus sampling parameter
-            stop: List of stop sequences
-            
-        Returns:
-            Generated text
-            
-        Raises:
-            NotImplementedError: Must be implemented by subclass
-        """
         raise NotImplementedError("Subclass must implement generate method")
 
 
 class PromptBuilder:
-    """
-    Builds prompts for Tree-of-Thought reasoning.
-    
-    Creates prompts for:
-    - Initial problem statement
-    - Thought expansion (generating candidate continuations)
-    - Evaluation (scoring partial reasoning paths)
-    - Answer extraction (extracting final answer from reasoning)
-    """
-    
     DEFAULT_SYSTEM_PROMPT = """You are a helpful assistant that reasons through problems
 by exploring multiple possible approaches. Consider different perspectives and evaluate
 which reasoning paths are most promising."""
@@ -876,24 +438,9 @@ which reasoning paths are most promising."""
         self,
         system_prompt: Optional[str] = None,
     ):
-        """
-        Initialize the PromptBuilder.
-        
-        Args:
-            system_prompt: Custom system prompt (uses default if None)
-        """
         self.system_prompt = system_prompt or self.DEFAULT_SYSTEM_PROMPT
     
     def build_initial_prompt(self, problem: str) -> str:
-        """
-        Build the initial prompt for the problem.
-        
-        Args:
-            problem: The problem statement
-            
-        Returns:
-            Complete initial prompt
-        """
         return f"{self.system_prompt}\n\nProblem: {problem}\n\nLet's explore different approaches to solve this problem."
     
     def build_expansion_prompt(
@@ -902,17 +449,6 @@ which reasoning paths are most promising."""
         partial_reasoning: str,
         num_branches: int,
     ) -> str:
-        """
-        Build a prompt for expanding a node (generating candidate thoughts).
-        
-        Args:
-            problem: The original problem statement
-            partial_reasoning: Current reasoning path so far
-            num_branches: Number of candidate thoughts to generate
-            
-        Returns:
-            Prompt for thought expansion
-        """
         return f"""{self.system_prompt}
 
 Problem: {problem}
@@ -934,16 +470,6 @@ Return them as numbered items:
         problem: str,
         partial_reasoning: str,
     ) -> str:
-        """
-        Build a prompt for evaluating a partial reasoning path.
-        
-        Args:
-            problem: The original problem statement
-            partial_reasoning: The reasoning path to evaluate
-            
-        Returns:
-            Prompt for evaluation
-        """
         return f"""Given this reasoning so far, rate its plausibility and likelihood of correctness from 1 to 10.
 
 Problem: {problem}
@@ -964,16 +490,6 @@ Score:"""
         problem: str,
         reasoning: str,
     ) -> str:
-        """
-        Build a prompt for extracting the final answer from reasoning.
-        
-        Args:
-            problem: The original problem statement
-            reasoning: Complete reasoning path
-            
-        Returns:
-            Prompt for answer extraction
-        """
         return f"""{self.system_prompt}
 
 Problem: {problem}
@@ -986,32 +502,16 @@ Format: {self.answer_prefix} [your answer]"""
     
     @property
     def answer_prefix(self) -> str:
-        """Get the answer prefix."""
         return "Final answer:"
 
 
 class NodeExpander:
-    """
-    Expands nodes by generating candidate thoughts using the LLM.
-    
-    Given a partial reasoning sequence, asks the LLM to propose
-    multiple different next steps (thoughts).
-    """
-    
     def __init__(
         self,
         llm: LLMBackend,
         config: ToTConfig,
         prompt_builder: PromptBuilder,
     ):
-        """
-        Initialize the NodeExpander.
-        
-        Args:
-            llm: LLM backend instance
-            config: ToT configuration
-            prompt_builder: PromptBuilder instance
-        """
         self.llm = llm
         self.config = config
         self.prompt_builder = prompt_builder
@@ -1021,16 +521,6 @@ class NodeExpander:
         problem: str,
         node: ThoughtNode,
     ) -> List[str]:
-        """
-        Expand a node by generating candidate thoughts.
-        
-        Args:
-            problem: The original problem statement
-            node: The node to expand
-            
-        Returns:
-            List of candidate thought strings
-        """
         partial_reasoning = node.get_path()
         
         prompt = self.prompt_builder.build_expansion_prompt(
@@ -1059,15 +549,6 @@ class NodeExpander:
             return []
     
     def _parse_thoughts(self, response: str) -> List[str]:
-        """
-        Parse numbered thoughts from LLM response.
-        
-        Args:
-            response: LLM response text
-            
-        Returns:
-            List of thought strings
-        """
         thoughts = []
         lines = response.strip().split('\n')
         
@@ -1096,27 +577,12 @@ class NodeExpander:
 
 
 class Evaluator:
-    """
-    Evaluates partial reasoning paths using the LLM as a heuristic.
-    
-    Uses the LLM to assign a score to each node's partial reasoning,
-    estimating the likelihood that extending this path will yield a correct answer.
-    """
-    
     def __init__(
         self,
         llm: LLMBackend,
         config: ToTConfig,
         prompt_builder: PromptBuilder,
     ):
-        """
-        Initialize the Evaluator.
-        
-        Args:
-            llm: LLM backend instance
-            config: ToT configuration
-            prompt_builder: PromptBuilder instance
-        """
         self.llm = llm
         self.config = config
         self.prompt_builder = prompt_builder
@@ -1126,16 +592,6 @@ class Evaluator:
         problem: str,
         node: ThoughtNode,
     ) -> float:
-        """
-        Evaluate a node's partial reasoning path.
-        
-        Args:
-            problem: The original problem statement
-            node: The node to evaluate
-            
-        Returns:
-            Score from 1.0 to 10.0
-        """
         partial_reasoning = node.get_path()
         
         prompt = self.prompt_builder.build_evaluation_prompt(
@@ -1161,15 +617,6 @@ class Evaluator:
             return 5.0  # Default neutral score
     
     def _parse_score(self, response: str) -> float:
-        """
-        Parse score from evaluation response.
-        
-        Args:
-            response: LLM response text
-            
-        Returns:
-            Score from 1.0 to 10.0
-        """
         import re
         
         # Try to find a number in the response
@@ -1186,24 +633,11 @@ class Evaluator:
 
 
 class AnswerExtractor:
-    """
-    Extracts the final answer from a complete reasoning path.
-    
-    When a leaf node produces a final answer, this extracts and verifies it.
-    """
-    
     def __init__(
         self,
         llm: LLMBackend,
         prompt_builder: PromptBuilder,
     ):
-        """
-        Initialize the AnswerExtractor.
-        
-        Args:
-            llm: LLM backend instance
-            prompt_builder: PromptBuilder instance
-        """
         self.llm = llm
         self.prompt_builder = prompt_builder
     
@@ -1212,16 +646,6 @@ class AnswerExtractor:
         problem: str,
         reasoning: str,
     ) -> str:
-        """
-        Extract final answer from reasoning path.
-        
-        Args:
-            problem: The original problem statement
-            reasoning: Complete reasoning path
-            
-        Returns:
-            Extracted final answer
-        """
         prompt = self.prompt_builder.build_answer_extraction_prompt(
             problem=problem,
             reasoning=reasoning,
@@ -1245,15 +669,6 @@ class AnswerExtractor:
             return reasoning  # Fallback to reasoning text
     
     def _extract_answer(self, response: str) -> str:
-        """
-        Extract answer from response text.
-        
-        Args:
-            response: LLM response text
-            
-        Returns:
-            Extracted answer
-        """
         prefix = self.prompt_builder.answer_prefix.lower()
         response_lower = response.lower()
         
@@ -1268,26 +683,12 @@ class AnswerExtractor:
 
 
 class SearchController:
-    """
-    Implements search algorithms for exploring the reasoning tree.
-    
-    Supports beam search, BFS, DFS, and MCTS strategies.
-    """
-    
     def __init__(
         self,
         config: ToTConfig,
         expander: NodeExpander,
         evaluator: Evaluator,
     ):
-        """
-        Initialize the SearchController.
-        
-        Args:
-            config: ToT configuration
-            expander: NodeExpander instance
-            evaluator: Evaluator instance
-        """
         self.config = config
         self.expander = expander
         self.evaluator = evaluator
@@ -1297,16 +698,6 @@ class SearchController:
         problem: str,
         root: ThoughtNode,
     ) -> List[ThoughtNode]:
-        """
-        Search the reasoning tree to find the best leaf nodes.
-        
-        Args:
-            problem: The original problem statement
-            root: Root node of the tree
-            
-        Returns:
-            List of best leaf nodes
-        """
         if self.config.search_strategy == SearchStrategy.BEAM:
             return self._beam_search(problem, root)
         elif self.config.search_strategy == SearchStrategy.BFS:
@@ -1326,16 +717,6 @@ class SearchController:
         problem: str,
         root: ThoughtNode,
     ) -> List[ThoughtNode]:
-        """
-        Perform beam search on the reasoning tree.
-        
-        Args:
-            problem: The original problem statement
-            root: Root node of the tree
-            
-        Returns:
-            List of best leaf nodes
-        """
         frontier = [root]
         
         for depth in range(self.config.max_depth):
@@ -1402,16 +783,6 @@ class SearchController:
         problem: str,
         root: ThoughtNode,
     ) -> List[ThoughtNode]:
-        """
-        Perform breadth-first search on the reasoning tree.
-        
-        Args:
-            problem: The original problem statement
-            root: Root node of the tree
-            
-        Returns:
-            List of leaf nodes at max depth
-        """
         queue = deque([root])
         leaves = []
         
@@ -1445,16 +816,6 @@ class SearchController:
         problem: str,
         root: ThoughtNode,
     ) -> List[ThoughtNode]:
-        """
-        Perform depth-first search on the reasoning tree.
-        
-        Args:
-            problem: The original problem statement
-            root: Root node of the tree
-            
-        Returns:
-            List of best leaf nodes
-        """
         stack = [root]
         best_leaves = []
         
@@ -1496,16 +857,6 @@ class SearchController:
         problem: str,
         root: ThoughtNode,
     ) -> List[ThoughtNode]:
-        """
-        Perform Monte-Carlo Tree Search on the reasoning tree.
-        
-        Args:
-            problem: The original problem statement
-            root: Root node of the tree
-            
-        Returns:
-            List of best leaf nodes
-        """
         for _ in range(self.config.mcts_simulations):
             # Selection: traverse to leaf using UCB1
             node = self._mcts_select(root)
@@ -1537,15 +888,6 @@ class SearchController:
         return self._mcts_get_best_leaves(root)
     
     def _mcts_select(self, node: ThoughtNode) -> ThoughtNode:
-        """
-        Select a node using UCB1 formula.
-        
-        Args:
-            node: Current node
-            
-        Returns:
-            Selected node
-        """
         while node.children:
             if not all(child.visit_count > 0 for child in node.children):
                 # Select unvisited child
@@ -1581,15 +923,6 @@ class SearchController:
         node: ThoughtNode,
         value: float,
     ) -> None:
-        """
-        Backpropagate value up the tree.
-        
-        Uses: Q(v) ← (N(v) · Q(v) + V_new) / (N(v) + 1)
-        
-        Args:
-            node: Node to start backpropagation from
-            value: Value to backpropagate
-        """
         while node is not None:
             # Update using EnhancedMCTS formula
             updated_value, updated_visits = EnhancedMCTS.update_value(
@@ -1602,15 +935,6 @@ class SearchController:
             node = node.parent
     
     def _mcts_get_best_leaves(self, root: ThoughtNode) -> List[ThoughtNode]:
-        """
-        Get best leaf nodes from MCTS tree.
-        
-        Args:
-            root: Root of the tree
-            
-        Returns:
-            List of best leaf nodes
-        """
         leaves = []
         stack = [root]
         
@@ -1630,19 +954,6 @@ class SearchController:
         problem: str,
         root: ThoughtNode,
     ) -> List[ThoughtNode]:
-        """
-        Perform quantum-inspired tree search with superposition of paths.
-        
-        Implements: |ψ_tree⟩ = Σ_path α_path |path⟩
-        Measurement: P(y | x) = |⟨y | ψ_tree⟩|²
-        
-        Args:
-            problem: The original problem statement
-            root: Root node of the tree
-            
-        Returns:
-            List of best leaf nodes (sampled from quantum measurement)
-        """
         # Build tree using beam search first
         leaves = self._beam_search(problem, root)
         
@@ -1682,28 +993,11 @@ class SearchController:
 
 
 class ToTReasoner:
-    """
-    Main Tree-of-Thought reasoning engine.
-    
-    Implements the core ToT algorithm:
-    1. Build initial prompt
-    2. Create root node
-    3. Search the tree using selected strategy
-    4. Extract final answer from best leaf
-    """
-    
     def __init__(
         self,
         llm: LLMBackend,
         config: Optional[ToTConfig] = None,
     ):
-        """
-        Initialize the ToTReasoner.
-        
-        Args:
-            llm: LLM backend instance
-            config: ToT configuration (uses defaults if None)
-        """
         self.llm = llm
         self.config = config or ToTConfig()
         
@@ -1733,15 +1027,6 @@ class ToTReasoner:
         self,
         problem: str,
     ) -> Dict[str, Any]:
-        """
-        Solve a problem using Tree-of-Thought reasoning.
-        
-        Args:
-            problem: Problem statement to solve
-            
-        Returns:
-            Dictionary with final_answer, reasoning, score, and optionally tree
-        """
         logger.info(f"Starting ToT reasoning for problem: {problem[:100]}...")
         
         # Create root node
@@ -1881,19 +1166,7 @@ class ToTReasoner:
 
 
 class AgentLLMAdapter(LLMBackend):
-    """
-    Adapter to use Agent's LLM with the ToT framework.
-    
-    Wraps the Agent's LLM interface to match the LLMBackend contract.
-    """
-    
     def __init__(self, agent: Any):
-        """
-        Initialize the adapter.
-        
-        Args:
-            agent: Agent instance with an LLM, or direct LLM instance/callable
-        """
         self.agent = agent
         # Handle both Agent instances and direct LLM instances
         if hasattr(agent, 'llm'):
@@ -1912,19 +1185,6 @@ class AgentLLMAdapter(LLMBackend):
         top_p: float = 0.9,
         stop: Optional[List[str]] = None,
     ) -> str:
-        """
-        Generate text using the Agent's LLM.
-        
-        Args:
-            prompt: Input prompt
-            max_tokens: Maximum tokens to generate
-            temperature: Sampling temperature
-            top_p: Nucleus sampling parameter
-            stop: List of stop sequences
-            
-        Returns:
-            Generated text
-        """
         # Try to get LLM from agent if not directly available
         llm = self.llm
         if llm is None and hasattr(self.agent, 'llm'):
@@ -1991,25 +1251,6 @@ class ToTAgent:
     """
     Tree-of-Thought Agent for exploring multiple reasoning paths.
     
-    This agent implements the Tree-of-Thought (ToT) reasoning framework,
-    which introduces a tree-structured latent variable R representing
-    multiple candidate reasoning paths and searches over this tree.
-    
-    Mathematical Foundation:
-        p_θ(y | x) = Σ_{R ∈ T} p_θ(R | x) · p_θ(y | R, x)
-        
-        Where:
-        - x = input (question, task description)
-        - y = final answer
-        - R = {r^(1), r^(2), ..., r^(k)} = set of candidate reasoning paths
-        - T = set of reasoning trees
-    
-    Attributes:
-        agent_name: Name of the agent
-        model_name: LLM model to use
-        config: ToT configuration
-        reasoner: Internal ToTReasoner instance
-    
     Example:
         >>> from swarms.agents import ToTAgent
         >>> # Using model_name
@@ -2030,25 +1271,16 @@ class ToTAgent:
     def __init__(
         self,
         agent_name: str = "tot-agent",
+        description: Optional[str] = None,
         model_name: Optional[str] = "gpt-4o",
         llm: Optional[Any] = None,
         system_prompt: Optional[str] = None,
+        global_system_prompt: Optional[str] = None,
+        secondary_system_prompt: Optional[str] = None,
         config: Optional[ToTConfig] = None,
         agent: Optional[Any] = None,
         **kwargs,
     ):
-        """
-        Initialize the ToTAgent.
-        
-        Args:
-            agent_name: Name of the agent
-            model_name: LLM model name (used if agent/llm not provided)
-            llm: Optional LLM instance or callable (takes precedence over model_name)
-            system_prompt: Optional custom system prompt
-            config: ToT configuration (uses defaults if None)
-            agent: Optional Agent instance to use (if provided, uses its LLM)
-            **kwargs: Additional arguments passed to Agent if creating one
-        """
         self.agent_name = agent_name
         self.model_name = model_name
         self.config = config or ToTConfig()
@@ -2070,12 +1302,24 @@ class ToTAgent:
             # Import Agent here to avoid circular imports
             from swarms.structs.agent import Agent
             
-            self.agent = Agent(
-                agent_name=agent_name,
-                model_name=model_name,
-                system_prompt=system_prompt,
+            # Prepare agent kwargs
+            agent_kwargs = {
+                "agent_name": agent_name,
+                "model_name": model_name,
                 **kwargs,
-            )
+            }
+            
+            # Add optional parameters if provided
+            if description is not None:
+                agent_kwargs["agent_description"] = description
+            if system_prompt is not None:
+                agent_kwargs["system_prompt"] = system_prompt
+            if global_system_prompt is not None:
+                agent_kwargs["global_system_prompt"] = global_system_prompt
+            if secondary_system_prompt is not None:
+                agent_kwargs["secondary_system_prompt"] = secondary_system_prompt
+            
+            self.agent = Agent(**agent_kwargs)
             llm_adapter = AgentLLMAdapter(self.agent)
         
         # Initialize the ToT reasoner
@@ -2084,21 +1328,20 @@ class ToTAgent:
             config=self.config,
         )
     
+    def step(self, task: str, *args, **kwargs) -> str:
+        result = self.run(task, return_tree=False, *args, **kwargs)
+        return result if isinstance(result, str) else result.get("answer", "")
+    
+    def __getattr__(self, name: str):
+        if hasattr(self, 'agent'):
+            return getattr(self.agent, name)
+        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+    
     def run(
         self,
         task: str,
         return_tree: Optional[bool] = None,
     ) -> Union[str, Dict[str, Any]]:
-        """
-        Run the Tree-of-Thought agent on a task.
-        
-        Args:
-            task: Task or question to solve
-            return_tree: Whether to return full result with tree (defaults to config setting)
-            
-        Returns:
-            Final answer string, or full result dict if return_tree=True
-        """
         # Temporarily override return_tree if specified
         original_return_tree = self.config.return_tree
         if return_tree is not None: