Merge pull request #124 from braintrustdata/ark/llm-convo-fix

squash BrainstoreTrace LLM convo bug

Author: realark

braintrust-sdk/src/main/java/dev/braintrust/trace/BrainstoreTrace.java

@@ -146,9 +146,14 @@ public List<Map<String, Object>> getSpans(@Nonnull String spanType) {
     public List<Map<String, Object>> getLLMConversationThread() {
         var allSpans = getSpans();
 
-        // Build children map: parent_id → children sorted by start_time
+        // Build children map: parent_id → children sorted by start_time, and collect the set of
+        // span IDs present in the fetched results.
         var children = new java.util.LinkedHashMap<String, List<Map<String, Object>>>();
+        var presentSpanIds = new java.util.HashSet<String>();
         for (var span : allSpans) {
+            if (span.get("span_id") instanceof String sid) {
+                presentSpanIds.add(sid);
+            }
             var parents = span.get("span_parents");
             if (parents instanceof List<?> parentList && !parentList.isEmpty()) {
                 if (parentList.get(0) instanceof String pid) {
@@ -163,24 +168,34 @@ public List<Map<String, Object>> getLLMConversationThread() {
                                         (a, b) ->
                                                 Double.compare(getStartTime(a), getStartTime(b))));
 
-        // Find root span (no parents)
-        var root =
+        // Find forest roots: spans with no parents, or whose parent is absent from the fetched
+        // set (their true ancestor hasn't been ingested yet). Sort by start time so the DFS
+        // visits orphan subtrees in chronological order.
+        var forestRoots =
                 allSpans.stream()
                         .filter(
                                 s -> {
                                     var p = s.get("span_parents");
-                                    return p == null || (p instanceof List<?> l && l.isEmpty());
+                                    if (p == null || (p instanceof List<?> l && l.isEmpty())) {
+                                        return true;
+                                    }
+                                    return p instanceof List<?> l
+                                            && l.get(0) instanceof String pid
+                                            && !presentSpanIds.contains(pid);
                                 })
-                        .findFirst()
-                        .orElse(null);
-        if (root == null) return List.of();
+                        .sorted((a, b) -> Double.compare(getStartTime(a), getStartTime(b)))
+                        .toList();
+        if (forestRoots.isEmpty()) return List.of();
 
         // Pre-order DFS to get all LLM spans in hierarchy order.
         // Prune entire subtrees rooted at scorer spans (purpose == "scorer") — these are
         // synthetic spans injected by the Braintrust backend and not part of the real trace.
         var llmSpansInOrder = new ArrayList<Map<String, Object>>();
         var stack = new java.util.ArrayDeque<Map<String, Object>>();
-        stack.push(root);
+        // Push forest roots in reverse start order so the earliest is processed first.
+        for (int i = forestRoots.size() - 1; i >= 0; i--) {
+            stack.push(forestRoots.get(i));
+        }
         while (!stack.isEmpty()) {
             var span = stack.pop();
             if ("automation".equals(getSpanType(span))) {

braintrust-sdk/src/test/java/dev/braintrust/trace/BrainstoreTraceTest.java

@@ -370,6 +370,60 @@ void getLLMConversationThreadPrunesAutomationSubtrees() {
         assertEquals(choice, thread.get(1));
     }
 
+    @Test
+    void getLLMConversationThreadWorksWithoutRootSpan() {
+        // Trace scorers run mid-flight: spans close bottom-up, so the parent/root span may not
+        // yet be ingested. Here the task span's parent ("root") is absent from the fetched set,
+        // making the task span a forest root. The thread must still be reconstructed.
+        var sysMsg = Map.<String, Object>of("role", "system", "content", "be helpful");
+        var userMsg = Map.<String, Object>of("role", "user", "content", "strawberry");
+        var assistantMsg = Map.<String, Object>of("role", "assistant", "content", "fruit");
+        var choice =
+                Map.<String, Object>of(
+                        "finish_reason", "stop", "index", 0, "message", assistantMsg);
+
+        // task span points at a "root" parent that was NOT fetched (root not ended/ingested yet)
+        var task = taskSpan("task", "root", 1.0);
+        var llm = llmSpan("llm1", "task", 1.1, List.of(sysMsg, userMsg), List.of(choice));
+
+        var trace = traceWithSpans(List.of(task, llm)); // no root span present
+        var thread = trace.getLLMConversationThread();
+
+        assertEquals(3, thread.size(), "thread must be reconstructed even without the root span");
+        assertEquals("system", thread.get(0).get("role"));
+        assertEquals("user", thread.get(1).get("role"));
+        assertEquals(assistantMsg, thread.get(2).get("message"));
+    }
+
+    @Test
+    void getLLMConversationThreadHandlesMultipleOrphanSubtrees() {
+        // Multiple orphan subtrees whose common ancestor ("root") is absent. Each subtree is a
+        // forest root; they must be visited in start-time order and their messages concatenated.
+        var user1 = Map.<String, Object>of("role", "user", "content", "Q1");
+        var asst1 = Map.<String, Object>of("role", "assistant", "content", "A1");
+        var choice1 = Map.<String, Object>of("finish_reason", "stop", "message", asst1);
+
+        var user2 = Map.<String, Object>of("role", "user", "content", "Q2");
+        var asst2 = Map.<String, Object>of("role", "assistant", "content", "A2");
+        var choice2 = Map.<String, Object>of("finish_reason", "stop", "message", asst2);
+
+        // Two task subtrees both parented at a missing "root". Provide out of order; turn2 starts
+        // later than turn1, so turn1's messages must come first.
+        var turn2 = taskSpan("turn2", "root", 2.0);
+        var llm2 = llmSpan("llm2", "turn2", 2.1, List.of(user2), List.of(choice2));
+        var turn1 = taskSpan("turn1", "root", 1.0);
+        var llm1 = llmSpan("llm1", "turn1", 1.1, List.of(user1), List.of(choice1));
+
+        var trace = traceWithSpans(List.of(turn2, llm2, turn1, llm1)); // root absent, out of order
+        var thread = trace.getLLMConversationThread();
+
+        assertEquals(4, thread.size());
+        assertEquals("Q1", thread.get(0).get("content"));
+        assertEquals(asst1, thread.get(1).get("message"));
+        assertEquals("Q2", thread.get(2).get("content"));
+        assertEquals(asst2, thread.get(3).get("message"));
+    }
+
     // -------------------------------------------------------------------------
     // fetchWithRetry — retry logic (via package-private constructor with custom supplier)
     // -------------------------------------------------------------------------