diff --git a/src/mistralai/extra/observability/__init__.py b/src/mistralai/extra/observability/__init__.py index 2087a7a5..772d55b7 100644 --- a/src/mistralai/extra/observability/__init__.py +++ b/src/mistralai/extra/observability/__init__.py @@ -7,6 +7,7 @@ from .telemetry import ( TelemetryConfigurationError, configure_telemetry, + get_telemetry_tracer, ) if TYPE_CHECKING: @@ -47,6 +48,7 @@ def set_tracer_provider( __all__ = [ "TelemetryConfigurationError", "configure_telemetry", + "get_telemetry_tracer", "set_tracer_provider", "trace", ] diff --git a/src/mistralai/extra/observability/telemetry.py b/src/mistralai/extra/observability/telemetry.py index 22f2922d..6845d2cc 100644 --- a/src/mistralai/extra/observability/telemetry.py +++ b/src/mistralai/extra/observability/telemetry.py @@ -11,7 +11,7 @@ from mistralai.client.utils import get_security_from_env -from .otel import OTEL_SERVICE_NAME +from .otel import MISTRAL_SDK_OTEL_TRACER_NAME, OTEL_SERVICE_NAME if TYPE_CHECKING: from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider @@ -125,6 +125,44 @@ def configure_telemetry( return True +def get_telemetry_tracer( + client: "Mistral", + name: str | None = None, +) -> otel_trace.Tracer: + """Return a tracer from the telemetry provider configured for a client. + + Custom and SDK-owned dedicated providers are used directly. Clients + explicitly configured for global telemetry use the OpenTelemetry global + provider. If telemetry is disabled or has not been configured for this + client, a TelemetryConfigurationError is raised. + """ + hooks = getattr(client.sdk_configuration, "_hooks", None) + if hooks is None: + raise ValueError("Cannot get telemetry tracer: SDK hooks not initialised.") + + hook = _get_tracing_hook(hooks) + tracer_name = name or MISTRAL_SDK_OTEL_TRACER_NAME + + if hook.tracer_provider is None and not hook._telemetry_use_global_provider: + configure_telemetry_for_hook( + hook, + client.sdk_configuration, + finalizer_owner=client, + respect_global_provider=True, + ) + + if hook.tracer_provider is not None: + return hook.tracer_provider.get_tracer(tracer_name) + if hook._telemetry_use_global_provider: + return otel_trace.get_tracer(tracer_name) + + raise TelemetryConfigurationError( + "Telemetry is not configured for this client. Call configure_telemetry(client) " + "or configure_telemetry(client, provider='global') before requesting a " + "telemetry tracer." + ) + + def configure_telemetry_for_hook( hook: "TracingHook", sdk_config: "SDKConfiguration", diff --git a/src/mistralai/extra/tests/test_otel_tracing.py b/src/mistralai/extra/tests/test_otel_tracing.py index b1543ae7..5a354071 100644 --- a/src/mistralai/extra/tests/test_otel_tracing.py +++ b/src/mistralai/extra/tests/test_otel_tracing.py @@ -2029,6 +2029,54 @@ async def _run(server_url: str): class TestPerInstanceTracerProvider(unittest.TestCase): """Tests for per-instance tracer_provider support via set_tracer_provider.""" + def test_get_telemetry_tracer_dedicated_provider_captures_app_spans(self): + from mistralai.extra.observability import ( + configure_telemetry, + get_telemetry_tracer, + ) + + exporter = InMemorySpanExporter() + dedicated_provider = TracerProvider() + dedicated_provider.add_span_processor(SimpleSpanProcessor(exporter)) + + with patch.dict(os.environ, {}, clear=True): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider", + return_value=dedicated_provider, + ): + with ( + _ChatCompletionTestServer() as server, + httpx.Client() as http_client, + ): + client = Mistral( + api_key="test-key", + client=http_client, + server_url=server.url, + ) + configure_telemetry(client) + tracer = get_telemetry_tracer(client, "my-agent") + + with tracer.start_as_current_span("invoke_agent"): + client.chat.complete( + model="mistral-small-latest", + messages=_make_user_messages("hello"), + ) + + with tracer.start_as_current_span( + "execute_tool web_search" + ): + pass + + spans = exporter.get_finished_spans() + invoke_span = next(s for s in spans if s.name == "invoke_agent") + chat_span = next(s for s in spans if s.name == "chat mistral-small-latest") + tool_span = next(s for s in spans if s.name == "execute_tool web_search") + + self.assertEqual(chat_span.parent.span_id, invoke_span.context.span_id) + self.assertEqual(tool_span.parent.span_id, invoke_span.context.span_id) + self.assertEqual(chat_span.context.trace_id, invoke_span.context.trace_id) + self.assertEqual(tool_span.context.trace_id, invoke_span.context.trace_id) + def test_custom_provider_captures_spans(self): """Spans go to the instance-specific exporter, not the global provider.""" # Create a standalone provider with its own exporter diff --git a/src/mistralai/extra/tests/test_telemetry.py b/src/mistralai/extra/tests/test_telemetry.py index f1071d26..feae82dd 100644 --- a/src/mistralai/extra/tests/test_telemetry.py +++ b/src/mistralai/extra/tests/test_telemetry.py @@ -10,7 +10,12 @@ from mistralai.client.models import Security from mistralai.client.sdkconfiguration import SDKConfiguration from mistralai.client.utils.logger import get_default_logger -from mistralai.extra.observability import configure_telemetry, set_tracer_provider +from mistralai.extra.observability import ( + configure_telemetry, + get_telemetry_tracer, + set_tracer_provider, +) +from mistralai.extra.observability.otel import MISTRAL_SDK_OTEL_TRACER_NAME from mistralai.extra.observability.telemetry import ( MISTRAL_TELEMETRY_ENDPOINT, MISTRAL_SDK_TELEMETRY_ENV, @@ -64,6 +69,14 @@ def _configure_for_hook( class FakeProvider: def __init__(self): self.shutdown_called = False + self.get_tracer_calls: list[str] = [] + self.tracers: dict[str, MagicMock] = {} + + def get_tracer(self, name: str): + self.get_tracer_calls.append(name) + if name not in self.tracers: + self.tracers[name] = MagicMock(name=f"tracer:{name}") + return self.tracers[name] def shutdown(self): self.shutdown_called = True @@ -277,6 +290,63 @@ def test_configure_telemetry_dedicated_replaces_custom_without_shutdown(self): self.assertIs(hook.tracer_provider, dedicated_provider) self.assertIs(hook._auto_telemetry_provider, dedicated_provider) + def test_get_telemetry_tracer_dedicated_uses_auto_provider(self): + provider = FakeProvider() + + with patch.dict(os.environ, {}, clear=True): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider", + return_value=provider, + ): + client = _make_client(api_key="test-key") + configure_telemetry(client) + tracer = get_telemetry_tracer(client) + + self.assertIs(tracer, provider.tracers[MISTRAL_SDK_OTEL_TRACER_NAME]) + self.assertEqual(provider.get_tracer_calls, [MISTRAL_SDK_OTEL_TRACER_NAME]) + + def test_get_telemetry_tracer_custom_provider_uses_custom_name(self): + provider = FakeProvider() + client = _make_client(api_key="test-key") + + configure_telemetry(client, provider=cast(TracerProvider, provider)) + tracer = get_telemetry_tracer(client, "my-agent") + + self.assertIs(tracer, provider.tracers["my-agent"]) + self.assertEqual(provider.get_tracer_calls, ["my-agent"]) + + def test_get_telemetry_tracer_global_uses_global_provider(self): + client = _make_client(api_key="test-key") + configure_telemetry(client, provider="global") + tracer = MagicMock() + + with patch( + "mistralai.extra.observability.telemetry.otel_trace.get_tracer", + return_value=tracer, + ) as get_tracer: + result = get_telemetry_tracer(client, "my-agent") + + self.assertIs(result, tracer) + get_tracer.assert_called_once_with("my-agent") + + def test_get_telemetry_tracer_disabled_raises_configuration_error(self): + env_cases = ({}, {MISTRAL_SDK_TELEMETRY_ENV: "false"}) + + for env in env_cases: + with self.subTest(env=env): + with patch.dict(os.environ, env, clear=True): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider" + ) as create_provider: + client = _make_client(api_key="test-key") + with self.assertRaisesRegex( + TelemetryConfigurationError, + "Telemetry is not configured", + ): + get_telemetry_tracer(client) + + create_provider.assert_not_called() + def test_internal_explicit_false_overrides_env_dedicated(self): with patch.dict( os.environ, {MISTRAL_SDK_TELEMETRY_ENV: "dedicated"}, clear=True