diff --git a/packages/opentelemetry-instrumentation-mistralai/opentelemetry/instrumentation/mistralai/__init__.py b/packages/opentelemetry-instrumentation-mistralai/opentelemetry/instrumentation/mistralai/__init__.py index aff3bf1d6d..b4b2d986b7 100644 --- a/packages/opentelemetry-instrumentation-mistralai/opentelemetry/instrumentation/mistralai/__init__.py +++ b/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", diff --git a/packages/opentelemetry-instrumentation-mistralai/tests/test_chunked_content.py b/packages/opentelemetry-instrumentation-mistralai/tests/test_chunked_content.py new file mode 100644 index 0000000000..b3b23f05f8 --- /dev/null +++ b/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}