Toggle Menu
Prev Up Next

# Source code for langchain_experimental.tot.controller

    from typing import Tuple

    from langchain_experimental.tot.memory import ToTDFSMemory
    from langchain_experimental.tot.thought import ThoughtValidity

    class ToTController:
        """
        Tree of Thought (ToT) controller.

        This is a version of a ToT controller, dubbed in the paper as a "Simple
        Controller".

        It has one parameter `c` which is the number of children to explore for each
        thought.
        """

        def __init__(self, c: int = 3):
            """
            Initialize the controller.

            Args:
                c: The number of children to explore at each node.
            """
            self.c = c

        def __call__(self, memory: ToTDFSMemory) -> Tuple[str, ...]:
            next_thought = memory.top()
            parent_thought = memory.top_parent()
            validity = (
                ThoughtValidity.VALID_INTERMEDIATE
                if next_thought is None
                else next_thought.validity
            )

            # 1 if the current partial solution is invalid, backtrack to the parent
            # thought.
            if validity == ThoughtValidity.INVALID:
                memory.pop()
                next_thought = memory.top()
                if next_thought and len(next_thought.children) >= self.c:
                    memory.pop()

            # 2 if the current partial solution is valid but C children were
            # explored and yet failed to find a final solution, backtrack to the
            # parent thought.
            elif (
                validity == ThoughtValidity.VALID_INTERMEDIATE
                and parent_thought
                and len(parent_thought.children) >= self.c
            ):
                memory.pop(2)

            return tuple(thought.text for thought in memory.current_path())