本教程将介绍 tree of thought 源码 中的Thought 和ToTDFSMemory 类

定义思维有效性

使用Enum模块来定义思维的有效性。

from enum import Enumclass ThoughtValidity(Enum):"""Enum for the validity of a thought."""VALID_INTERMEDIATE = 0VALID_FINAL = 1INVALID = 2

定义思维节点

使用Pydantic的BaseModel来定义思维节点。

from __future__ import annotations
from typing import Set
from pydantic import BaseModel, Fieldclass Thought(BaseModel):"""A thought in the ToT."""text: strvalidity: ThoughtValiditychildren: Set[Thought] = Field(default_factory=set)def __hash__(self) -> int:return id(self)

测试代码

def test_thought():thought1 = Thought(text="This is a valid thought.", validity=ThoughtValidity.VALID_FINAL)thought2 = Thought(text="This is first subvalid thought.", validity=ThoughtValidity.VALID_INTERMEDIATE)thought3 = Thought(text="This is the second subvalid thought.", validity=ThoughtValidity.VALID_INTERMEDIATE)# 添加子思维thought1.children.add(thought2)thought1.children.add(thought3)# 输出thought1的内容print(thought1.model_dump_json())  # 检查子思维print("Children of thought1:")for child in thought1.children:print(child.model_dump_json())  # 测试哈希print(f"Hash of thought1: {hash(thought1)}")print(f"Hash of thought2: {hash(thought2)}")print(f"Hash of thought3: {hash(thought3)}")test_thought()

输出：

{"text":"This is a valid thought.","validity":1,"children":[{"text":"This is first subvalid thought.","validity":0,"children":[]},{"text":"This is the second subvalid thought.","validity":0,"children":[]}]}
Children of thought1:
{"text":"This is first subvalid thought.","validity":0,"children":[]}
{"text":"This is the second subvalid thought.","validity":0,"children":[]}
Hash of thought1: 2202338415536
Hash of thought2: 2202338415936
Hash of thought2: 2202338423216

构建ToT DFS记忆结构

定义一个用于深度优先搜索的记忆结构。

from typing import List, Optionalclass ToTDFSMemory:"""Memory for the Tree of Thought (ToT) chain.It is implemented as a stack of thoughts. This allows for a depth first search (DFS) of the ToT."""def __init__(self, stack: Optional[List[Thought]] = None):self.stack: List[Thought] = stack or []def top(self) -> Optional[Thought]:"Get the top of the stack without popping it."return self.stack[-1] if len(self.stack) > 0 else Nonedef pop(self, n: int = 1) -> Optional[Thought]:"Pop the top n elements of the stack and return the last one."if len(self.stack) < n:return Nonefor _ in range(n):node = self.stack.pop()return nodedef top_parent(self) -> Optional[Thought]:"Get the parent of the top of the stack without popping it."return self.stack[-2] if len(self.stack) > 1 else Nonedef store(self, node: Thought) -> None:"Add a node on the top of the stack."if len(self.stack) > 0:self.stack[-1].children.add(node)self.stack.append(node)@propertydef level(self) -> int:"Return the current level of the stack."return len(self.stack)def current_path(self) -> List[Thought]:"Return the thoughts path."return self.stack[:]

测试ToT DFS记忆结构

def test_tot_dfs_memory():# 创建一些思维实例thought1 = Thought(text="Root thought", validity=ThoughtValidity.VALID_FINAL)thought2 = Thought(text="Child thought 1", validity=ThoughtValidity.VALID_INTERMEDIATE)thought3 = Thought(text="Child thought 2", validity=ThoughtValidity.VALID_INTERMEDIATE)thought4 = Thought(text="Grandchild thought", validity=ThoughtValidity.VALID_FINAL)# 初始化 ToTDFSMemorymemory = ToTDFSMemory()# 存储思维memory.store(thought1)print(f"Stored: {thought1.text}, Current level: {memory.level}")memory.store(thought2)print(f"Stored: {thought2.text}, Current level: {memory.level}")memory.store(thought3)print(f"Stored: {thought3.text}, Current level: {memory.level}")# 查看顶部思维top_thought = memory.top()print(f"Top thought: {top_thought.text if top_thought else 'None'}")# 存储一个子思维memory.store(thought4)print(f"Stored: {thought4.text}, Current level: {memory.level}")# 查看当前路径current_path = memory.current_path()print("Current path:")for thought in current_path:print(f"- {thought.text}")# 获取父思维parent_thought = memory.top_parent()print(f"Parent of top thought: {parent_thought.text if parent_thought else 'None'}")# 弹出思维popped_thought = memory.pop()print(f"Popped thought: {popped_thought.text if popped_thought else 'None'}, Current level: {memory.level}")# 再次查看当前路径current_path = memory.current_path()print("Current path after popping:")for thought in current_path:print(f"- {thought.text}")test_tot_dfs_memory()

输出：

Stored: Root thought, Current level: 1
Stored: Child thought 1, Current level: 2
Stored: Child thought 2, Current level: 3
Top thought: Child thought 2
Stored: Grandchild thought, Current level: 4
Current path:
- Root thought
- Child thought 1
- Child thought 2
- Grandchild thought
Parent of top thought: Child thought 2
Popped thought: Grandchild thought, Current level: 3
Current path after popping:
- Root thought
- Child thought 1
- Child thought 2

参考链接：https://github.com/langchain-ai/langchain-experimental/blob/main/libs/experimental/langchain_experimental/tot/thought.py
如果有任何问题，欢迎在评论区提问。