promptlayer 1.3.1 → 1.3.3

package/vendor/claude-agents/trace/hooks/py/stop_parser.py

@@ -68,6 +68,22 @@ def message_text(content):
     return ""
 
 
+def message_thinking(content):
+    if isinstance(content, list):
+        parts = []
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "thinking":
+                thinking = block.get("thinking")
+                if isinstance(thinking, str) and thinking:
+                    parts.append(thinking)
+        return "\n".join(parts).strip()
+    if isinstance(content, dict) and content.get("type") == "thinking":
+        thinking = content.get("thinking")
+        if isinstance(thinking, str):
+            return thinking
+    return ""
+
+
 def flatten_indexed(prefix, items, out):
     for i, item in enumerate(items):
         for key, value in item.items():
@@ -101,6 +117,69 @@ def is_tool_result_user(rec):
     )
 
 
+def assistant_message_id(rec):
+    if rec.get("type") != "assistant":
+        return ""
+    msg = rec.get("message", {})
+    if not isinstance(msg, dict):
+        return ""
+    return stringify(msg.get("id"))
+
+
+def merge_content_blocks(existing, incoming):
+    if isinstance(existing, list) and isinstance(incoming, list):
+        return existing + incoming
+    if existing is None:
+        return incoming
+    if incoming is None:
+        return existing
+    if isinstance(existing, list):
+        return existing + [incoming]
+    if isinstance(incoming, list):
+        return [existing] + incoming
+    return [existing, incoming]
+
+
+def coalesce_assistant_message_fragments(records):
+    coalesced = []
+    for rec in records:
+        if not coalesced:
+            coalesced.append(rec)
+            continue
+
+        prev = coalesced[-1]
+        rec_msg_id = assistant_message_id(rec)
+        prev_msg_id = assistant_message_id(prev)
+        same_assistant_message = (
+            rec_msg_id
+            and prev_msg_id == rec_msg_id
+            and rec.get("sessionId") == prev.get("sessionId")
+        )
+        if not same_assistant_message:
+            coalesced.append(rec)
+            continue
+
+        prev_msg = prev.get("message", {})
+        rec_msg = rec.get("message", {})
+        if not isinstance(prev_msg, dict) or not isinstance(rec_msg, dict):
+            coalesced.append(rec)
+            continue
+
+        prev_msg["content"] = merge_content_blocks(
+            prev_msg.get("content"),
+            rec_msg.get("content"),
+        )
+        for key in ("model", "stop_reason", "stop_sequence", "stop_details", "context_management"):
+            if rec_msg.get(key) is not None:
+                prev_msg[key] = rec_msg.get(key)
+        if rec_msg.get("usage"):
+            prev_msg["usage"] = rec_msg.get("usage")
+        if rec.get("timestamp"):
+            prev["timestamp"] = rec.get("timestamp")
+
+    return coalesced
+
+
 def parse_transcript(transcript_path, turn_start_fallback, pending_payloads, expected_session_id=None):
     records = []
     with open(transcript_path, encoding="utf-8") as f:
@@ -118,6 +197,8 @@ def parse_transcript(transcript_path, turn_start_fallback, pending_payloads, exp
             except Exception:
                 continue
 
+    records = coalesce_assistant_message_fragments(records)
+
     if not records:
         now_ns = int(datetime.now(timezone.utc).timestamp() * 1_000_000_000)
         return {"turn": {"start_ns": now_ns, "end_ns": now_ns}, "tools": [], "llms": []}
@@ -137,7 +218,6 @@ def parse_transcript(transcript_path, turn_start_fallback, pending_payloads, exp
     llms = []
     pending_tool_uses = []
     pending_payload_idx = 0
-    saw_human_input = False
 
     turn_start_ns = turn_start_fallback
     turn_end_ns = turn_start_fallback
@@ -160,7 +240,6 @@ def parse_transcript(transcript_path, turn_start_fallback, pending_payloads, exp
                 if content:
                     append_history_item(history, {"role": "user", "content": content})
                     last_input_ns = timestamp_ns or last_input_ns
-                    saw_human_input = True
             continue
 
         if rec_type == "user":
@@ -231,7 +310,6 @@ def parse_transcript(transcript_path, turn_start_fallback, pending_payloads, exp
             user_text = content_to_text(content)
             append_history_item(history, {"role": "user", "content": user_text})
             last_input_ns = timestamp_ns or last_input_ns
-            saw_human_input = True
             continue
 
         if rec_type != "assistant":
@@ -248,6 +326,7 @@ def parse_transcript(transcript_path, turn_start_fallback, pending_payloads, exp
         prompt_tokens = safe_int(usage.get("input_tokens"), 0)
         completion_tokens = safe_int(usage.get("output_tokens"), 0)
         output_text = message_text(msg.get("content"))
+        output_thinking = message_thinking(msg.get("content"))
 
         tool_calls = []
         content_blocks = msg.get("content")
@@ -277,7 +356,7 @@ def parse_transcript(transcript_path, turn_start_fallback, pending_payloads, exp
                     }
                 )
 
-        if not output_text and not tool_calls:
+        if not output_text and not output_thinking and not tool_calls:
             continue
 
         llm_start_ns = last_input_ns or timestamp_ns or turn_start_ns
@@ -309,11 +388,13 @@ def parse_transcript(transcript_path, turn_start_fallback, pending_payloads, exp
         flatten_indexed("gen_ai.prompt", history, attrs)
 
         completion_item = {"role": "assistant", "content": output_text}
+        if output_thinking:
+            completion_item["thinking"] = output_thinking
         if tool_calls:
             completion_item["tool_calls"] = tool_calls
         flatten_indexed("gen_ai.completion", [completion_item], attrs)
 
-        span_name = "LLM Call (User)" if saw_human_input else "LLM call"
+        span_name = "LLM call"
 
         if emit_for_turn:
             llms.append(
@@ -326,10 +407,11 @@ def parse_transcript(transcript_path, turn_start_fallback, pending_payloads, exp
             )
 
         assistant_history = {"role": "assistant", "content": output_text}
+        if output_thinking:
+            assistant_history["thinking"] = output_thinking
         if tool_calls:
             assistant_history["tool_calls"] = tool_calls
         history.append(assistant_history)
-        saw_human_input = False
 
     if turn_start_ns is None:
         turn_start_ns = int(datetime.now(timezone.utc).timestamp() * 1_000_000_000)
@@ -369,14 +451,14 @@ def build_stop_hook_span_specs(
             trace_id=trace_id,
             span_id=session_span_id,
             parent_span_id=session_parent_span_id,
-            name="Claude Code session",
+            name="LLM session",
             kind="1",
             start_ns=str(session_start_ns),
             end_ns=turn_end_ns,
             attrs={
                 "source": "claude-code",
                 "hook": session_hook_attr,
-                "node_type": "WORKFLOW",
+                "node_type": "LLM_SESSION",
                 "session.lifecycle": session_lifecycle_attr,
             },
         )