Spaces:
Paused
Paused
| 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 | |