Spaces:
Paused
Paused
File size: 3,045 Bytes
fced70d 70f0982 |
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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import copy
from typing import Any, Dict, List, Optional, Self
import uuid
class ResearchNode:
def __init__(self, query: str = "_", parent: Optional[Self] = None, depth: int = 0):
self.parent = parent
self.id = str(uuid.uuid4())
self.query = query
self.depth = depth
self.children: List[ResearchNode] = []
self.data: List[Dict[str, Any]] = []
def find_node(self, node_id: str) -> Optional[Self]:
"""
Returns the node with the given id.
If not found, returns None.
"""
if self.id == node_id:
return self
for child in self.children:
found = child.find_node(node_id)
if found:
return found
return None
def add_child(self, query: str, node: Optional[Self] = None) -> Self:
if node:
child = node
child.parent = self
child.depth = self.depth + 1
else:
child = ResearchNode(query, parent=self, depth=self.depth + 1)
self.children.append(child)
return child
def get_path_to_root(self) -> List[str]:
"""
Returns the path from this node to the root node.
List[str]: [root.query, ..., self.query]
"""
path = [self.query]
current = self
while current.parent:
current = current.parent
path.append(current.query)
return list(reversed(path))
def max_depth(self) -> int:
if not self.children:
return self.depth
return max([child.max_depth() for child in self.children])
def total_children(self) -> int:
if not self.children:
return 0
return len(self.children) + sum([child.total_children() for child in self.children])
def get_all_data(self) -> List[Dict[str, Any]]:
data = copy.deepcopy(self.data)
for child in self.children:
data.extend(child.get_all_data())
return data
# Build research tree structure
def build_tree_structure(self) -> Dict:
if not self:
return {}
sources = {d["url"]: d["text"] for d in self.data if d.get("url") and d.get("text")}
return {
"query": self.query,
"depth": self.depth,
"sources": sources,
"children": [child.build_tree_structure() for child in self.children],
}
# Return deep copy with node pointers | Isolated function
def deep_copy_tree(root: Optional[Self] = None) -> Self:
"""
Returns a deep copy of the tree starting from this node.
"""
if root is None:
return None
new_node = ResearchNode(root.query, depth=root.depth)
new_node.id = root.id
new_node.data = copy.deepcopy(root.data)
for child in root.children:
new_child = ResearchNode.deep_copy_tree(child)
new_child.parent = new_node
new_node.children.append(new_child)
return new_node
|