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":
@@ -2950,6 +2959,10 @@ def annotate(
                                         "mime_type" and one of "content" (base64-encoded image) or "attachment_key".
                            - embedding spans: accepts a string, list of strings, or a dictionary of form
                                               {"text": "...", ...} or a list of dictionaries with the same signature.
+                           - agent, workflow, task, step and tool spans: any JSON serializable type. A message-shaped
+                                              payload whose messages carry "image_parts" or "audio_parts" is recorded
+                                              as typed messages alongside the collapsed value string, and keys outside
+                                              the message schema described above are dropped.
                            - other: any JSON serializable type.
         :param output_data: A single output string, dictionary, or a list of dictionaries based on the span kind:
                            - llm spans: accepts a string, or a dictionary of form {"content": "...", "role": "...",
@@ -2963,6 +2976,10 @@ def annotate(
                            - retrieval spans: a dictionary containing any of the key value pairs
                                               {"name": str, "id": str, "text": str, "score": float},
                                               or a list of dictionaries with the same signature.
+                           - agent, workflow, task, step and tool spans: any JSON serializable type. A message-shaped
+                                              payload whose messages carry "image_parts" or "audio_parts" is recorded
+                                              as typed messages alongside the collapsed value string, and keys outside
+                                              the message schema described above are dropped.
                            - other: any JSON serializable type.
         :param metadata: Dictionary of JSON serializable key-value metadata pairs relevant to the input/output operation
                          described by the LLMObs span.
@@ -3077,6 +3094,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,14 +795,21 @@ 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.input persists across annotate calls, and the emit path prefers messages
+                # over value. Clearing the sibling keeps a later annotate authoritative rather
+                # than leaving a stale representation to win.
+                meta[LLMOBS_STRUCT.INPUT].pop(LLMOBS_STRUCT.VALUE, None)
                 meta[LLMOBS_STRUCT.INPUT][LLMOBS_STRUCT.MESSAGES] = input_messages
             else:
+                meta[LLMOBS_STRUCT.INPUT].pop(LLMOBS_STRUCT.MESSAGES, None)
                 meta[LLMOBS_STRUCT.INPUT][LLMOBS_STRUCT.VALUE] = safe_json(input_messages, ensure_ascii=False) or ""
         if input_value is not None:
+            meta[LLMOBS_STRUCT.INPUT].pop(LLMOBS_STRUCT.MESSAGES, None)
             meta[LLMOBS_STRUCT.INPUT][LLMOBS_STRUCT.VALUE] = safe_json(input_value, ensure_ascii=False) or ""
         if input_documents is not None:
             meta[LLMOBS_STRUCT.INPUT][LLMOBS_STRUCT.DOCUMENTS] = input_documents
@@ -811,11 +818,14 @@ 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].pop(LLMOBS_STRUCT.VALUE, None)
                 meta[LLMOBS_STRUCT.OUTPUT][LLMOBS_STRUCT.MESSAGES] = output_messages
             else:
+                meta[LLMOBS_STRUCT.OUTPUT].pop(LLMOBS_STRUCT.MESSAGES, None)
                 meta[LLMOBS_STRUCT.OUTPUT][LLMOBS_STRUCT.VALUE] = safe_json(output_messages, ensure_ascii=False) or ""
         if output_value is not None:
+            meta[LLMOBS_STRUCT.OUTPUT].pop(LLMOBS_STRUCT.MESSAGES, None)
             meta[LLMOBS_STRUCT.OUTPUT][LLMOBS_STRUCT.VALUE] = safe_json(output_value, ensure_ascii=False) or ""
         if output_documents is not None:
             meta[LLMOBS_STRUCT.OUTPUT][LLMOBS_STRUCT.DOCUMENTS] = output_documents
@@ -851,6 +861,116 @@ 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 anything other than a non-empty list of dicts
+    does not count, which leaves malformed input on whatever path it takes today: routing it
+    here instead would hand it to Messages(), which raises TypeError on a non-dict part, and
+    annotate turns that into an LLMObsAnnotateSpanError in the caller's own code.
+    """
+    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 and all(isinstance(part, dict) for part in 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]
+
+
+_VALUE_OMITTED_MARKER = "[value omitted: event size limit]"
+
+
+def _media_payload_chars(messages: list[Message]) -> int:
+    """Characters the inline media parts contribute, without serializing the whole list."""
+    total = 0
+    for message in messages:
+        if not isinstance(message, dict):
+            continue
+        for key in _MEDIA_PART_KEYS:
+            parts = message.get(key)
+            if not isinstance(parts, list):
+                continue
+            for part in parts:
+                if isinstance(part, dict):
+                    total += len(part.get("content", "") or "")
+    return total
+
+
+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"
+    ):
+        value = stripped[0].get("content", "") or ""
+    else:
+        value = safe_json(stripped, ensure_ascii=False) or ""
+    # The messages already carry this text, so the derived value duplicates it and halves the
+    # headroom the media needs. Past the limit the writer replaces input and output wholesale
+    # (_writer._truncate_span_event), discarding the very media this field exists to accompany,
+    # so drop the duplicate instead and keep the payload.
+    if _media_payload_chars(messages) + 2 * len(value) > config._llmobs_event_size_limit:
+        return _VALUE_OMITTED_MARKER
+    return value
+
+
 def validate_tags_list(tags: list[str]) -> None:
     if not isinstance(tags, list):
         raise TypeError("Tags must be a list of strings")

ddtrace/llmobs/decorators.py

@@ -36,6 +36,25 @@ def _get_span_inputs(args: OrderedDict) -> dict:
     return {arg: value for arg, value in args.items() if arg != "self"}
 
 
+def _annotate_or_log(span, operation_kind: str, io_kind: str, **annotation) -> None:
+    """Annotate a decorated span without letting an annotation failure reach the caller.
+
+    Automatic annotation is on by default, so the decorator inspects whatever the user's
+    arguments and return value happen to be. A shape annotate rejects must not take down
+    the function being traced, especially on the input side, where this runs before the
+    function body.
+    """
+    try:
+        LLMObs.annotate(span=span, **annotation)
+    except LLMObsAnnotateSpanError:
+        log.debug(
+            "Failed to auto-annotate %s for @%s decorated function. Use LLMObs.annotate() to manually annotate the %s.",
+            io_kind,
+            operation_kind,
+            io_kind,
+        )
+
+
 async def yield_from_async_gen(func, span, args, kwargs):
     try:
         gen = func(*args, **kwargs)
@@ -235,7 +254,9 @@ def generator_wrapper(*args, **kwargs):
                     func_signature = signature(func)
                     bound_args = func_signature.bind_partial(*args, **kwargs)
                     if _automatic_io_annotation and bound_args.arguments:
-                        LLMObs.annotate(span=span, input_data=_get_span_inputs(bound_args.arguments))
+                        _annotate_or_log(
+                            span, operation_kind, "input", input_data=_get_span_inputs(bound_args.arguments)
+                        )
                     return yield_from_async_gen(func, span, args, kwargs)
 
                 @wraps(func)
@@ -256,11 +277,13 @@ async def wrapper(*args, **kwargs):
                         func_signature = signature(func)
                         bound_args = func_signature.bind_partial(*args, **kwargs)
                         if _automatic_io_annotation and bound_args.arguments:
-                            LLMObs.annotate(span=span, input_data=_get_span_inputs(bound_args.arguments))
+                            _annotate_or_log(
+                                span, operation_kind, "input", input_data=_get_span_inputs(bound_args.arguments)
+                            )
                         resp = await func(*args, **kwargs)
                         if _automatic_io_annotation and resp is not None and operation_kind != "retrieval":
                             if get_llmobs_output_value(span) is None:
-                                LLMObs.annotate(span=span, output_data=resp)
+                                _annotate_or_log(span, operation_kind, "output", output_data=resp)
                         return resp
 
             else:
@@ -284,7 +307,9 @@ def generator_wrapper(*args, **kwargs):
                         func_signature = signature(func)
                         bound_args = func_signature.bind_partial(*args, **kwargs)
                         if _automatic_io_annotation and bound_args.arguments:
-                            LLMObs.annotate(span=span, input_data=_get_span_inputs(bound_args.arguments))
+                            _annotate_or_log(
+                                span, operation_kind, "input", input_data=_get_span_inputs(bound_args.arguments)
+                            )
                         try:
                             yield from func(*args, **kwargs)
                         except (StopIteration, GeneratorExit):
@@ -314,11 +339,13 @@ def wrapper(*args, **kwargs):
                         func_signature = signature(func)
                         bound_args = func_signature.bind_partial(*args, **kwargs)
                         if _automatic_io_annotation and bound_args.arguments:
-                            LLMObs.annotate(span=span, input_data=_get_span_inputs(bound_args.arguments))
+                            _annotate_or_log(
+                                span, operation_kind, "input", input_data=_get_span_inputs(bound_args.arguments)
+                            )
                         resp = func(*args, **kwargs)
                         if _automatic_io_annotation and resp is not None and operation_kind != "retrieval":
                             if get_llmobs_output_value(span) is None:
-                                LLMObs.annotate(span=span, output_data=resp)
+                                _annotate_or_log(span, operation_kind, "output", output_data=resp)
                         return resp
 
             return generator_wrapper if (isgeneratorfunction(func) or isasyncgenfunction(func)) else wrapper