|
78 | 78 | UserMessage, |
79 | 79 | ) |
80 | 80 | 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 | +) |
82 | 85 | from mistralai.extra.run.tools import ( |
83 | 86 | RunFunction, |
84 | 87 | create_function_result, |
@@ -1735,6 +1738,77 @@ def test_no_conversation_id_without_baggage(self): |
1735 | 1738 | span = self._get_single_span() |
1736 | 1739 | self.assertNotIn("gen_ai.conversation.id", span.attributes) |
1737 | 1740 |
|
| 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 | + |
1738 | 1812 | # -- Concurrency: interleaved requests on shared hook ---------------------- |
1739 | 1813 |
|
1740 | 1814 | def test_concurrent_async_requests_get_correct_spans(self): |
|
0 commit comments