Spaces:
Runtime error
Runtime error
File size: 1,467 Bytes
ed4d993 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
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.
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 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[:]
|