feat(llmobs): carry typed messages with media on non-LLM span kinds

ddtrace/llmobs/_constants.py

@@ -234,6 +234,8 @@ class LLMOBS_STRUCT:
     EXPECTED_OUTPUT: Final = "expected_output"
     VALUE: Final = "value"
     MESSAGES: Final = "messages"
+    AUDIO_PARTS: Final = "audio_parts"
+    IMAGE_PARTS: Final = "image_parts"
     DOCUMENTS: Final = "documents"
     AGENT_MANIFEST: Final = "agent_manifest"
     SPAN: Final = "span"

ddtrace/llmobs/_llmobs.py

@@ -159,6 +159,7 @@
 from ddtrace.llmobs._utils import _trace_id_to_wire
 from ddtrace.llmobs._utils import _validate_prompt
 from ddtrace.llmobs._utils import add_span_link
+from ddtrace.llmobs._utils import collapse_messages_to_value
 from ddtrace.llmobs._utils import enforce_message_role
 from ddtrace.llmobs._utils import get_asyncio
 from ddtrace.llmobs._utils import get_llmobs_ml_app
@@ -174,6 +175,7 @@
 from ddtrace.llmobs._utils import resolve_llmobs_git_metadata
 from ddtrace.llmobs._utils import resolve_ml_app
 from ddtrace.llmobs._utils import safe_json
+from ddtrace.llmobs._utils import span_kind_keeps_messages
 from ddtrace.llmobs._writer import LLMObsAPIClient
 from ddtrace.llmobs._writer import LLMObsEvalMetricWriter
 from ddtrace.llmobs._writer import LLMObsExperimentsClient
@@ -424,7 +426,7 @@ def _build_llmobs_span(
         llmobs_span.input = [Message(content=safe_json(input_value, ensure_ascii=False) or "", role="")]
 
     input_messages = llmobs_input.get(LLMOBS_STRUCT.MESSAGES)
-    if span_kind == "llm" and input_messages is not None:
+    if input_messages is not None and span_kind_keeps_messages(span_kind, input_messages):
         input_type = "messages"
         llmobs_span.input = enforce_message_role(input_messages)
 
@@ -439,7 +441,7 @@ def _build_llmobs_span(
         llmobs_span.output = [Message(content=safe_json(output_value, ensure_ascii=False) or "", role="")]
 
     output_messages = llmobs_output.get(LLMOBS_STRUCT.MESSAGES)
-    if span_kind == "llm" and output_messages is not None:
+    if output_messages is not None and span_kind_keeps_messages(span_kind, output_messages):
         output_type = "messages"
         llmobs_span.output = enforce_message_role(output_messages)
 
@@ -541,6 +543,11 @@ def _normalize_llmobs_meta(
 
     if input_type == "messages":
         meta_input[LLMOBS_STRUCT.MESSAGES] = llmobs_span.input
+        if span_kind != "llm":
+            # Non-LLM spans carry both: value for readers that only understand value, and
+            # messages for the typed media parts value cannot represent. Derived here, after
+            # the user span processor, so processor edits reach both fields.
+            meta_input[LLMOBS_STRUCT.VALUE] = collapse_messages_to_value(span_kind, llmobs_span.input)
     elif input_type == "value" and llmobs_span.input:
         meta_input[LLMOBS_STRUCT.VALUE] = llmobs_span.input[0].get("content", "")
     elif input_type == "documents":
@@ -552,6 +559,8 @@ def _normalize_llmobs_meta(
 
     if output_type == "messages":
         meta_output[LLMOBS_STRUCT.MESSAGES] = llmobs_span.output
+        if span_kind != "llm":
+            meta_output[LLMOBS_STRUCT.VALUE] = collapse_messages_to_value(span_kind, llmobs_span.output)
     elif output_type == "value" and llmobs_span.output:
         meta_output[LLMOBS_STRUCT.VALUE] = llmobs_span.output[0].get("content", "")
     elif output_type == "documents":
@@ -3077,6 +3086,22 @@ def annotate(
                     )
                 elif span_kind == "experiment":
                     cls._tag_freeform_io(span, input_value=input_data, output_value=output_data)
+                elif span_kind_keeps_messages(span_kind, input_data) or span_kind_keeps_messages(
+                    span_kind, output_data
+                ):
+                    # Media-bearing sides route through the message tagger so the typed parts
+                    # survive. Each side is decided on its own: a plain-string side keeps the
+                    # value tagging it has today rather than being reshaped into a message.
+                    media_input = input_data if span_kind_keeps_messages(span_kind, input_data) else None
+                    media_output = output_data if span_kind_keeps_messages(span_kind, output_data) else None
+                    annotation_error_message, error = cls._tag_llm_io(
+                        span, input_messages=media_input, output_messages=media_output
+                    )
+                    cls._tag_text_io(
+                        span,
+                        input_value=input_data if media_input is None else None,
+                        output_value=output_data if media_output is None else None,
+                    )
                 else:
                     cls._tag_text_io(span, input_value=input_data, output_value=output_data)
             if _linked_spans and isinstance(_linked_spans, list):

ddtrace/llmobs/_utils.py

@@ -795,10 +795,11 @@ def _annotate_llmobs_span_data(
             llmobs_span_data[LLMOBS_STRUCT.SPAN_LINKS] = span_links
         if config is not None:
             llmobs_span_data[LLMOBS_STRUCT.CONFIG] = config
-        # Add I/O messages to messages field only for LLM spans, otherwise add to value field
-        is_llm = meta[LLMOBS_STRUCT.SPAN].get(LLMOBS_STRUCT.KIND) == "llm"
+        # Add I/O messages to the messages field for LLM spans, and for the non-LLM kinds that
+        # keep messages when a message carries media; every other case collapses to value.
+        annotated_span_kind = meta[LLMOBS_STRUCT.SPAN].get(LLMOBS_STRUCT.KIND)
         if input_messages is not None:
-            if is_llm:
+            if span_kind_keeps_messages(annotated_span_kind, input_messages):
                 meta[LLMOBS_STRUCT.INPUT][LLMOBS_STRUCT.MESSAGES] = input_messages
             else:
                 meta[LLMOBS_STRUCT.INPUT][LLMOBS_STRUCT.VALUE] = safe_json(input_messages, ensure_ascii=False) or ""
@@ -811,7 +812,7 @@ def _annotate_llmobs_span_data(
             existing_prompt.update(cast(Prompt, prompt))
             span._set_ctx_item(INPUT_PROMPT, existing_prompt)
         if output_messages is not None:
-            if is_llm:
+            if span_kind_keeps_messages(annotated_span_kind, output_messages):
                 meta[LLMOBS_STRUCT.OUTPUT][LLMOBS_STRUCT.MESSAGES] = output_messages
             else:
                 meta[LLMOBS_STRUCT.OUTPUT][LLMOBS_STRUCT.VALUE] = safe_json(output_messages, ensure_ascii=False) or ""
@@ -851,6 +852,87 @@ def enforce_message_role(messages: list[Message]) -> list[Message]:
     return messages
 
 
+# Non-LLM span kinds allowed to keep typed messages alongside the collapsed value string.
+# Adding a kind here requires the serving API to populate messages for it first, otherwise the
+# messages are dropped on read and the inline media has spent the per-event size budget for
+# something that can never render. The serving API now builds messages for every kind below:
+# agent has its own builder, and workflow / task / step / tool route through defaultSpanFromEvent,
+# which populates them as of the non-LLM span builder change.
+MEDIA_MESSAGE_SPAN_KINDS: frozenset = frozenset(("agent", "workflow", "task", "step", "tool"))
+
+_SCALAR_VALUE_SPAN_KINDS: frozenset = frozenset(("agent", "workflow", "task", "step"))
+_MEDIA_PART_KEYS = (LLMOBS_STRUCT.AUDIO_PARTS, LLMOBS_STRUCT.IMAGE_PARTS)
+
+
+def messages_carry_media(messages: Any) -> bool:
+    """True when any message carries a non-empty audio_parts or image_parts list.
+
+    Runs on the annotate path against unvalidated user input, so it tolerates any shape
+    rather than raising. A media key holding something other than a non-empty list does
+    not count, which leaves malformed input on whatever path it takes today.
+    """
+    if isinstance(messages, dict):
+        messages = [messages]
+    if not isinstance(messages, list):
+        return False
+    for message in messages:
+        if not isinstance(message, dict):
+            continue
+        for key in _MEDIA_PART_KEYS:
+            parts = message.get(key)
+            if isinstance(parts, list) and parts:
+                return True
+    return False
+
+
+def span_kind_keeps_messages(span_kind: Optional[str], messages: Any) -> bool:
+    """True when message-shaped I/O should be stored as typed messages for this span kind.
+
+    LLM spans always keep messages. The other kinds keep them only when a message actually
+    carries media, mirroring the trace indexer, which retains messages on a non-LLM span
+    only for the typed parts the collapsed value string cannot represent.
+    """
+    if span_kind == "llm":
+        return True
+    return span_kind in MEDIA_MESSAGE_SPAN_KINDS and messages_carry_media(messages)
+
+
+def _strip_media_parts(messages: list[Message]) -> list[dict]:
+    """Drop the media part lists so base64 payloads never reach the value string."""
+    return [{k: v for k, v in message.items() if k not in _MEDIA_PART_KEYS} for message in messages]
+
+
+def _messages_have_tool_structure(messages: list[dict]) -> bool:
+    """True when any message carries tool_calls or tool_results.
+
+    Plain-text rendering cannot represent tool structure without dropping it, so its
+    presence forces the JSON form of the value string.
+    """
+    return any("tool_calls" in message or "tool_results" in message for message in messages)
+
+
+def collapse_messages_to_value(span_kind: str, messages: list[Message]) -> str:
+    """Render the scalar value string for a non-LLM span that also carries typed messages.
+
+    Mirrors nonLLMValueString in the trace indexer so a natively instrumented span and an
+    OTel one read back the same way. Agent, workflow, task and step render a lone
+    plain-text, non-system message as readable text; anything else keeps the JSON form,
+    which preserves roles, turns and tool structure. Media parts are stripped either way,
+    so the value stays small and the payload lives only on messages.
+    """
+    if not messages:
+        return ""
+    stripped = _strip_media_parts(messages)
+    if (
+        span_kind in _SCALAR_VALUE_SPAN_KINDS
+        and len(stripped) == 1
+        and not _messages_have_tool_structure(stripped)
+        and stripped[0].get("role") != "system"
+    ):
+        return stripped[0].get("content", "") or ""
+    return safe_json(stripped, ensure_ascii=False) or ""
+
+
 def validate_tags_list(tags: list[str]) -> None:
     if not isinstance(tags, list):
         raise TypeError("Tags must be a list of strings")

releasenotes/notes/llmobs-agent-span-media-parts-fc6e57c41f7c178f.yaml

@@ -0,0 +1,8 @@
+---
+fixes:
+  - |
+    LLM Observability: Fixes an issue where images and audio annotated on non-LLM spans were lost.
+    Messages annotated via ``LLMObs.annotate`` on an ``agent``, ``workflow``, ``task``, ``tool`` or
+    ``step`` span that carry ``image_parts`` or ``audio_parts`` are now recorded as structured
+    messages alongside the span's input and output values, so the media is collected and rendered
+    instead of being flattened into text. Spans without media, and non-message input, are unchanged.