fix(mistralai): record chunked assistant content instead of dropping it

packages/opentelemetry-instrumentation-mistralai/opentelemetry/instrumentation/mistralai/__init__.py

@@ -133,6 +133,21 @@ def _set_model_input_attributes(span, to_wrap, kwargs):
     )
 
 
+def _content_as_str(content):
+    # Reasoning models answer with a list of content chunks, which are SDK models
+    # rather than plain dicts, so they are dumped before being serialised.
+    if content is None or isinstance(content, str):
+        return content
+    return json.dumps(
+        [
+            chunk.model_dump(mode="json", exclude_none=True)
+            if hasattr(chunk, "model_dump")
+            else chunk
+            for chunk in content
+        ]
+    )
+
+
 @dont_throw
 def _set_response_attributes(span, llm_request_type, response):
     if llm_request_type == LLMRequestTypeValues.EMBEDDING or not span.is_recording():
@@ -149,11 +164,7 @@ def _set_response_attributes(span, llm_request_type, response):
             _set_span_attribute(
                 span,
                 f"{prefix}.content",
-                (
-                    choice.message.content
-                    if isinstance(choice.message.content, str)
-                    else json.dumps(choice.message.content)
-                ),
+                _content_as_str(choice.message.content),
             )
             _set_span_attribute(
                 span,
@@ -217,6 +228,54 @@ def _set_model_response_attributes(span, llm_request_type, response):
             )
 
 
+def _append_chunk(chunks, chunk):
+    # A delta that continues the trailing chunk of the same type extends it, so a
+    # streamed answer ends up with the same chunks as the non-streamed one.
+    last = chunks[-1] if chunks else None
+    if last is not None and last.get("type") == chunk.get("type"):
+        if chunk.get("type") == "text":
+            last["text"] = last.get("text", "") + chunk.get("text", "")
+            return
+        if chunk.get("type") == "thinking":
+            thoughts = last.setdefault("thinking", [])
+            for thought in chunk.get("thinking") or []:
+                if (
+                    thoughts
+                    and thoughts[-1].get("type") == "text"
+                    and thought.get("type") == "text"
+                ):
+                    thoughts[-1]["text"] = thoughts[-1].get("text", "") + thought.get(
+                        "text", ""
+                    )
+                else:
+                    thoughts.append(dict(thought))
+            return
+    chunks.append(dict(chunk))
+
+
+def _merge_delta_content(current, delta):
+    # Text deltas are joined as before. Reasoning models stream lists of chunks,
+    # which are collected as plain data so the whole answer can be serialised.
+    if not delta:
+        return current
+    if isinstance(delta, str):
+        if isinstance(current, list):
+            _append_chunk(current, {"type": "text", "text": delta})
+            return current
+        return (current or "") + delta
+    chunks = current if isinstance(current, list) else []
+    if isinstance(current, str) and current:
+        chunks = [{"type": "text", "text": current}]
+    for chunk in delta:
+        _append_chunk(
+            chunks,
+            chunk.model_dump(mode="json", exclude_none=True)
+            if hasattr(chunk, "model_dump")
+            else chunk,
+        )
+    return chunks
+
+
 def _accumulate_streaming_response(span, event_logger, llm_request_type, response):
     accumulated_response = ChatCompletionResponse(
         id="",
@@ -251,7 +310,9 @@ def _accumulate_streaming_response(span, event_logger, llm_request_type, respons
                 )
 
             accumulated_response.choices[idx].finish_reason = choice.finish_reason
-            accumulated_response.choices[idx].message.content += choice.delta.content
+            accumulated_response.choices[idx].message.content = _merge_delta_content(
+                accumulated_response.choices[idx].message.content, choice.delta.content
+            )
             accumulated_response.choices[idx].message.role = choice.delta.role
 
     _handle_response(span, event_logger, llm_request_type, accumulated_response)
@@ -295,7 +356,9 @@ async def _aaccumulate_streaming_response(
                 )
 
             accumulated_response.choices[idx].finish_reason = choice.finish_reason
-            accumulated_response.choices[idx].message.content += choice.delta.content
+            accumulated_response.choices[idx].message.content = _merge_delta_content(
+                accumulated_response.choices[idx].message.content, choice.delta.content
+            )
             accumulated_response.choices[idx].message.role = choice.delta.role
 
     _handle_response(span, event_logger, llm_request_type, accumulated_response)
@@ -360,7 +423,7 @@ def _emit_choice_events(
                 ChoiceEvent(
                     index=choice.index,
                     message={
-                        "content": choice.message.content,
+                        "content": _content_as_str(choice.message.content),
                         "role": choice.message.role or "assistant",
                     },
                     finish_reason=choice.finish_reason or "unknown",

packages/opentelemetry-instrumentation-mistralai/tests/test_chunked_content.py

@@ -0,0 +1,172 @@
+import json
+
+import pytest
+from mistralai.models import ChatCompletionResponse
+from opentelemetry.instrumentation.mistralai import _set_response_attributes
+from opentelemetry.instrumentation.mistralai.utils import TRACELOOP_TRACE_CONTENT
+from opentelemetry.semconv._incubating.attributes import (
+    gen_ai_attributes as GenAIAttributes,
+)
+from opentelemetry.semconv_ai import LLMRequestTypeValues
+
+REASONING = "The user asks for the capital of France. It is Paris."
+ANSWER = "Paris."
+
+# The shape magistral models return: the assistant content is a list of chunks,
+# a thinking chunk (a list of text chunks) followed by the answer.
+MAGISTRAL_RESPONSE = {
+    "id": "b3f1d1f0e8d94f8f9d4b8d7e5a1c2b3d",
+    "object": "chat.completion",
+    "created": 1758000000,
+    "model": "magistral-medium-latest",
+    "choices": [
+        {
+            "index": 0,
+            "message": {
+                "role": "assistant",
+                "content": [
+                    {"type": "thinking", "thinking": [{"type": "text", "text": REASONING}]},
+                    {"type": "text", "text": ANSWER},
+                ],
+                "tool_calls": None,
+            },
+            "finish_reason": "stop",
+        }
+    ],
+    "usage": {"prompt_tokens": 20, "total_tokens": 61, "completion_tokens": 41},
+}
+
+
+@pytest.fixture
+def trace_content(monkeypatch):
+    monkeypatch.setenv(TRACELOOP_TRACE_CONTENT, "true")
+
+
+def test_chunked_content_is_recorded_as_json(tracer_provider, span_exporter, trace_content):
+    response = ChatCompletionResponse.model_validate(MAGISTRAL_RESPONSE)
+    tracer = tracer_provider.get_tracer("test")
+    with tracer.start_as_current_span("mistralai.chat") as span:
+        _set_response_attributes(span, LLMRequestTypeValues.CHAT, response)
+
+    attributes = span_exporter.get_finished_spans()[0].attributes
+    content = attributes.get(f"{GenAIAttributes.GEN_AI_COMPLETION}.0.content")
+    assert content is not None, "the completion content was dropped"
+    chunks = json.loads(content)
+    assert chunks[0]["type"] == "thinking"
+    assert chunks[0]["thinking"][0]["text"] == REASONING
+    assert chunks[1] == {"type": "text", "text": ANSWER}
+    assert attributes.get(f"{GenAIAttributes.GEN_AI_COMPLETION}.0.role") == "assistant"
+    assert attributes.get(f"{GenAIAttributes.GEN_AI_COMPLETION}.0.finish_reason") == "stop"
+
+
+def test_string_content_is_unchanged(tracer_provider, span_exporter, trace_content):
+    plain = dict(MAGISTRAL_RESPONSE)
+    plain["choices"] = [
+        {
+            "index": 0,
+            "message": {"role": "assistant", "content": ANSWER, "tool_calls": None},
+            "finish_reason": "stop",
+        }
+    ]
+    response = ChatCompletionResponse.model_validate(plain)
+    tracer = tracer_provider.get_tracer("test")
+    with tracer.start_as_current_span("mistralai.chat") as span:
+        _set_response_attributes(span, LLMRequestTypeValues.CHAT, response)
+
+    attributes = span_exporter.get_finished_spans()[0].attributes
+    assert attributes.get(f"{GenAIAttributes.GEN_AI_COMPLETION}.0.content") == ANSWER
+
+
+def _magistral_stream():
+    from mistralai.models import (
+        CompletionChunk,
+        CompletionEvent,
+        CompletionResponseStreamChoice,
+        DeltaMessage,
+    )
+
+    def event(content, finish_reason=None):
+        return CompletionEvent(
+            data=CompletionChunk(
+                id="b3f1d1f0e8d94f8f9d4b8d7e5a1c2b3d",
+                model="magistral-medium-latest",
+                choices=[
+                    CompletionResponseStreamChoice(
+                        index=0,
+                        delta=DeltaMessage(role="assistant", content=content),
+                        finish_reason=finish_reason,
+                    )
+                ],
+            )
+        )
+
+    # A streamed reasoning answer: the thinking arrives in pieces, then the text.
+    return [
+        event([{"type": "thinking", "thinking": [{"type": "text", "text": "The user asks "}]}]),
+        event([{"type": "thinking", "thinking": [{"type": "text", "text": "for Paris."}]}]),
+        event([{"type": "text", "text": "Par"}]),
+        event([{"type": "text", "text": "is."}]),
+        event(None, finish_reason="stop"),
+    ]
+
+
+def _assert_streamed_completion(span_exporter):
+    attributes = span_exporter.get_finished_spans()[0].attributes
+    content = attributes.get(f"{GenAIAttributes.GEN_AI_COMPLETION}.0.content")
+    assert content is not None, "the streamed completion was dropped"
+    chunks = json.loads(content)
+    # Token-by-token deltas are folded into one chunk per block, as in the non-streamed response.
+    assert [c["type"] for c in chunks] == ["thinking", "text"]
+    assert chunks[0]["thinking"] == [{"type": "text", "text": "The user asks for Paris."}]
+    assert chunks[1]["text"] == ANSWER
+    assert attributes.get(f"{GenAIAttributes.GEN_AI_COMPLETION}.0.finish_reason") == "stop"
+
+
+def test_streamed_chunked_content_is_accumulated(tracer_provider, span_exporter, trace_content):
+    from opentelemetry.instrumentation.mistralai import _accumulate_streaming_response
+
+    span = tracer_provider.get_tracer("test").start_span("mistralai.chat")
+    events = list(
+        _accumulate_streaming_response(span, None, LLMRequestTypeValues.CHAT, iter(_magistral_stream()))
+    )
+    assert len(events) == 5, "every event must still reach the caller"
+    _assert_streamed_completion(span_exporter)
+
+
+@pytest.mark.asyncio
+async def test_async_streamed_chunked_content_is_accumulated(
+    tracer_provider, span_exporter, trace_content
+):
+    from opentelemetry.instrumentation.mistralai import _aaccumulate_streaming_response
+
+    async def stream():
+        for event in _magistral_stream():
+            yield event
+
+    span = tracer_provider.get_tracer("test").start_span("mistralai.chat")
+    events = [
+        e
+        async for e in _aaccumulate_streaming_response(
+            span, None, LLMRequestTypeValues.CHAT, stream()
+        )
+    ]
+    assert len(events) == 5
+    _assert_streamed_completion(span_exporter)
+
+
+def test_chunked_content_is_serialised_on_the_events_path(
+    log_exporter, logger_provider, trace_content, monkeypatch
+):
+    from opentelemetry.instrumentation.mistralai import _emit_choice_events
+    from opentelemetry.instrumentation.mistralai.config import Config
+
+    monkeypatch.setattr(Config, "use_legacy_attributes", False)
+    response = ChatCompletionResponse.model_validate(MAGISTRAL_RESPONSE)
+    _emit_choice_events(response, logger_provider.get_logger("test"))
+
+    logs = log_exporter.get_finished_logs()
+    assert len(logs) == 1, "the choice event was dropped"
+    body = logs[0].log_record.body
+    chunks = json.loads(body["message"]["content"])
+    assert chunks[0]["type"] == "thinking"
+    assert chunks[1] == {"type": "text", "text": ANSWER}