mirror of
https://github.com/hwchase17/langchain
synced 2024-10-31 15:20:26 +00:00
e7e5cb9d08
# [WIP] Tree of Thought introducing a new ToTChain. This PR adds a new chain called ToTChain that implements the ["Large Language Model Guided Tree-of-Though"](https://arxiv.org/pdf/2305.08291.pdf) paper. There's a notebook example `docs/modules/chains/examples/tot.ipynb` that shows how to use it. Implements #4975 ## Who can review? Community members can review the PR once tests pass. Tag maintainers/contributors who might be interested: - @hwchase17 - @vowelparrot --------- Co-authored-by: Vadim Gubergrits <vgubergrits@outbox.com> Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import List, Optional
|
|
|
|
from langchain_experimental.tot.thought import Thought
|
|
|
|
|
|
class ToTDFSMemory:
|
|
"""
|
|
Memory for the Tree of Thought (ToT) chain. 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 None
|
|
|
|
def 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 None
|
|
for _ in range(n):
|
|
node = self.stack.pop()
|
|
return node
|
|
|
|
def 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 None
|
|
|
|
def 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)
|
|
|
|
@property
|
|
def 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[:]
|