fix: correct end_nanos accounting for async model calls

This at least fixes the block timestamps for async model calls. i don’t think it’s possible to do this generally (at least with my understanding of python futures), because the pdl_lazy logic obscures the original asyncio.Future… but this future is available at the point of fire for the async model calls… so i’ve fixed that part.

It’s imperfect, because e.g. an outer Text block… doesn’t know that there are Futures under the hood. but still this seems like an improvement over what we had. and that might be something we can patch over in the UI.

This includes said patch to the UI (see src/view/timeline/model.ts, the comment that references #683)
Signed-off-by: Nick Mitchell <[email protected]>

Author: starpit

pdl-live-react/src/view/timeline/model.ts

@@ -75,12 +75,25 @@ function computeModelIter(
         block,
       }
 
-  return [
-    ...(ignoreRoot ? [] : [root]),
-    ...childrenOf(block)
-      .filter(nonNullable)
-      .flatMap((child) => computeModelIter(child, root)),
-  ].filter(nonNullable)
+  const childrenModel = childrenOf(block)
+    .filter(nonNullable)
+    .flatMap((child) => computeModelIter(child, root))
+
+  // Correct for anomalies in the trace where a child may have an
+  // earlier end timestamp than its children. See
+  // https://github.com/IBM/prompt-declaration-language/pull/683
+  if (root) {
+    const maxEnd = childrenModel.reduce(
+      (maxEnd, child) => Math.max(maxEnd, child.block.pdl__timing.end_nanos),
+      0,
+    )
+    root.block.pdl__timing.end_nanos = Math.max(
+      root.block.pdl__timing.end_nanos,
+      maxEnd,
+    )
+  }
+
+  return [...(ignoreRoot ? [] : [root]), ...childrenModel].filter(nonNullable)
 }
 
 export function childrenOf(block: NonScalarPdlBlock) {

src/pdl/pdl_dumper.py

@@ -97,8 +97,6 @@ def block_to_dict(  # noqa: C901
             context = block.context
         if len(context) > 0:
             d["context"] = context
-    if block.pdl__timing is not None:
-        d["pdl__timing"] = timing_to_dict(block.pdl__timing)
     if block.description is not None:
         d["description"] = block.description
     if block.role is not None:
@@ -239,6 +237,12 @@ def block_to_dict(  # noqa: C901
     #     d["location"] = location_to_dict(block.location)
     if block.fallback is not None:
         d["fallback"] = block_to_dict(block.fallback, json_compatible)
+    # Warning: remember to update timing here at the end! this ensures
+    # that any logic that updates timestamps when futures
+    # finish... has a chance to do its work before we record the
+    # timestamps to the trace
+    if block.pdl__timing is not None:
+        d["pdl__timing"] = timing_to_dict(block.pdl__timing)
     return d
 
 

src/pdl/pdl_llms.py

@@ -100,6 +100,16 @@ def generate_text(
         pdl_future: PdlLazy[tuple[dict[str, Any], Any]] = PdlConst(future)
         message = lazy_apply((lambda x: x[0]), pdl_future)
         response = lazy_apply((lambda x: x[1]), pdl_future)
+
+        # update the end timestamp when the future is done
+        def update_end_nanos(future):
+            import time
+
+            if block.pdl__timing is not None:
+                block.pdl__timing.end_nanos = time.time_ns()
+
+        future.add_done_callback(update_end_nanos)
+
         return message, response
 
     @staticmethod