Skip to content

Commit b770bb8

Browse files
authored
feat(otel): read workflow execution id from baggage (#611)
1 parent 22f9040 commit b770bb8

2 files changed

Lines changed: 88 additions & 7 deletions

File tree

‎src/mistralai/extra/observability/otel.py‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,11 @@
4040
os.getenv("MISTRAL_SDK_DEBUG_TRACING", "false").lower() == "true"
4141
)
4242
DEBUG_HINT: str = "To see detailed tracing logs, set MISTRAL_SDK_DEBUG_TRACING=true."
43+
WORKFLOW_EXECUTION_ID_ATTRIBUTE: str = "wf.workflow.execution_id"
44+
BAGGAGE_SPAN_ATTRIBUTES: tuple[str, ...] = (
45+
gen_ai_attributes.GEN_AI_CONVERSATION_ID,
46+
WORKFLOW_EXECUTION_ID_ATTRIBUTE,
47+
)
4348

4449

4550
class MistralAIAttributes:
@@ -127,6 +132,13 @@ def set_available_attributes(span: Span, attributes: dict[str, Any]) -> None:
127132
span.set_attribute(attribute, value)
128133

129134

135+
def _set_baggage_attributes(span: Span) -> None:
136+
for attribute in BAGGAGE_SPAN_ATTRIBUTES:
137+
value = get_baggage(attribute)
138+
if value:
139+
span.set_attribute(attribute, str(value))
140+
141+
130142
def _set_http_attributes(span: Span, operation_id: str, request: httpx.Request) -> None:
131143
"""Set HTTP and server attributes on the span."""
132144
if not request.url.port:
@@ -462,12 +474,7 @@ def get_traced_request_and_span(
462474
try:
463475
span = tracer.start_span(name=operation_id)
464476
span.set_attributes({"agent.trace.public": ""})
465-
# Propagate gen_ai.conversation.id from OTEL baggage if present
466-
conversation_id = get_baggage(gen_ai_attributes.GEN_AI_CONVERSATION_ID)
467-
if conversation_id:
468-
span.set_attribute(
469-
gen_ai_attributes.GEN_AI_CONVERSATION_ID, str(conversation_id)
470-
)
477+
_set_baggage_attributes(span)
471478
# Inject the span context into the request headers to be used by the backend service to continue the trace
472479
propagate.inject(request.headers, context=set_span_in_context(span))
473480
span = enrich_span_from_request(span, operation_id, request)

‎src/mistralai/extra/tests/test_otel_tracing.py‎

Lines changed: 75 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,10 @@
7878
UserMessage,
7979
)
8080
from mistralai.client.sdk import Mistral
81-
from mistralai.extra.observability.otel import TracedResponse
81+
from mistralai.extra.observability.otel import (
82+
WORKFLOW_EXECUTION_ID_ATTRIBUTE,
83+
TracedResponse,
84+
)
8285
from mistralai.extra.run.tools import (
8386
RunFunction,
8487
create_function_result,
@@ -1735,6 +1738,77 @@ def test_no_conversation_id_without_baggage(self):
17351738
span = self._get_single_span()
17361739
self.assertNotIn("gen_ai.conversation.id", span.attributes)
17371740

1741+
# -- Baggage propagation: wf.workflow.execution_id ------------------------
1742+
1743+
def test_workflow_execution_id_from_baggage(self):
1744+
"""When wf.workflow.execution_id is set in OTEL baggage, it must appear as a span attribute."""
1745+
request = ChatCompletionRequest(
1746+
model="mistral-small-latest",
1747+
messages=[UserMessage(content="Hello from a workflow activity")],
1748+
)
1749+
response = ChatCompletionResponse(
1750+
id="cmpl-workflow-baggage-001",
1751+
object="chat.completion",
1752+
model="mistral-small-latest",
1753+
created=1700000012,
1754+
choices=[
1755+
ChatCompletionChoice(
1756+
index=0,
1757+
message=AssistantMessage(content="Hi!", tool_calls=None),
1758+
finish_reason="stop",
1759+
),
1760+
],
1761+
usage=UsageInfo(prompt_tokens=8, completion_tokens=2, total_tokens=10),
1762+
)
1763+
1764+
expected_execution_id = "workflow-exec-from-baggage-123"
1765+
ctx = set_baggage(WORKFLOW_EXECUTION_ID_ATTRIBUTE, expected_execution_id)
1766+
token = context_api.attach(ctx)
1767+
try:
1768+
self._run_hook_lifecycle(
1769+
"chat_completion_v1_chat_completions_post",
1770+
request,
1771+
response,
1772+
)
1773+
finally:
1774+
context_api.detach(token)
1775+
1776+
span = self._get_single_span()
1777+
self.assertEqual(
1778+
span.attributes[WORKFLOW_EXECUTION_ID_ATTRIBUTE],
1779+
expected_execution_id,
1780+
)
1781+
1782+
def test_no_workflow_execution_id_without_baggage(self):
1783+
"""When no baggage is set, wf.workflow.execution_id must NOT appear on a chat span."""
1784+
request = ChatCompletionRequest(
1785+
model="mistral-small-latest",
1786+
messages=[UserMessage(content="Hello outside a workflow")],
1787+
)
1788+
response = ChatCompletionResponse(
1789+
id="cmpl-workflow-nobag-001",
1790+
object="chat.completion",
1791+
model="mistral-small-latest",
1792+
created=1700000013,
1793+
choices=[
1794+
ChatCompletionChoice(
1795+
index=0,
1796+
message=AssistantMessage(content="Hi!", tool_calls=None),
1797+
finish_reason="stop",
1798+
),
1799+
],
1800+
usage=UsageInfo(prompt_tokens=6, completion_tokens=2, total_tokens=8),
1801+
)
1802+
1803+
self._run_hook_lifecycle(
1804+
"chat_completion_v1_chat_completions_post",
1805+
request,
1806+
response,
1807+
)
1808+
1809+
span = self._get_single_span()
1810+
self.assertNotIn(WORKFLOW_EXECUTION_ID_ATTRIBUTE, span.attributes)
1811+
17381812
# -- Concurrency: interleaved requests on shared hook ----------------------
17391813

17401814
def test_concurrent_async_requests_get_correct_spans(self):

0 commit comments

Comments
 (0)