diff --git a/src/mistralai/extra/observability/otel.py b/src/mistralai/extra/observability/otel.py index f71fb301..d637350f 100644 --- a/src/mistralai/extra/observability/otel.py +++ b/src/mistralai/extra/observability/otel.py @@ -40,6 +40,11 @@ os.getenv("MISTRAL_SDK_DEBUG_TRACING", "false").lower() == "true" ) DEBUG_HINT: str = "To see detailed tracing logs, set MISTRAL_SDK_DEBUG_TRACING=true." +WORKFLOW_EXECUTION_ID_ATTRIBUTE: str = "wf.workflow.execution_id" +BAGGAGE_SPAN_ATTRIBUTES: tuple[str, ...] = ( + gen_ai_attributes.GEN_AI_CONVERSATION_ID, + WORKFLOW_EXECUTION_ID_ATTRIBUTE, +) class MistralAIAttributes: @@ -127,6 +132,13 @@ def set_available_attributes(span: Span, attributes: dict[str, Any]) -> None: span.set_attribute(attribute, value) +def _set_baggage_attributes(span: Span) -> None: + for attribute in BAGGAGE_SPAN_ATTRIBUTES: + value = get_baggage(attribute) + if value: + span.set_attribute(attribute, str(value)) + + def _set_http_attributes(span: Span, operation_id: str, request: httpx.Request) -> None: """Set HTTP and server attributes on the span.""" if not request.url.port: @@ -462,12 +474,7 @@ def get_traced_request_and_span( try: span = tracer.start_span(name=operation_id) span.set_attributes({"agent.trace.public": ""}) - # Propagate gen_ai.conversation.id from OTEL baggage if present - conversation_id = get_baggage(gen_ai_attributes.GEN_AI_CONVERSATION_ID) - if conversation_id: - span.set_attribute( - gen_ai_attributes.GEN_AI_CONVERSATION_ID, str(conversation_id) - ) + _set_baggage_attributes(span) # Inject the span context into the request headers to be used by the backend service to continue the trace propagate.inject(request.headers, context=set_span_in_context(span)) span = enrich_span_from_request(span, operation_id, request) diff --git a/src/mistralai/extra/tests/test_otel_tracing.py b/src/mistralai/extra/tests/test_otel_tracing.py index 5a354071..80f6ca24 100644 --- a/src/mistralai/extra/tests/test_otel_tracing.py +++ b/src/mistralai/extra/tests/test_otel_tracing.py @@ -78,7 +78,10 @@ UserMessage, ) from mistralai.client.sdk import Mistral -from mistralai.extra.observability.otel import TracedResponse +from mistralai.extra.observability.otel import ( + WORKFLOW_EXECUTION_ID_ATTRIBUTE, + TracedResponse, +) from mistralai.extra.run.tools import ( RunFunction, create_function_result, @@ -1735,6 +1738,77 @@ def test_no_conversation_id_without_baggage(self): span = self._get_single_span() self.assertNotIn("gen_ai.conversation.id", span.attributes) + # -- Baggage propagation: wf.workflow.execution_id ------------------------ + + def test_workflow_execution_id_from_baggage(self): + """When wf.workflow.execution_id is set in OTEL baggage, it must appear as a span attribute.""" + request = ChatCompletionRequest( + model="mistral-small-latest", + messages=[UserMessage(content="Hello from a workflow activity")], + ) + response = ChatCompletionResponse( + id="cmpl-workflow-baggage-001", + object="chat.completion", + model="mistral-small-latest", + created=1700000012, + choices=[ + ChatCompletionChoice( + index=0, + message=AssistantMessage(content="Hi!", tool_calls=None), + finish_reason="stop", + ), + ], + usage=UsageInfo(prompt_tokens=8, completion_tokens=2, total_tokens=10), + ) + + expected_execution_id = "workflow-exec-from-baggage-123" + ctx = set_baggage(WORKFLOW_EXECUTION_ID_ATTRIBUTE, expected_execution_id) + token = context_api.attach(ctx) + try: + self._run_hook_lifecycle( + "chat_completion_v1_chat_completions_post", + request, + response, + ) + finally: + context_api.detach(token) + + span = self._get_single_span() + self.assertEqual( + span.attributes[WORKFLOW_EXECUTION_ID_ATTRIBUTE], + expected_execution_id, + ) + + def test_no_workflow_execution_id_without_baggage(self): + """When no baggage is set, wf.workflow.execution_id must NOT appear on a chat span.""" + request = ChatCompletionRequest( + model="mistral-small-latest", + messages=[UserMessage(content="Hello outside a workflow")], + ) + response = ChatCompletionResponse( + id="cmpl-workflow-nobag-001", + object="chat.completion", + model="mistral-small-latest", + created=1700000013, + choices=[ + ChatCompletionChoice( + index=0, + message=AssistantMessage(content="Hi!", tool_calls=None), + finish_reason="stop", + ), + ], + usage=UsageInfo(prompt_tokens=6, completion_tokens=2, total_tokens=8), + ) + + self._run_hook_lifecycle( + "chat_completion_v1_chat_completions_post", + request, + response, + ) + + span = self._get_single_span() + self.assertNotIn(WORKFLOW_EXECUTION_ID_ATTRIBUTE, span.attributes) + # -- Concurrency: interleaved requests on shared hook ---------------------- def test_concurrent_async_requests_get_correct_spans(self):