| 4 | |
| 5 | |
| 6 | class DAGManager: |
| 7 | NODE_TYPES = {"cookie", "master", "cURL"} |
| 8 | |
| 9 | def __init__(self): |
| 10 | self.graph = nx.DiGraph() |
| 11 | self.root_id = None |
| 12 | def add_node( |
| 13 | self, |
| 14 | node_type: Literal["cookie", "master", "cURL", "not found"], |
| 15 | content: Optional[dict] = None, |
| 16 | dynamic_parts: Optional[List[str]] = None, |
| 17 | extracted_parts: Optional[List[str]] = None, |
| 18 | input_variables: Optional[Dict[str, str]] = None, |
| 19 | ): |
| 20 | node_id = str(uuid.uuid4()) |
| 21 | self.graph.add_node(node_id, node_type=node_type, content=content, dynamic_parts=dynamic_parts, extracted_parts=extracted_parts, input_variables=input_variables) |
| 22 | return node_id |
| 23 | |
| 24 | def update_node( |
| 25 | self, |
| 26 | node_id: str, |
| 27 | **attributes: Optional[List[str]]): |
| 28 | |
| 29 | for attr, value in attributes.items(): |
| 30 | if value is not None: |
| 31 | self.graph.nodes[node_id][attr] = value |
| 32 | |
| 33 | def detect_cycles(self): |
| 34 | """ |
| 35 | Detects if there are cycles in the DAG managed by this class. |
| 36 | If a cycle is found, it returns the list of nodes involved in the cycle. |
| 37 | If no cycle is found, it returns None. |
| 38 | |
| 39 | Returns: |
| 40 | - A list of nodes forming a cycle, or None if no cycles are found. |
| 41 | """ |
| 42 | try: |
| 43 | cycle = list(nx.find_cycle(self.graph, orientation='original')) |
| 44 | print("Cycle detected:") |
| 45 | return cycle |
| 46 | except nx.exception.NetworkXNoCycle: |
| 47 | return None |
| 48 | |
| 49 | def get_node(self, node_id: str) -> Optional[Dict]: |
| 50 | """ |
| 51 | Retrieves the attributes of the specified node. |
| 52 | |
| 53 | :param node_id: ID of the node to retrieve. |
| 54 | :return: Dictionary of node attributes or None if the node does not exist. |
| 55 | """ |
| 56 | return self.graph.nodes.get(node_id, None) |
| 57 | |
| 58 | def add_edge(self, from_node_id: str, to_node_id: str): |
| 59 | self.graph.add_edge(from_node_id, to_node_id) |
| 60 | |
| 61 | def __str__(self): |
| 62 | nodes_info = [] |
| 63 | for node_id in self.graph.nodes: |