Upload runtime_state_manager.py

examples/runtime_state_manager.py

@@ -0,0 +1,75 @@
+from aworld.runners.state_manager import RuntimeStateManager, RunNodeBusiType, RunNodeStatus, RunNode
+from typing import List
+
+
+def test_runtime_state_manager():
+    state_manager = RuntimeStateManager()
+    session_id = "1"
+
+    state_manager.create_node(busi_type=RunNodeBusiType.TASK,
+                              busi_id="1", session_id=session_id, msg_id="1")
+
+    state_manager.run_node(busi_type=RunNodeBusiType.TASK, busi_id="1")
+    node = state_manager.get_node(busi_type=RunNodeBusiType.TASK, busi_id="1")
+    assert node.status == RunNodeStatus.RUNNING
+
+    state_manager.break_node(busi_type=RunNodeBusiType.TASK, busi_id="1")
+    node = state_manager.get_node(busi_type=RunNodeBusiType.TASK, busi_id="1")
+    assert node.status == RunNodeStatus.BREAKED
+
+    state_manager.run_succeed(busi_type=RunNodeBusiType.TASK, busi_id="1")
+    node = state_manager.get_node(busi_type=RunNodeBusiType.TASK, busi_id="1")
+    assert node.status == RunNodeStatus.SUCCESS
+
+    state_manager.create_node(busi_type=RunNodeBusiType.TASK,
+                              busi_id="2", session_id=session_id, msg_id="2", msg_from="1")
+
+    state_manager.run_node(busi_type=RunNodeBusiType.TASK, busi_id="2")
+    state_manager.run_failed(busi_type=RunNodeBusiType.TASK, busi_id="2")
+    node = state_manager.get_node(busi_type=RunNodeBusiType.TASK, busi_id="2")
+    assert node.status == RunNodeStatus.FAILED
+
+    state_manager.create_node(busi_type=RunNodeBusiType.TASK,
+                              busi_id="3", session_id=session_id, msg_id="3", msg_from="1")
+    state_manager.run_node(busi_type=RunNodeBusiType.TASK, busi_id="3")
+    state_manager.run_timeout(busi_type=RunNodeBusiType.TASK, busi_id="3")
+    node = state_manager.get_node(busi_type=RunNodeBusiType.TASK, busi_id="3")
+    assert node.status == RunNodeStatus.TIMEOUNT
+
+    state_manager.create_node(busi_type=RunNodeBusiType.TASK,
+                              busi_id="4", session_id=session_id, msg_id="4", msg_from="3")
+    state_manager.run_succeed(busi_type=RunNodeBusiType.TASK, busi_id="1")
+
+    nodes = state_manager.get_nodes(session_id=session_id)
+    build_run_flow(nodes)
+
+
+def build_run_flow(nodes: List[RunNode]):
+    graph = {}
+    start_nodes = []
+
+    for node in nodes:
+        if hasattr(node, 'parent_node_id') and node.parent_node_id:
+            if node.parent_node_id not in graph:
+                graph[node.parent_node_id] = []
+            graph[node.parent_node_id].append(node.node_id)
+        else:
+            start_nodes.append(node.node_id)
+
+    for start in start_nodes:
+        print("-----------------------------------")
+        _print_tree(graph, start, "", True)
+        print("-----------------------------------")
+
+
+def _print_tree(graph, node_id, prefix, is_last):
+    print(prefix + ("└── " if is_last else "├── ") + node_id)
+    if node_id in graph:
+        children = graph[node_id]
+        for i, child in enumerate(children):
+            _print_tree(graph, child, prefix +
+                        ("    " if is_last else "│   "), i == len(children)-1)
+
+
+if __name__ == "__main__":
+    test_runtime_state_manager()