Spaces:
Paused
Paused
File size: 2,017 Bytes
02298d2 b636e8f 9155a62 2d96b3b b636e8f 2d96b3b b636e8f 62283c0 2d96b3b 0a3d9b7 2d96b3b 73fba58 2d96b3b 54a7d14 afce8a8 73fba58 54a7d14 88139f0 54a7d14 9155a62 27a07a9 b636e8f 27a07a9 b636e8f 27a07a9 |
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 |
import copy
from typing import Any, Dict, List, Optional, Self
class ResearchNode:
def __init__(self, query: str = "_", parent: Optional[Self] = None, depth: int = 0):
self.query = query
self.parent = parent
self.depth = depth
self.children: List[ResearchNode] = []
self.data: List[Dict[str, Any]] = []
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],
}
|