Add public telemetry tracer helper (#568)

* Add public telemetry tracer helper

* Fix telemetry helper test typing

Author: rbarbadillo

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",
 ]

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",

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

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