diff --git a/mbodied/tree/README.md b/mbodied/tree/README.md new file mode 100644 index 0000000..2f240c3 --- /dev/null +++ b/mbodied/tree/README.md @@ -0,0 +1,187 @@ +# Tree of Thought - README + +## Overview +This repository implements a **Tree of Thought** (ToT) framework inspired by the decision-making process in AI systems. The Tree of Thought leverages Large Language Models (LLMs) to generate, evaluate, and expand thoughts (actions or decisions) recursively. The system allows for the exploration of multiple possible actions and evaluates each step to find the best path through a tree-structured reasoning process. + +For reference, see the [Tree of Thoughts: Deliberate Problem Solving with LLMs](https://arxiv.org/pdf/2305.10601). + +## Table of Contents +- [Key Concepts](#key-concepts) +- [Usage](#usage) +- [Tree of Thought Components](#tree-of-thought-components) + - [ThoughtNode](#thoughtnode) + - [TreeOfThought](#treeofthought) +- [Embedding with PCA](#embedding-with-pca) +- [Best Path Calculation](#best-path-calculation) +- [Visualization](#visualization) + + +## Key Concepts +- **ThoughtNode**: A node in the tree representing a thought (or action). Each node has an evaluation score and can have child nodes representing subsequent thoughts. +- **Tree of Thought (ToT)**: A tree structure that explores various paths of decisions, where each node represents a thought, and the branches represent possible follow-up actions. +- **PCA (Principal Component Analysis)**: A method used in this framework to reduce the dimensionality of thought embeddings for performance optimization. + +## Usage +To run the system, you will need to set up a **LanguageAgent** and provide a task/instruction. You can optionally pass an image to be processed along with the instructions. + +### Example +```python +from mbodied.agents import LanguageAgent +from mbodied.types.sense.vision import Image +from tree_of_thought import TreeOfThought + +image = Image(path="resources/color_image.png") +cognition = LanguageAgent( + context="You are an embodied planner that responds with a python list of strings and nothing else.", + api_key=os.getenv("OPENAI_API_KEY"), + model_src="openai", + recorder="auto", +) +tree_of_thought = TreeOfThought(language_agent=cognition, n_components=10, max_depth=3,) + +tree_of_thought.generate_thoughts(instruction="Switch the position of the remote and the fork", image=image) +tree_of_thought.traverse() +tree_of_thought.get_actions() +``` + +### Output: +``` +Level 0: Thought: Start Evaluation: 0.5 + Level 1: Thought: Pick up the remote Evaluation: 0.5 + Level 2: Thought: move forward Evaluation: 8.0 + Level 3: Thought: grasp remote Evaluation: 9.0 + Level 3: Thought: lift remote Evaluation: 8.0 + Level 2: Thought: grasp remote Evaluation: 9.0 + Level 3: Thought: lift remote Evaluation: 8.0 + Level 2: Thought: lift remote Evaluation: 8.0 + Level 1: Thought: Place the remote where the fork is Evaluation: 0.5 + Level 2: Thought: move to fork location Evaluation: 8.0 + Level 3: Thought: place remote Evaluation: 9.0 + Level 2: Thought: place remote Evaluation: 9.0 + Level 1: Thought: Pick up the fork Evaluation: 0.5 + Level 2: Thought: move to fork Evaluation: 8.0 + Level 3: Thought: grasp fork Evaluation: 9.0 + Level 3: Thought: lift fork Evaluation: 8.0 + Level 2: Thought: grasp fork Evaluation: 9.0 + Level 3: Thought: lift fork Evaluation: 8.0 + Level 2: Thought: lift fork Evaluation: 8.0 + Level 1: Thought: Place the fork where the remote was Evaluation: 0.5 + Level 2: Thought: move to remote location Evaluation: 8.0 + Level 3: Thought: place fork Evaluation: 9.0 + Level 2: Thought: place fork Evaluation: 9.0 + + Best Path: + Action: move forward + Action: grasp remote + Action: lift remote + Action: move to fork location + Action: place remote + Action: move to fork + Action: grasp fork + Action: lift fork + Action: move to remote location + Action: place fork +``` + +### The best action path can also be generated using the language_agent +```python +tree_of_thought.get_actions_with_llm() +``` + +### The structure of the actions in the tree +```python +Root +│ +├── Action: "Pick up the remote" +│ ├── Thought: "move arm to the right" (Evaluation: 0.9) +│ ├── Thought: "lower arm" (Evaluation: 0.9) +│ ├── Thought: "grasp remote" (Evaluation: 1.0) +│ └── Thought: "lift arm" (Evaluation: 0.9) +│ +├── Action: "Place the remote where the fork is" +│ ├── Thought: "lower arm" (Evaluation: 0.9) +│ ├── Thought: "grasp remote" (Evaluation: 1.0) +│ └── Thought: "lift arm" (Evaluation: 0.9) +│ +├── Action: "Pick up the fork" +│ ├── Thought: "grasp fork" (Evaluation: 1.0) +│ └── Thought: "lift arm" (Evaluation: 0.9) +│ +└── Action: "Place the fork where the remote was" + ├── Thought: "release fork" (Evaluation: 1.0) + └── Thought: "lift arm" (Evaluation: 0.9) +``` + +## Tree of Thought Components + +# ThoughtNode + +A `ThoughtNode` represents an individual decision/action in the reasoning process. It contains the following attributes: + +- **thought**: The actual action or decision. +- **embedding**: A high-dimensional representation of the thought (optional). +- **evaluation**: A score representing how promising the thought is. +- **children**: A list of child nodes (follow-up actions). +- **reduced_embedding**: Embedding reduced via PCA for optimization. + +### Methods + +- **`add_child`**: Adds a child node to the current thought node. +- **`is_leaf`**: Checks if the node has any children. + +## TreeOfThought + +The `TreeOfThought` manages the tree and recursively expands on thoughts using the `LanguageAgent`. It generates embeddings for thoughts and uses PCA to reduce the dimensionality of these embeddings. + +### Parameters + +- **language_agent**: The agent responsible for generating new thoughts based on the input instruction. +- **n_components**: Number of PCA components used to reduce the dimensionality of embeddings. +- **max_depth**: Maximum depth for the tree exploration. +- **embed**: Flag to indicate whether embeddings should be generated for thoughts. + +### Core Methods + +- **`generate_thoughts`**: Initializes the thought tree by querying the `LanguageAgent` with an instruction. +- **`get_actions`**: Retrieves the best action path in the thought tree using a combination of BFS and DFS. +- **`traverse`**: Traverses and prints the structure of the thought tree. + +## Embedding with PCA + +Each thought can be transformed into a high-dimensional embedding using a SentenceTransformer model. To optimize performance, PCA is used to reduce the embedding dimensionality: + +- **Embedding**: Captures the semantic meaning of a thought. +- **PCA Reduction**: Reduces the number of dimensions while retaining as much information as possible. + +### Example + +```python +thought_node = ThoughtNode("Find the optimal solution", embedding=embedding_vector, n_components=10) +``` + +## Thought pathfinding system + +This system explores a thought tree using a combination of **Breadth-First Search (BFS)** and **Depth-First Search (DFS)** to identify the most promising path based on evaluation scores and depth. + +The thought tree consists of nodes representing decisions or actions, each with an associated evaluation score. The goal of the system is to traverse the tree and find the optimal path, avoiding redundant or low-value thoughts unless they are terminal actions (leaf nodes). + +- Prioritizing **Depth-First Search (DFS)** to explore the deepest nodes in the tree first, collecting evaluated thoughts from the deepest strategies before backtracking. +- Using **Breadth-First Search (BFS)** within each level to ensure all possible actions at the current depth are explored before backtracking to higher levels. +- Skipping nodes with an evaluation of 0.5 unless they are leaf nodes. Nodes with a score of 0.5` are considered neutral and are skipped unless they represent terminal actions with no further decisions (leaf nodes). + +### Example of Best Path Calculation + +```python +tree_of_thought.get_actions(recompute=True) +``` + +### Visualization + +The **traverse** function enables visualization of the thought tree. It prints out the thought at each level and display the corresponding evaluation score. +```python +tree_of_thought.traverse() +``` + + + + diff --git a/mbodied/tree/__init__.py b/mbodied/tree/__init__.py new file mode 100644 index 0000000..c8a35eb --- /dev/null +++ b/mbodied/tree/__init__.py @@ -0,0 +1,4 @@ +from .prompt import generate_prompt, create_llm_prompt +from .tree_of_thought import ThoughtNode, TreeOfThought + +__all__ = ["generate_prompt", "create_llm_prompt", "ThoughtNode", "TreeOfThought"] \ No newline at end of file diff --git a/mbodied/tree/prompt.py b/mbodied/tree/prompt.py new file mode 100644 index 0000000..c409d91 --- /dev/null +++ b/mbodied/tree/prompt.py @@ -0,0 +1,100 @@ +from mbodied.types.message import Message +from mbodied.types.sample import Sample +from typing import Optional, List +from pydantic import Field + + +class Thought(Sample): + thought: List[str] = Field( + description="The actions to be taken by the robot" + ) + evaluation: Optional[List[float]] = Field( + description="The evaluation of the each thought. It can be a number between 0.1 and 1.0 being 0.1 the worst and 1.0 the best." + ) + +def generate_prompt(instruction: str) -> List[Message]: + """ + Generate the prompt to send to the LLM based on the instruction. + Parameter: + - instruction: The instruction from the user. + return: Context to send to the language agent. + """ + TREE_OF_THOUGHT_PROMPT = f""" + You are a robot command agent designed to solve complex tasks by generating, evaluating, and refining actions. Your task is to generate solutions, evaluate the quality of your thoughts, and return a response in a structured format. + - Use the action provided e.g "turn right" to generate an ACCURATE list of next steps for the robot to follow after the action in the format: ["move forward", "turn left", ...] based on the instruction: {instruction}. + - Generate only steps that follow the action given and not previous steps taken to get to the action while telling the robot exactly what to do. + - Return an empty list if there is no further action required to reach the goal from the action provided. + + ### Instructions: + + 1. **Understand the Action:** + - Analyze the initial action or thought provided. + - Break it down into smaller, manageable steps if necessary. + - Ensure you fully comprehend the context before proceeding. + + 2. **Generate New Actions:** + - Based on the provided thought, generate a LIST of new actions or reasoning steps. + - Ensure that the actions are logical, well-justified, and relevant to the initial problem. + + 3. **Self-Evaluation:** + - After generating each action, evaluate it for relevance, feasibility, safety, and efficiency. + - Strictly follow the instructions below to assign an evaluation score between 1 and 10: + - **1 to 4:** Flawed or irrelevant action. + - **5 to 7:** Actions that can be broken into further steps e.g "Place object on the table". + - **8 to 10:** Only actions that cannot be broken down any further e.g "move forward". + + 4. **Iterate on Actions:** + - Based on the actions generated, further reason and expand on each step. + - Re-evaluate new actions as necessary. + + Respond in the following json schema: + {Thought.model_json_schema()} + """ + """Please provide the response as a JSON array e.g + { + "thought": ["move forward", "turn left", "pick up apple"], + "evaluation": [8, 9, 6] + } + """ + + context = [ + Message( + role="user", + content=TREE_OF_THOUGHT_PROMPT, + ), + Message(role="assistant", content="Understood!"), + ] + + return context + +def create_llm_prompt(formatted_thought_tree: str, instruction: str) -> List[Message]: + """ + Create the prompt to send to the LLM based on the formatted thought tree. + Parameter: + - formatted_thought_tree: The formatted thought sequence tree. + - instruction: The instruction from the user. + return: Context to send to the language agent. + """ + prompt = f""" + Given the following thought tree, pick the best path through the actions that completes the goal in the instruction {instruction}. + The path should: + - Not pick duplicate actions/steps. + - Not skip important steps. + - Prioritize steps that cannot be broken down any further e.g "move forward" over "move to table". + - Prioritize lower-level actions as they are most likely not steps that can be broken down any further e.g "grasp object". + + Here is the thought tree: + {formatted_thought_tree} + + DO NOT PROVIDE ANY RESPONSE EXCEPT A LIST IN THE FOLLOWING FORMAT: + ["action1", "action2", "action3",...] + """ + + context = [ + Message( + role="user", + content=prompt, + ), + Message(role="assistant", content="Understood!"), + ] + + return context \ No newline at end of file diff --git a/mbodied/tree/tree_of_thought.py b/mbodied/tree/tree_of_thought.py new file mode 100644 index 0000000..2e4f793 --- /dev/null +++ b/mbodied/tree/tree_of_thought.py @@ -0,0 +1,400 @@ +import logging +import uuid +import re +from mbodied.agents import LanguageAgent +from mbodied.types.sense.vision import Image +from mbodied.tree.prompt import create_llm_prompt, generate_prompt, Thought +from typing import List, Optional, Callable +from sklearn.decomposition import PCA +from sentence_transformers import SentenceTransformer + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +class ThoughtNode: + def __init__( + self, + thought: str, + embedding = None, + parent: 'ThoughtNode' = None, + evaluation: Optional[float] = None, + n_components: int = 10 + ): + """ + Initialize a ThoughtNode representing a decision/action embedding. + + Parameters: + - thought: The action or decision. + - embedding: The high-dimensional embedding of the action. + - parent: The parent node if any. + - evaluation: The score of this thought. + - n_components: Number of components for reducing dimensionality. + """ + self.id = str(uuid.uuid4()) + self.embedding = embedding + self.reduced_embedding = None + self.children: List['ThoughtNode'] = [] + self.thought = thought + self.parent = parent + self.evaluation = evaluation if evaluation is not None else 0.5 + if embedding is not None: + self.pca_model = PCA(n_components=n_components) + self.reduced_embedding = self.pca_model.fit_transform([embedding])[0] + + def __repr__(self): + return f"Thought: {self.thought}, Eval: {self.evaluation}" + + def add_child(self, child_node: 'ThoughtNode') -> 'ThoughtNode': + """ + Add a child node representing the next decision in the thought process. + + Parameters: + - child_node: The ThoughtNode to be added to the children + """ + self.children.append(child_node) + logger.info(f"Added child node {child_node.id}") + return child_node + + def is_leaf(self): + return len(self.children) == 0 + + +class TreeOfThought: + def __init__( + self, + language_agent: 'LanguageAgent', + n_components: int = 10, + max_depth: int = 3, + embed: bool = False + ): + """ + For reference, see: https://arxiv.org/pdf/2305.10601 + Initialize a Tree of Thought with a root node. + + Parameters: + - language_agent: The LanguageAgent used to generate thoughts/actions. + - n_components: Number of components/dimensions for PCA to reduce embeddings. + - max_depth: Maximum depth allowed in the tree. + - embed: Whether to generate embeddings for thoughts + """ + self.root = None + self.instruction = None + self.image = None + self.language_agent = language_agent + self.n_components = n_components + self.max_depth = max_depth + self.embed = embed + self.embedding_model = SentenceTransformer('paraphrase-MiniLM-L6-v2') + self.best_path = [] + self.llm_best_path = [] + + def generate_thoughts( + self, + instruction, + image = None, + parse_function: Optional[Callable[[str], List[str]]] = None + ): + """ + Query the Language Agent to generate actions/thoughts based on an instruction. + + Parameters: + - instruction: The instruction to be sent to the Language Agent. + - depth: The current depth of the thought tree (defaults to 0). + - parse_function: A user-defined function to process the output of the LanguageAgent. + This function takes the raw output as input and returns a list of thoughts or child states. + - image: An image passed to the language agent. + + Returns: + - A list of actions generated by the Language Agent. + """ + self.instruction = instruction + self.image = image + try: + response = self.language_agent.act(instruction, image) + except Exception as e: + logger.error(f"Language agent failed: {e}") + return + response = response.replace('```python\n', '').replace('\n```', '') + actions = parse_function(response) if parse_function else [step.strip().strip('"').strip("'") for step in response.strip('[]').split(',')] + self.language_agent.context = generate_prompt(instruction) + logger.info(f"Generated actions: {actions}") + self.root = ThoughtNode("Start") + self._recursive_reasoning( + node = self.root, + thoughts = actions, + evaluations = None, + depth = 0, + image = image + ) + + def _recursive_reasoning( + self, + node: 'ThoughtNode', + thoughts: List[str], + evaluations: List[float], + depth: int, + image: Image + ): + """ + Perform recursive reasoning by generating further thoughts from the current node. + + Parameters: + - node: The current ThoughtNode to reason from. + - depth: Current depth in the tree. + """ + if depth >= self.max_depth or len(thoughts) == 0: + logger.info("Max depth reached, stopping recursion.") + return + + for i, thought in enumerate(thoughts): + evaluation = evaluations[i] if evaluations else None + embedding = None + if self.embed: + embedding = self._generate_embedding(thought) + + child_node = ThoughtNode( + thought=thought, + embedding=embedding, + parent=node, + evaluation=evaluation, + n_components=self.n_components + ) + child_node = node.add_child(child_node) + + try: + response = self.language_agent.act_and_parse(thought, image, Thought) + except Exception as e: + logger.error(f"Language agent failed: {e}") + return + logger.info(f"Generated thought: {response.thought}, Evaluation: {response.evaluation}") + + self._recursive_reasoning( + node = child_node, + thoughts = response.thought, + evaluations = response.evaluation, + depth = depth + 1, + image=image, + ) + + def _generate_embedding(self, thought: str): + """ + Generate an embedding for a given thought using a pre-trained SentenceTransformer. + + Parameters: + - thought: The thought/action to generate an embedding for. + + Returns: + - A high-dimensional embedding for the thought. + """ + embedding = self.embedding_model.encode(thought, convert_to_tensor=False) + return embedding + + def get_actions(self, recompute: bool = False): + """ + Fetch best action path + + parameters: + - recompute: Boolean value which determines whether to recalculate best path + """ + if recompute or len(self.best_path) == 0: + best_path = self._traverse(self.root) + self.best_path = best_path + print("Best Path:") + for action in self.best_path: + print(f"Action: {action}") + + def get_actions_from_llm(self, recompute: bool = False): + """ + Generate the best path based on the given ThoughtNode tree using the LLM. + parameters: + - recompute: Boolean value which determines whether to recalculate best path + """ + if recompute or len(self.llm_best_path) == 0: + formatted_tree = self._format_thought_tree(self.root) + self.language_agent.context = create_llm_prompt(formatted_tree, self.instruction) + try: + response = self.language_agent.act(self.instruction, self.image) + except Exception as e: + logger.error(f"Language agent failed: {e}") + return + response = response.replace('```python\n', '').replace('\n```', '') + match = re.search(r'\[(.*?)\]', response) + if match: + command_list = match.group(0) + self.llm_best_path = [step.strip().strip('"').strip("'") for step in command_list[1:-1].split(',')] + + print("Best Path:") + for action in self.llm_best_path: + print(f"Action: {action}") + + def get_depth(self, node: ThoughtNode) -> int: + """ + Compute the depth of a node in the tree. + + Parameters: + - node: The ThoughtNode to compute depth for. + + Returns: + - Integer representing the depth. + """ + depth = 0 + current = node + while current != self.root: + depth += 1 + current = current.parent + return depth + + def traverse(self, node=None, level=0): + """ + Traverse the tree of thought and print the structure. + + Parameters: + - node: The current ThoughtNode to start traversal from. Defaults to root. + - level: Depth level of the node in the tree. + """ + if node is None: + node = self.root + + print(" " * level + f"Level {level}: Thought: {node.thought} Evaluation: {node.evaluation}") + + for child in node.children: + self.traverse(child, level + 1) + + def _action_exists(self, action: str, path: List[str]) -> bool: + """ + Checks if the action already exists in the current path to avoid redundant steps. + + Args: + action (str): The action to check. + path (List[str]): The current path containing previously selected actions. + + Returns: + bool: True if the action exists in the path, False otherwise. + """ + return any(action in step for step in path) + + def _traverse(self, node: ThoughtNode) -> List[str]: + """ + Performs depth-first traversal but avoids adding nodes with evaluation 0.5 unless they are leaf nodes. + Gathers all children at each level, processes the breadth, and backtracks when required. + + Args: + node (ThoughtNode): The root node of the thought tree. + + Returns: + List[str]: The ordered list of actions based on depth-first traversal rules. + """ + result = [] + visited = set() + + self._dfs_collect(node, result, visited) + + return result + + def _dfs_collect(self, node: ThoughtNode, result: List[str], visited: set): + """ + Helper function to perform the DFS and collect thoughts as per the given rules. + + Args: + node (ThoughtNode): Current node being traversed. + result (List[str]): The result list where actions are being added. + visited (set): Tracks visited thoughts to avoid duplicates. + """ + if node.evaluation == 0.5 and not node.is_leaf(): + for child in node.children: + self._dfs_collect(child, result, visited) + else: + if node.thought not in visited: + visited.add(node.thought) + result.append(node.thought) + + for child in node.children: + self._dfs_collect(child, result, visited) + + def _format_thought_tree(self, node: ThoughtNode, level: int = 0) -> str: + """ + Format the tree structure of thoughts into a readable string for the LLM. + Parameters: + - node: The root node of the thought tree. + - level: The current level of depth in the tree (used for formatting). + return: A string representing the thought tree. + """ + formatted_str = f"{' ' * level}Level {level}: Thought: {node.thought} Evaluation: {node.evaluation}\n" + for child in node.children: + formatted_str += self._format_thought_tree(child, level + 1) + + return formatted_str + +if __name__ == "__main__": + def run_thought_tree(instruction: str, language_agent, n_components: int = 3, max_depth: int = 3, image=None): + """ + Initializes and runs the Tree of Thought process for a given instruction. + + Parameters: + - instruction: The problem or task for the agent to reason about. + - initial_embedding: The embedding for the root thought. + - language_agent: The agent responsible for generating thoughts. + - n_components: Number of components for PCA (default 3). + - max_depth: Maximum depth of the thought tree (default 5). + + Returns: + - The final tree of thought nodes. + """ + thought_tree = TreeOfThought(language_agent=language_agent, n_components=n_components, max_depth=max_depth) + + thought_tree.generate_thoughts(instruction, image) + + return thought_tree + + import os + + image = Image(path="resources/color_image.png") + cognition = LanguageAgent( + context="You are an embodied planner that responds with a python list of strings and nothing else.", + api_key=os.getenv("OPENAI_API_KEY"), + model_src="openai", + recorder="auto", + ) + + # Example Usage + thought_tree = run_thought_tree("switch the positions of the remote and the fork", language_agent=cognition, image=image) + thought_tree.traverse() + thought_tree.get_actions() + + """ + Level 0: Thought: Start Evaluation: 0.5 + Level 1: Thought: Pick up the remote Evaluation: 0.5 + Level 2: Thought: move forward Evaluation: 8.0 + Level 3: Thought: grasp remote Evaluation: 9.0 + Level 3: Thought: lift remote Evaluation: 8.0 + Level 2: Thought: grasp remote Evaluation: 9.0 + Level 3: Thought: lift remote Evaluation: 8.0 + Level 2: Thought: lift remote Evaluation: 8.0 + Level 1: Thought: Place the remote where the fork is Evaluation: 0.5 + Level 2: Thought: move to fork location Evaluation: 8.0 + Level 3: Thought: place remote Evaluation: 9.0 + Level 2: Thought: place remote Evaluation: 9.0 + Level 1: Thought: Pick up the fork Evaluation: 0.5 + Level 2: Thought: move to fork Evaluation: 8.0 + Level 3: Thought: grasp fork Evaluation: 9.0 + Level 3: Thought: lift fork Evaluation: 8.0 + Level 2: Thought: grasp fork Evaluation: 9.0 + Level 3: Thought: lift fork Evaluation: 8.0 + Level 2: Thought: lift fork Evaluation: 8.0 + Level 1: Thought: Place the fork where the remote was Evaluation: 0.5 + Level 2: Thought: move to remote location Evaluation: 8.0 + Level 3: Thought: place fork Evaluation: 9.0 + Level 2: Thought: place fork Evaluation: 9.0 + + Best Path: + Action: move forward + Action: grasp remote + Action: lift remote + Action: move to fork location + Action: place remote + Action: move to fork + Action: grasp fork + Action: lift fork + Action: move to remote location + Action: place fork + """ \ No newline at end of file diff --git a/tests/test_tree_of_thought.py b/tests/test_tree_of_thought.py new file mode 100644 index 0000000..081ee99 --- /dev/null +++ b/tests/test_tree_of_thought.py @@ -0,0 +1,61 @@ +from unittest import mock +from mbodied.agents.language.language_agent import LanguageAgent +from mbodied.tree.tree_of_thought import TreeOfThought, ThoughtNode + +@mock.patch("mbodied.agents.language.language_agent.LanguageAgent.act", side_effect=[ + "['Pick up the remote', 'Place the remote where the fork is', 'Pick up the fork', 'Place the fork where the remote was']", + "['move forward', 'grasp remote']", + "['move forward', 'grasp remote']", + "[]", + "['move to fork location', 'place remote']", + "['place remote']", + "[]", + "['grasp fork']", + "[]" + "['move to remotes original location', 'place fork']", + "['place fork']", + "[]" + +]) +def test_generate_thoughts(mock_act): + + language_agent = mock.Mock() + + language_agent.act.side_effect = [ + "['Pick up the remote', 'Place the remote where the fork is', 'Pick up the fork', 'Place the fork where the remote was']", + "['move forward', 'grasp remote']", + "['move forward', 'grasp remote']", + "[]", + "['move to fork location', 'place remote']", + "['place remote']", + "[]", + "['grasp fork']", + "[]" + "['move to remotes original location', 'place fork']", + "['place fork']", + "[]" + ] + + language_agent = LanguageAgent() + tot = TreeOfThought(language_agent=language_agent, max_depth=2) + tot.generate_thoughts(instruction="pick up spoon") + assert tot.root is not None, "Expected root node" + assert tot.root.thought == "Start", f"Expected Start but got {tot.root.thought}" + tot.traverse() + tot.get_actions() + +def test_add_child(): + parent_node = ThoughtNode("Parent") + child_node = ThoughtNode("Child") + parent_node.add_child(child_node) + + assert parent_node.children[0] == child_node, f"Expected {parent_node} to contain {child_node}" + +def test_is_leaf(): + leaf_node = ThoughtNode("Leaf") + assert leaf_node.is_leaf(), "Expected to be a leaf node" + + non_leaf_node = ThoughtNode("Non-leaf") + non_leaf_node.add_child(ThoughtNode("Child")) + assert not non_leaf_node.is_leaf(), "Expected to have children nodes" +