Add tree structuring to task and trace api functions

Author: monoxgas

dreadnode/api/client.py

@@ -7,7 +7,12 @@
 from pydantic import BaseModel
 from ulid import ULID
 
-from dreadnode.api.util import process_run, process_task
+from dreadnode.api.util import (
+    convert_flat_tasks_to_tree,
+    convert_flat_trace_to_tree,
+    process_run,
+    process_task,
+)
 from dreadnode.util import logger
 from dreadnode.version import VERSION
 
@@ -20,9 +25,11 @@
     RunSummary,
     StatusFilter,
     Task,
+    TaskTree,
     TimeAggregationType,
     TimeAxisType,
     TraceSpan,
+    TraceTree,
     UserDataCredentials,
 )
 
@@ -134,13 +141,41 @@ def _get_run(self, run: str | ULID) -> RawRun:
     def get_run(self, run: str | ULID) -> Run:
         return process_run(self._get_run(run))
 
-    def get_run_tasks(self, run: str | ULID) -> list[Task]:
+    TraceFormat = t.Literal["tree", "flat"]
+
+    @t.overload
+    def get_run_tasks(
+        self, run: str | ULID, *, format: t.Literal["tree"] = "tree"
+    ) -> list[TaskTree]: ...
+
+    @t.overload
+    def get_run_tasks(
+        self, run: str | ULID, *, format: t.Literal["flat"] = "flat"
+    ) -> list[Task]: ...
+
+    def get_run_tasks(
+        self, run: str | ULID, *, format: TraceFormat = "flat"
+    ) -> list[Task] | list[TaskTree]:
         raw_run = self._get_run(run)
         response = self.request("GET", f"/strikes/projects/runs/{run!s}/tasks/full")
         raw_tasks = [RawTask(**task) for task in response.json()]
-        return [process_task(task, raw_run) for task in raw_tasks]
-
-    def get_run_trace(self, run: str | ULID) -> list[Task | TraceSpan]:
+        tasks = [process_task(task, raw_run) for task in raw_tasks]
+        tasks = sorted(tasks, key=lambda x: x.timestamp)
+        return tasks if format == "flat" else convert_flat_tasks_to_tree(tasks)
+
+    @t.overload
+    def get_run_trace(
+        self, run: str | ULID, *, format: t.Literal["tree"] = "tree"
+    ) -> list[TraceTree]: ...
+
+    @t.overload
+    def get_run_trace(
+        self, run: str | ULID, *, format: t.Literal["flat"] = "flat"
+    ) -> list[Task | TraceSpan]: ...
+
+    def get_run_trace(
+        self, run: str | ULID, *, format: TraceFormat = "flat"
+    ) -> list[Task | TraceSpan] | list[TraceTree]:
         raw_run = self._get_run(run)
         response = self.request("GET", f"/strikes/projects/runs/{run!s}/spans/full")
         trace: list[Task | TraceSpan] = []
@@ -149,7 +184,9 @@ def get_run_trace(self, run: str | ULID) -> list[Task | TraceSpan]:
                 trace.append(process_task(RawTask(**item), raw_run))
             else:
                 trace.append(TraceSpan(**item))
-        return trace
+
+        trace = sorted(trace, key=lambda x: x.timestamp)
+        return trace if format == "flat" else convert_flat_trace_to_tree(trace)
 
     # Data exports
 

dreadnode/api/models.py

@@ -285,11 +285,9 @@ class TaskTree(BaseModel):
     children: list["TaskTree"] = []
 
 
-class SpanTree(BaseModel):
-    """Tree representation of a trace span with its children"""
-
+class TraceTree(BaseModel):
     span: Task | TraceSpan
-    children: list["SpanTree"] = []
+    children: list["TraceTree"] = []
 
 
 # User data credentials

dreadnode/api/util.py

@@ -9,12 +9,49 @@
     RawTask,
     Run,
     Task,
+    TaskTree,
+    TraceSpan,
+    TraceTree,
     V0Object,
 )
 
 logger = getLogger(__name__)
 
 
+def convert_flat_tasks_to_tree(tasks: list[Task]) -> list[TaskTree]:
+    tree_nodes: dict[str, TaskTree] = {}
+    for task in tasks:
+        tree_nodes[task.span_id] = TaskTree(task=task)
+
+    roots: list[TaskTree] = []
+    for task in tasks:
+        if task.parent_task_span_id not in tree_nodes:
+            roots.append(tree_nodes[task.span_id])
+        else:
+            parent_node = tree_nodes.get(task.parent_task_span_id)
+            if parent_node:
+                parent_node.children.append(tree_nodes[task.span_id])
+
+    return roots
+
+
+def convert_flat_trace_to_tree(trace: list[Task | TraceSpan]) -> list[TraceTree]:
+    tree_nodes: dict[str, TraceTree] = {}
+    for span in trace:
+        tree_nodes[span.span_id] = TraceTree(span=span)
+
+    roots: list[TraceTree] = []
+    for span in trace:
+        if span.parent_span_id not in tree_nodes:
+            roots.append(tree_nodes[span.span_id])
+        else:
+            parent_node = tree_nodes.get(span.parent_span_id)
+            if parent_node:
+                parent_node.children.append(tree_nodes[span.span_id])
+
+    return roots
+
+
 def process_run(run: RawRun) -> Run:
     inputs: dict[str, Object] = {}
     outputs: dict[str, Object] = {}