diff --git a/pyproject.toml b/pyproject.toml index be262286..c61f4d43 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,10 @@ agents = [ realtime = [ "websockets >=13.0", ] +telemetry = [ + "opentelemetry-sdk (>=1.33.1,<2.0.0)", + "opentelemetry-exporter-otlp-proto-http (>=1.33.1,<2.0.0)", +] workflow_payload_offloading_azure = [ "azure-storage-blob[aio]>=12.28.0,<13.0.0", @@ -60,6 +64,7 @@ dev = [ "pyyaml>=6.0.2,<7", "mypy==1.15.0", "opentelemetry-sdk (>=1.33.1,<2.0.0)", + "opentelemetry-exporter-otlp-proto-http (>=1.33.1,<2.0.0)", "pylint==3.2.3", "pytest>=8.2.2,<9", "pytest-asyncio>=0.23.7,<0.24", diff --git a/src/mistralai/client/_hooks/registration.py b/src/mistralai/client/_hooks/registration.py index 4da6eb7c..e6650de5 100644 --- a/src/mistralai/client/_hooks/registration.py +++ b/src/mistralai/client/_hooks/registration.py @@ -16,6 +16,8 @@ def init_hooks(hooks: Hooks): with an instance of a hook that implements that specific Hook interface Hooks are registered per SDK instance, and are valid for the lifetime of the SDK instance """ + # Always register tracing: it also supports app-owned global OTel providers + # and per-client providers configured after SDK construction. tracing_hook = TracingHook() workflow_encoding_hook = WorkflowEncodingHook() hooks.register_before_request_hook(CustomUserAgentHook()) diff --git a/src/mistralai/client/_hooks/tracing.py b/src/mistralai/client/_hooks/tracing.py index 14c8cffc..4fb18e0d 100644 --- a/src/mistralai/client/_hooks/tracing.py +++ b/src/mistralai/client/_hooks/tracing.py @@ -1,5 +1,6 @@ import logging -from typing import Optional, Tuple, Union +import weakref +from typing import Any, Optional, Tuple, Union import httpx from opentelemetry import trace @@ -11,6 +12,7 @@ get_traced_request_and_span, get_traced_response, ) +from mistralai.extra.observability.telemetry import configure_telemetry_for_hook from .types import ( AfterErrorContext, AfterErrorHook, @@ -29,11 +31,19 @@ class TracingHook(BeforeRequestHook, AfterSuccessHook, AfterErrorHook): def __init__(self) -> None: self.tracer_provider: Optional[trace.TracerProvider] = None + self._auto_telemetry_provider: Optional[Any] = None + self._telemetry_finalizer: Optional[weakref.finalize] = None + self._telemetry_auto_disabled: bool = False self.tracing_enabled, self.tracer = get_or_create_otel_tracer() def before_request( self, hook_ctx: BeforeRequestContext, request: httpx.Request ) -> Union[httpx.Request, Exception]: + configure_telemetry_for_hook( + self, + hook_ctx.config, + respect_global_provider=True, + ) # Refresh tracer/provider per request so tracing can be enabled if the # application configures OpenTelemetry after the client is instantiated. self.tracing_enabled, self.tracer = get_or_create_otel_tracer( diff --git a/src/mistralai/extra/observability/__init__.py b/src/mistralai/extra/observability/__init__.py index d3ae6cd9..c2d1e0a5 100644 --- a/src/mistralai/extra/observability/__init__.py +++ b/src/mistralai/extra/observability/__init__.py @@ -4,6 +4,11 @@ from opentelemetry import trace as otel_trace from .otel import MISTRAL_SDK_OTEL_TRACER_NAME +from .telemetry import ( + TelemetryConfigurationError, + configure_telemetry, + resolve_telemetry_enabled, +) if TYPE_CHECKING: from mistralai.client.sdk import Mistral @@ -25,6 +30,9 @@ def set_tracer_provider( When set, all SDK spans produced by *client* will be emitted through *provider* instead of the global TracerProvider. + This helper is kept for compatibility. New code can call + configure_telemetry(client, provider=provider) directly. + Usage:: from opentelemetry.sdk.trace import TracerProvider @@ -34,22 +42,13 @@ def set_tracer_provider( client = Mistral(api_key="...") set_tracer_provider(client, TracerProvider()) """ - from mistralai.client._hooks.tracing import TracingHook - - hooks = getattr(client.sdk_configuration, "_hooks", None) - if hooks is None: - raise ValueError( - "Cannot set tracer_provider: SDK hooks not initialised on this client." - ) - - for hook in hooks.before_request_hooks: - if isinstance(hook, TracingHook): - hook.tracer_provider = provider - return - - raise ValueError( - "Cannot set tracer_provider: TracingHook not found in the client's hooks." - ) + configure_telemetry(client, provider=provider) -__all__ = ["trace", "set_tracer_provider"] +__all__ = [ + "TelemetryConfigurationError", + "configure_telemetry", + "resolve_telemetry_enabled", + "set_tracer_provider", + "trace", +] diff --git a/src/mistralai/extra/observability/telemetry.py b/src/mistralai/extra/observability/telemetry.py new file mode 100644 index 00000000..418f1d8e --- /dev/null +++ b/src/mistralai/extra/observability/telemetry.py @@ -0,0 +1,388 @@ +"""Opt-in OpenTelemetry SDK configuration for Mistral telemetry.""" + +from __future__ import annotations + +import logging +import os +import weakref +from typing import TYPE_CHECKING, Any, Literal + +from opentelemetry import trace as otel_trace + +from mistralai.client.utils import get_security_from_env + +from .otel import OTEL_SERVICE_NAME + +if TYPE_CHECKING: + from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider + + from mistralai.client.sdk import Mistral + from mistralai.client.sdkconfiguration import SDKConfiguration + from mistralai.client._hooks.tracing import TracingHook + + +MISTRAL_SDK_TELEMETRY_ENV = "MISTRAL_SDK_TELEMETRY" +MISTRAL_TELEMETRY_ENDPOINT = "https://api.mistral.ai/telemetry/v1/traces" +OTEL_EXPORTER_OTLP_ENDPOINT_ENV = "OTEL_EXPORTER_OTLP_ENDPOINT" +OTEL_EXPORTER_OTLP_TRACES_ENDPOINT_ENV = "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT" +TELEMETRY_PROVIDER_DEDICATED = "dedicated" +TELEMETRY_PROVIDER_GLOBAL = "global" + +_TRUE_VALUES = {"1", "true", "yes", "on"} +_FALSE_VALUES = {"0", "false", "no", "off"} +_PROVIDER_VALUES = {TELEMETRY_PROVIDER_DEDICATED, TELEMETRY_PROVIDER_GLOBAL} + +TelemetryProviderMode = Literal["dedicated", "global"] +TelemetrySetting = bool | str | None + +logger = logging.getLogger(__name__) + + +class TelemetryConfigurationError(RuntimeError): + """Raised when opt-in telemetry cannot be configured.""" + + +def resolve_telemetry_enabled(telemetry: TelemetrySetting = None) -> bool: + """Resolve the telemetry opt-in flag from an explicit value or environment.""" + return _resolve_telemetry_configuration(telemetry)[0] is not None + + +def _resolve_telemetry_configuration( + telemetry: TelemetrySetting = None, +) -> tuple[TelemetryProviderMode | None, bool]: + """Return telemetry provider mode and whether to use OTel env config.""" + mode = ( + _resolve_telemetry_mode(telemetry) + if telemetry is not None + else _resolve_mistral_telemetry_env() + ) + return ( + mode, + mode == TELEMETRY_PROVIDER_DEDICATED + and _has_otel_exporter_endpoint_env(), + ) + + +def _resolve_telemetry_mode(value: bool | str) -> TelemetryProviderMode | None: + if isinstance(value, bool): + return TELEMETRY_PROVIDER_DEDICATED if value else None + + normalized = value.strip().lower() + if normalized in _TRUE_VALUES or normalized == TELEMETRY_PROVIDER_DEDICATED: + return TELEMETRY_PROVIDER_DEDICATED + if normalized == TELEMETRY_PROVIDER_GLOBAL: + return TELEMETRY_PROVIDER_GLOBAL + if normalized in _FALSE_VALUES: + return None + + accepted_values = ", ".join(sorted(_TRUE_VALUES | _FALSE_VALUES | _PROVIDER_VALUES)) + raise TelemetryConfigurationError( + f"Invalid telemetry setting {value!r}. Expected one of: {accepted_values}." + ) + + +def _resolve_provider_mode(value: str) -> TelemetryProviderMode: + normalized = value.strip().lower() + if normalized == TELEMETRY_PROVIDER_DEDICATED: + return TELEMETRY_PROVIDER_DEDICATED + if normalized == TELEMETRY_PROVIDER_GLOBAL: + return TELEMETRY_PROVIDER_GLOBAL + + accepted_values = ", ".join(sorted(_PROVIDER_VALUES)) + raise TelemetryConfigurationError( + f"Invalid telemetry provider {value!r}. Expected one of: {accepted_values}." + ) + + +def _resolve_mistral_telemetry_env() -> TelemetryProviderMode | None: + env_value = os.getenv(MISTRAL_SDK_TELEMETRY_ENV) + if env_value is None or env_value == "": + return None + + try: + return _resolve_telemetry_mode(env_value) + except TelemetryConfigurationError as exc: + accepted_values = ", ".join( + sorted(_TRUE_VALUES | _FALSE_VALUES | _PROVIDER_VALUES) + ) + raise TelemetryConfigurationError( + f"Invalid {MISTRAL_SDK_TELEMETRY_ENV}={env_value!r}. " + f"Expected one of: {accepted_values}." + ) from exc + + +def _has_otel_exporter_endpoint_env() -> bool: + return any( + bool(os.getenv(env_name, "").strip()) + for env_name in ( + OTEL_EXPORTER_OTLP_TRACES_ENDPOINT_ENV, + OTEL_EXPORTER_OTLP_ENDPOINT_ENV, + ) + ) + + +def configure_telemetry( + client: "Mistral", + provider: str | otel_trace.TracerProvider = TELEMETRY_PROVIDER_DEDICATED, +) -> bool: + """Configure telemetry provider mode for a Mistral client. + + By default, this creates an SDK-owned telemetry provider/exporter. Passing + provider="global" clears the per-client provider so SDK spans use the + global OpenTelemetry provider. Passing a TracerProvider attaches it to this + client without taking ownership of its lifecycle. + """ + hooks = getattr(client.sdk_configuration, "_hooks", None) + if hooks is None: + raise ValueError("Cannot configure telemetry: SDK hooks not initialised.") + + hook = _get_tracing_hook(hooks) + if isinstance(provider, str): + provider_mode = _resolve_provider_mode(provider) + if provider_mode == TELEMETRY_PROVIDER_GLOBAL: + return _use_global_tracer_provider(hook, replace_existing=True) + + return configure_telemetry_for_hook( + hook, + client.sdk_configuration, + telemetry=provider_mode, + finalizer_owner=client, + replace_existing=True, + ) + + if isinstance(provider, bool): + raise TelemetryConfigurationError( + "Invalid telemetry provider bool. Expected 'dedicated', 'global', " + "or an OpenTelemetry TracerProvider." + ) + + _attach_custom_tracer_provider(hook, provider) + return True + + +def configure_telemetry_for_hook( + hook: "TracingHook", + sdk_config: "SDKConfiguration", + telemetry: TelemetrySetting = None, + finalizer_owner: Any | None = None, + respect_global_provider: bool = False, + replace_existing: bool = False, +) -> bool: + """Configure telemetry for a tracing hook when the user has opted in.""" + # Fast path: already resolved and no explicit override requested. + if hook._auto_telemetry_provider is not None and telemetry is None: + return True + if telemetry is None and hook._telemetry_auto_disabled: + return False + + telemetry_setting = telemetry + if telemetry_setting is None: + config_setting = getattr(sdk_config, "telemetry", None) + telemetry_setting = ( + config_setting if isinstance(config_setting, (bool, str)) else None + ) + using_env_setting = telemetry_setting is None + + provider_mode, use_otel_env_exporter = _resolve_telemetry_configuration( + telemetry_setting + ) + if provider_mode is None: + _shutdown_telemetry_provider(hook) + hook._telemetry_auto_disabled = True + return False + + if provider_mode == TELEMETRY_PROVIDER_GLOBAL: + return _use_global_tracer_provider( + hook, + replace_existing=replace_existing or not using_env_setting, + ) + + if ( + provider_mode == TELEMETRY_PROVIDER_DEDICATED + and respect_global_provider + and using_env_setting + and _has_real_global_tracer_provider() + ): + logger.debug( + "Skipping Mistral SDK telemetry auto-configuration because a global " + "OpenTelemetry provider is already configured. Call " + "configure_telemetry(client, provider='dedicated') to attach an " + "SDK-owned provider for this client." + ) + hook._telemetry_auto_disabled = True + return False + + if hook._auto_telemetry_provider is not None: + return True + + if hook.tracer_provider is not None: + if not replace_existing: + return False + hook.tracer_provider = None + + api_key = ( + None + if use_otel_env_exporter + else _resolve_api_key_from_security(getattr(sdk_config, "security", None)) + ) + provider = _create_telemetry_tracer_provider( + api_key=api_key, + use_otel_env_exporter=use_otel_env_exporter, + ) + _attach_telemetry_provider(hook, provider, finalizer_owner or sdk_config) + return True + + +def set_tracing_hook_provider( + client: "Mistral", + provider: otel_trace.TracerProvider, +) -> None: + """Attach a provider to the client's tracing hook, replacing auto telemetry.""" + configure_telemetry(client, provider=provider) + + +def _get_tracing_hook(hooks: Any) -> "TracingHook": + from mistralai.client._hooks.tracing import TracingHook + + for hook in hooks.before_request_hooks: + if isinstance(hook, TracingHook): + return hook + + raise ValueError( + "Cannot configure telemetry: TracingHook not found in the client's hooks." + ) + + +def _resolve_api_key_from_security(security: Any) -> str: + from mistralai.client.models import Security + + if callable(security): + security = security() + + if getattr(security, "api_key", None) is None: + security = None + + security = get_security_from_env(security, Security) + api_key = getattr(security, "api_key", None) if security is not None else None + if api_key is None: + raise TelemetryConfigurationError( + "Mistral telemetry requires an API key. Pass api_key=... to the " + "client or set MISTRAL_API_KEY." + ) + + return str(api_key) + + +def _create_telemetry_tracer_provider( + *, + api_key: str | None, + use_otel_env_exporter: bool, +) -> "SDKTracerProvider": + ( + batch_span_processor_cls, + otlp_span_exporter_cls, + resource_cls, + tracer_provider_cls, + ) = _load_otel_sdk() + + if use_otel_env_exporter: + exporter = otlp_span_exporter_cls() + else: + if api_key is None: + raise TelemetryConfigurationError( + "Mistral telemetry requires an API key. Pass api_key=... to the " + "client or set MISTRAL_API_KEY." + ) + exporter = otlp_span_exporter_cls( + endpoint=MISTRAL_TELEMETRY_ENDPOINT, + headers={"Authorization": _as_bearer_token(api_key)}, + ) + provider = tracer_provider_cls( + resource=resource_cls.create({"service.name": OTEL_SERVICE_NAME}) + ) + provider.add_span_processor(batch_span_processor_cls(exporter)) + return provider + + +def _load_otel_sdk(): + try: + from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( + OTLPSpanExporter, + ) + from opentelemetry.sdk.resources import Resource + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import BatchSpanProcessor + except ImportError as exc: + raise TelemetryConfigurationError( + "Mistral telemetry requires optional OpenTelemetry SDK/exporter " + "dependencies. Install them with `pip install 'mistralai[telemetry]'` " + "or `uv add 'mistralai[telemetry]'`." + ) from exc + + return BatchSpanProcessor, OTLPSpanExporter, Resource, TracerProvider + + +def _has_real_global_tracer_provider() -> bool: + return not isinstance( + otel_trace.get_tracer_provider(), + otel_trace.ProxyTracerProvider, + ) + + +def _attach_telemetry_provider( + hook: "TracingHook", + provider: "SDKTracerProvider", + finalizer_owner: Any, +) -> None: + _shutdown_telemetry_provider(hook) + hook.tracer_provider = provider + hook._auto_telemetry_provider = provider + hook._telemetry_auto_disabled = False + hook._telemetry_finalizer = weakref.finalize( + finalizer_owner, provider.shutdown + ) + + +def _attach_custom_tracer_provider( + hook: "TracingHook", + provider: otel_trace.TracerProvider, +) -> None: + _shutdown_telemetry_provider(hook) + hook.tracer_provider = provider + hook._telemetry_auto_disabled = False + + +def _use_global_tracer_provider( + hook: "TracingHook", + *, + replace_existing: bool, +) -> bool: + if ( + hook.tracer_provider is not None + and hook._auto_telemetry_provider is None + and not replace_existing + ): + return False + + _shutdown_telemetry_provider(hook) + hook.tracer_provider = None + hook._telemetry_auto_disabled = True + return True + + +def _shutdown_telemetry_provider(hook: "TracingHook") -> None: + finalizer = hook._telemetry_finalizer + if finalizer is not None: + finalizer.detach() + hook._telemetry_finalizer = None + + provider = hook._auto_telemetry_provider + if provider is not None: + provider.shutdown() + if hook.tracer_provider is provider: + hook.tracer_provider = None + hook._auto_telemetry_provider = None + + +def _as_bearer_token(api_key: str) -> str: + return api_key if api_key.lower().startswith("bearer ") else f"Bearer {api_key}" diff --git a/src/mistralai/extra/tests/test_telemetry.py b/src/mistralai/extra/tests/test_telemetry.py new file mode 100644 index 00000000..4751a587 --- /dev/null +++ b/src/mistralai/extra/tests/test_telemetry.py @@ -0,0 +1,544 @@ +import os +import unittest +from typing import TYPE_CHECKING, cast +from unittest.mock import MagicMock, patch + +from opentelemetry.sdk.trace import TracerProvider + +from mistralai.client._hooks import SDKHooks +from mistralai.client._hooks.tracing import TracingHook +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.telemetry import ( + MISTRAL_TELEMETRY_ENDPOINT, + MISTRAL_SDK_TELEMETRY_ENV, + OTEL_EXPORTER_OTLP_ENDPOINT_ENV, + OTEL_EXPORTER_OTLP_TRACES_ENDPOINT_ENV, + TelemetryConfigurationError, + _create_telemetry_tracer_provider, + configure_telemetry_for_hook, + resolve_telemetry_enabled, +) + +if TYPE_CHECKING: + from mistralai.client.sdk import Mistral + + +class FakeClient: + def __init__(self, sdk_configuration: SDKConfiguration): + self.sdk_configuration = sdk_configuration + + +def _make_client(api_key: str | None = "test-key") -> "Mistral": + sdk_configuration = SDKConfiguration( + client=None, + client_supplied=True, + async_client=None, + async_client_supplied=True, + debug_logger=get_default_logger(), + security=Security(api_key=api_key), + ) + sdk_configuration.__dict__["_hooks"] = SDKHooks() + return cast("Mistral", FakeClient(sdk_configuration)) + + +def _get_tracing_hook(client: "Mistral") -> TracingHook: + hooks = client.sdk_configuration.__dict__["_hooks"] + tracing_hooks = [h for h in hooks.before_request_hooks if isinstance(h, TracingHook)] + assert len(tracing_hooks) == 1 + return tracing_hooks[0] + + +def _configure_for_hook( + client: "Mistral", + telemetry: bool | str | None = None, +) -> bool: + """Helper to call configure_telemetry_for_hook via a client.""" + return configure_telemetry_for_hook( + _get_tracing_hook(client), + client.sdk_configuration, + telemetry=telemetry, + ) + + +class FakeProvider: + def __init__(self): + self.shutdown_called = False + + def shutdown(self): + self.shutdown_called = True + + +class FakeExporter: + instances: list["FakeExporter"] = [] + + def __init__(self, *args: object, **kwargs: object): + self.args = args + self.kwargs = kwargs + FakeExporter.instances.append(self) + + +class FakeResource: + @classmethod + def create(cls, attributes: dict[str, str]) -> dict[str, dict[str, str]]: + return {"resource": attributes} + + +class FakeSpanProcessor: + def __init__(self, exporter: FakeExporter): + self.exporter = exporter + + +class FakeTracerProvider: + def __init__(self, *, resource: object): + self.resource = resource + self.span_processors: list[FakeSpanProcessor] = [] + + def add_span_processor(self, span_processor: FakeSpanProcessor): + self.span_processors.append(span_processor) + + +class TestTelemetryConfiguration(unittest.TestCase): + def setUp(self): + FakeExporter.instances.clear() + + def test_resolve_telemetry_enabled_defaults_to_false(self): + with patch.dict(os.environ, {}, clear=True): + self.assertFalse(resolve_telemetry_enabled()) + + def test_resolve_telemetry_enabled_parses_env_values(self): + for value in ("1", "true", "yes", "on", "dedicated", "global"): + with self.subTest(value=value): + with patch.dict(os.environ, {MISTRAL_SDK_TELEMETRY_ENV: value}): + self.assertTrue(resolve_telemetry_enabled()) + + for value in ("0", "false", "no", "off"): + with self.subTest(value=value): + with patch.dict(os.environ, {MISTRAL_SDK_TELEMETRY_ENV: value}): + self.assertFalse(resolve_telemetry_enabled()) + + def test_otel_traces_endpoint_env_does_not_enable_telemetry(self): + with patch.dict( + os.environ, + {OTEL_EXPORTER_OTLP_TRACES_ENDPOINT_ENV: "http://collector:4318/v1/traces"}, + clear=True, + ): + self.assertFalse(resolve_telemetry_enabled()) + + def test_otel_endpoint_env_does_not_enable_telemetry(self): + with patch.dict( + os.environ, + {OTEL_EXPORTER_OTLP_ENDPOINT_ENV: "http://collector:4318"}, + clear=True, + ): + self.assertFalse(resolve_telemetry_enabled()) + + def test_mistral_env_false_disables_otel_endpoint_env_autoconfiguration(self): + env = { + MISTRAL_SDK_TELEMETRY_ENV: "false", + OTEL_EXPORTER_OTLP_TRACES_ENDPOINT_ENV: "http://collector:4318/v1/traces", + } + + with patch.dict(os.environ, env, clear=True): + self.assertFalse(resolve_telemetry_enabled()) + + def test_invalid_mistral_env_value_raises_configuration_error(self): + with patch.dict(os.environ, {MISTRAL_SDK_TELEMETRY_ENV: "maybe"}, clear=True): + with self.assertRaisesRegex( + TelemetryConfigurationError, + r"dedicated.*global", + ): + resolve_telemetry_enabled() + + def test_configure_telemetry_attaches_per_client_provider(self): + provider = FakeProvider() + + with patch.dict(os.environ, {}, clear=True): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider", + return_value=provider, + ) as create_provider: + client = _make_client(api_key="test-key") + configured = configure_telemetry(client) + + self.assertTrue(configured) + create_provider.assert_called_once_with( + api_key="test-key", + use_otel_env_exporter=False, + ) + self.assertIs(_get_tracing_hook(client).tracer_provider, provider) + + def test_configure_telemetry_accepts_explicit_dedicated_provider_mode(self): + provider = FakeProvider() + + with patch.dict(os.environ, {}, clear=True): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider", + return_value=provider, + ) as create_provider: + client = _make_client(api_key="test-key") + configured = configure_telemetry(client, provider="dedicated") + + self.assertTrue(configured) + create_provider.assert_called_once_with( + api_key="test-key", + use_otel_env_exporter=False, + ) + self.assertIs(_get_tracing_hook(client).tracer_provider, provider) + + def test_configure_telemetry_global_provider_mode_clears_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) + + configured = configure_telemetry(client, provider="global") + + hook = _get_tracing_hook(client) + self.assertTrue(configured) + self.assertTrue(provider.shutdown_called) + self.assertIsNone(hook.tracer_provider) + self.assertIsNone(hook._auto_telemetry_provider) + self.assertTrue(hook._telemetry_auto_disabled) + + def test_configure_telemetry_custom_provider_replaces_auto_without_shutdown(self): + auto_provider = FakeProvider() + custom_provider = FakeProvider() + + with patch.dict(os.environ, {}, clear=True): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider", + return_value=auto_provider, + ): + client = _make_client(api_key="test-key") + configure_telemetry(client) + + configured = configure_telemetry(client, provider=custom_provider) + + hook = _get_tracing_hook(client) + self.assertTrue(configured) + self.assertTrue(auto_provider.shutdown_called) + self.assertFalse(custom_provider.shutdown_called) + self.assertIs(hook.tracer_provider, custom_provider) + self.assertIsNone(hook._auto_telemetry_provider) + + def test_configure_telemetry_dedicated_replaces_custom_without_shutdown(self): + custom_provider = FakeProvider() + dedicated_provider = FakeProvider() + client = _make_client(api_key="test-key") + configure_telemetry(client, provider=custom_provider) + + with patch.dict(os.environ, {}, clear=True): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider", + return_value=dedicated_provider, + ): + configured = configure_telemetry(client, provider="dedicated") + + hook = _get_tracing_hook(client) + self.assertTrue(configured) + self.assertFalse(custom_provider.shutdown_called) + self.assertIs(hook.tracer_provider, dedicated_provider) + self.assertIs(hook._auto_telemetry_provider, dedicated_provider) + + def test_internal_explicit_false_overrides_env_true(self): + with patch.dict(os.environ, {MISTRAL_SDK_TELEMETRY_ENV: "true"}, clear=True): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider" + ) as create_provider: + client = _make_client(api_key="test-key") + configured = _configure_for_hook(client, telemetry=False) + + self.assertFalse(configured) + create_provider.assert_not_called() + self.assertIsNone(_get_tracing_hook(client).tracer_provider) + + def test_internal_explicit_false_disables_auto_telemetry_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) + configured = _configure_for_hook(client, telemetry=False) + + self.assertFalse(configured) + self.assertTrue(provider.shutdown_called) + self.assertIsNone(_get_tracing_hook(client).tracer_provider) + + def test_env_global_uses_global_provider_mode(self): + with patch.dict( + os.environ, + {MISTRAL_SDK_TELEMETRY_ENV: "global"}, + clear=True, + ): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider" + ) as create_provider: + client = _make_client(api_key="test-key") + configured = _configure_for_hook(client) + + hook = _get_tracing_hook(client) + self.assertTrue(configured) + create_provider.assert_not_called() + self.assertIsNone(hook.tracer_provider) + self.assertTrue(hook._telemetry_auto_disabled) + + def test_env_global_does_not_replace_manual_provider(self): + manual_provider = FakeProvider() + client = _make_client(api_key="test-key") + configure_telemetry(client, provider=manual_provider) + + with patch.dict( + os.environ, + {MISTRAL_SDK_TELEMETRY_ENV: "global"}, + clear=True, + ): + configured = _configure_for_hook(client) + + self.assertFalse(configured) + self.assertIs(_get_tracing_hook(client).tracer_provider, manual_provider) + self.assertFalse(manual_provider.shutdown_called) + + def test_env_true_uses_mistral_api_key_fallback(self): + provider = FakeProvider() + env = { + MISTRAL_SDK_TELEMETRY_ENV: "true", + "MISTRAL_API_KEY": "env-key", + } + + with patch.dict(os.environ, env, clear=True): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider", + return_value=provider, + ) as create_provider: + client = _make_client(api_key=None) + configured = _configure_for_hook(client) + + self.assertTrue(configured) + create_provider.assert_called_once_with( + api_key="env-key", + use_otel_env_exporter=False, + ) + self.assertIs(_get_tracing_hook(client).tracer_provider, provider) + + def test_otel_endpoint_env_configures_without_mistral_api_key(self): + provider = FakeProvider() + env = { + OTEL_EXPORTER_OTLP_TRACES_ENDPOINT_ENV: "http://collector:4318/v1/traces", + } + + with patch.dict(os.environ, env, clear=True): + with patch( + "mistralai.extra.observability.telemetry._resolve_api_key_from_security" + ) as resolve_api_key: + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider", + return_value=provider, + ) as create_provider: + client = _make_client(api_key=None) + configured = configure_telemetry(client) + + self.assertTrue(configured) + resolve_api_key.assert_not_called() + create_provider.assert_called_once_with( + api_key=None, + use_otel_env_exporter=True, + ) + self.assertIs(_get_tracing_hook(client).tracer_provider, provider) + + def test_env_true_prefers_otel_endpoint_env_over_mistral_endpoint(self): + provider = FakeProvider() + env = { + MISTRAL_SDK_TELEMETRY_ENV: "true", + OTEL_EXPORTER_OTLP_ENDPOINT_ENV: "http://collector:4318", + } + + with patch.dict(os.environ, env, clear=True): + with patch( + "mistralai.extra.observability.telemetry._resolve_api_key_from_security" + ) as resolve_api_key: + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider", + return_value=provider, + ) as create_provider: + client = _make_client(api_key=None) + configured = _configure_for_hook(client) + + self.assertTrue(configured) + resolve_api_key.assert_not_called() + create_provider.assert_called_once_with( + api_key=None, + use_otel_env_exporter=True, + ) + + def test_sdk_config_global_uses_global_provider_mode(self): + client = _make_client(api_key="test-key") + client.sdk_configuration.__dict__["telemetry"] = "global" + hook = _get_tracing_hook(client) + + with patch.dict(os.environ, {}, clear=True): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider" + ) as create_provider: + configured = configure_telemetry_for_hook(hook, client.sdk_configuration) + + self.assertTrue(configured) + create_provider.assert_not_called() + self.assertIsNone(hook.tracer_provider) + self.assertTrue(hook._telemetry_auto_disabled) + + def test_missing_optional_dependencies_raise_install_hint(self): + with patch.dict(os.environ, {}, clear=True): + with patch( + "mistralai.extra.observability.telemetry._load_otel_sdk", + side_effect=TelemetryConfigurationError( + "Install them with `pip install 'mistralai[telemetry]'`." + ), + ): + client = _make_client(api_key="test-key") + with self.assertRaisesRegex( + TelemetryConfigurationError, + r"mistralai\[telemetry\]", + ): + configure_telemetry(client) + + def test_manual_provider_replaces_auto_telemetry_provider(self): + provider = FakeProvider() + manual_provider = TracerProvider() + + 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) + + set_tracer_provider(client, manual_provider) + + self.assertTrue(provider.shutdown_called) + self.assertIs(_get_tracing_hook(client).tracer_provider, manual_provider) + + def test_configure_telemetry_for_hook_reads_sdk_config_telemetry_flag(self): + provider = FakeProvider() + client = _make_client(api_key="test-key") + client.sdk_configuration.__dict__["telemetry"] = True + hook = _get_tracing_hook(client) + + with patch.dict(os.environ, {}, clear=True): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider", + return_value=provider, + ): + configured = configure_telemetry_for_hook(hook, client.sdk_configuration) + + self.assertTrue(configured) + self.assertIs(hook.tracer_provider, provider) + + def test_auto_configuration_skips_existing_manual_provider(self): + hook = TracingHook() + manual_provider = MagicMock() + hook.tracer_provider = manual_provider + client = _make_client(api_key="test-key") + + with patch.dict(os.environ, {MISTRAL_SDK_TELEMETRY_ENV: "true"}, clear=True): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider" + ) as create_provider: + configured = configure_telemetry_for_hook( + hook, + client.sdk_configuration, + ) + + self.assertFalse(configured) + create_provider.assert_not_called() + self.assertIs(hook.tracer_provider, manual_provider) + + def test_env_auto_configuration_respects_existing_global_provider(self): + client = _make_client(api_key="test-key") + hook = _get_tracing_hook(client) + + with patch.dict(os.environ, {MISTRAL_SDK_TELEMETRY_ENV: "true"}, clear=True): + with patch( + "mistralai.extra.observability.telemetry._has_real_global_tracer_provider", + return_value=True, + ): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider" + ) as create_provider: + with self.assertLogs( + "mistralai.extra.observability.telemetry", + level="DEBUG", + ) as logs: + configured = configure_telemetry_for_hook( + hook, + client.sdk_configuration, + respect_global_provider=True, + ) + + self.assertFalse(configured) + create_provider.assert_not_called() + self.assertIsNone(hook.tracer_provider) + self.assertTrue(hook._telemetry_auto_disabled) + self.assertIn("global OpenTelemetry provider", logs.output[0]) + + def test_mistral_exporter_uses_mistral_endpoint_and_auth(self): + with patch( + "mistralai.extra.observability.telemetry._load_otel_sdk", + return_value=( + FakeSpanProcessor, + FakeExporter, + FakeResource, + FakeTracerProvider, + ), + ): + provider = _create_telemetry_tracer_provider( + api_key="test-key", + use_otel_env_exporter=False, + ) + + self.assertIsInstance(provider, FakeTracerProvider) + self.assertEqual(len(FakeExporter.instances), 1) + self.assertEqual(FakeExporter.instances[0].args, ()) + self.assertEqual( + FakeExporter.instances[0].kwargs, + { + "endpoint": MISTRAL_TELEMETRY_ENDPOINT, + "headers": {"Authorization": "Bearer test-key"}, + }, + ) + + def test_otel_env_exporter_uses_exporter_environment_defaults(self): + with patch( + "mistralai.extra.observability.telemetry._load_otel_sdk", + return_value=( + FakeSpanProcessor, + FakeExporter, + FakeResource, + FakeTracerProvider, + ), + ): + provider = _create_telemetry_tracer_provider( + api_key=None, + use_otel_env_exporter=True, + ) + + self.assertIsInstance(provider, FakeTracerProvider) + self.assertEqual(len(FakeExporter.instances), 1) + self.assertEqual(FakeExporter.instances[0].args, ()) + self.assertEqual(FakeExporter.instances[0].kwargs, {}) + + +if __name__ == "__main__": + unittest.main() diff --git a/uv.lock b/uv.lock index d1b64c13..cd0d41ce 100644 --- a/uv.lock +++ b/uv.lock @@ -839,6 +839,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c6/97/451d55e05487a5cd6279a01a7e34921858b16f7dc8aa38a2c684743cd2b3/google_auth-2.45.0-py2.py3-none-any.whl", hash = "sha256:82344e86dc00410ef5382d99be677c6043d72e502b625aa4f4afa0bdacca0f36", size = 233312, upload-time = "2025-12-15T22:58:40.777Z" }, ] +[[package]] +name = "googleapis-common-protos" +version = "1.75.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b5/c8/f439cffde755cffa462bfbb156278fa6f9d09119719af9814b858fd4f81f/googleapis_common_protos-1.75.0.tar.gz", hash = "sha256:53a062ff3c32552fbd62c11fe23768b78e4ddf0494d5e5fd97d3f4689c75fbbd", size = 151035, upload-time = "2026-05-07T08:04:49.423Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/c8/e2645aa8ed02fd4c7a2f59d68783b65b1f3cbdfe39a6308e156509d1fee8/googleapis_common_protos-1.75.0-py3-none-any.whl", hash = "sha256:961ed60399c457ceb0ee8f285a84c870aabc9c6a832b9d37bb281b5bebde43ed", size = 300631, upload-time = "2026-05-07T08:03:30.345Z" }, +] + [[package]] name = "griffe" version = "1.15.0" @@ -1061,6 +1073,10 @@ gcp = [ realtime = [ { name = "websockets" }, ] +telemetry = [ + { name = "opentelemetry-exporter-otlp-proto-http" }, + { name = "opentelemetry-sdk" }, +] workflow-payload-encryption = [ { name = "cryptography" }, ] @@ -1088,6 +1104,7 @@ dev = [ { name = "invoke" }, { name = "mcp" }, { name = "mypy" }, + { name = "opentelemetry-exporter-otlp-proto-http" }, { name = "opentelemetry-sdk" }, { name = "pylint" }, { name = "pytest" }, @@ -1122,6 +1139,8 @@ requires-dist = [ { name = "mistralai", extras = ["workflow-payload-offloading-gcs"], marker = "extra == 'workflow-payload-offloading'" }, { name = "mistralai", extras = ["workflow-payload-offloading-s3"], marker = "extra == 'workflow-payload-offloading'" }, { name = "opentelemetry-api", specifier = ">=1.33.1,<2.0.0" }, + { name = "opentelemetry-exporter-otlp-proto-http", marker = "extra == 'telemetry'", specifier = ">=1.33.1,<2.0.0" }, + { name = "opentelemetry-sdk", marker = "extra == 'telemetry'", specifier = ">=1.33.1,<2.0.0" }, { name = "opentelemetry-semantic-conventions", specifier = ">=0.60b1,<0.61" }, { name = "pydantic", specifier = ">=2.11.2" }, { name = "python-dateutil", specifier = ">=2.8.2" }, @@ -1129,7 +1148,7 @@ requires-dist = [ { name = "typing-inspection", specifier = ">=0.4.0" }, { name = "websockets", marker = "extra == 'realtime'", specifier = ">=13.0" }, ] -provides-extras = ["gcp", "agents", "realtime", "workflow-payload-offloading-azure", "workflow-payload-offloading-gcs", "workflow-payload-offloading-s3", "workflow-payload-offloading", "workflow-payload-encryption"] +provides-extras = ["gcp", "agents", "realtime", "telemetry", "workflow-payload-offloading-azure", "workflow-payload-offloading-gcs", "workflow-payload-offloading-s3", "workflow-payload-offloading", "workflow-payload-encryption"] [package.metadata.requires-dev] dev = [ @@ -1138,6 +1157,7 @@ dev = [ { name = "invoke", specifier = ">=2.2.0,<3" }, { name = "mcp", specifier = ">=1.0,<2" }, { name = "mypy", specifier = "==1.15.0" }, + { name = "opentelemetry-exporter-otlp-proto-http", specifier = ">=1.33.1,<2.0.0" }, { name = "opentelemetry-sdk", specifier = ">=1.33.1,<2.0.0" }, { name = "pylint", specifier = "==3.2.3" }, { name = "pytest", specifier = ">=8.2.2,<9" }, @@ -1387,6 +1407,48 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cf/df/d3f1ddf4bb4cb50ed9b1139cc7b1c54c34a1e7ce8fd1b9a37c0d1551a6bd/opentelemetry_api-1.39.1-py3-none-any.whl", hash = "sha256:2edd8463432a7f8443edce90972169b195e7d6a05500cd29e6d13898187c9950", size = 66356, upload-time = "2025-12-11T13:32:17.304Z" }, ] +[[package]] +name = "opentelemetry-exporter-otlp-proto-common" +version = "1.39.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-proto" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e9/9d/22d241b66f7bbde88a3bfa6847a351d2c46b84de23e71222c6aae25c7050/opentelemetry_exporter_otlp_proto_common-1.39.1.tar.gz", hash = "sha256:763370d4737a59741c89a67b50f9e39271639ee4afc999dadfe768541c027464", size = 20409, upload-time = "2025-12-11T13:32:40.885Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8c/02/ffc3e143d89a27ac21fd557365b98bd0653b98de8a101151d5805b5d4c33/opentelemetry_exporter_otlp_proto_common-1.39.1-py3-none-any.whl", hash = "sha256:08f8a5862d64cc3435105686d0216c1365dc5701f86844a8cd56597d0c764fde", size = 18366, upload-time = "2025-12-11T13:32:20.2Z" }, +] + +[[package]] +name = "opentelemetry-exporter-otlp-proto-http" +version = "1.39.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "googleapis-common-protos" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-exporter-otlp-proto-common" }, + { name = "opentelemetry-proto" }, + { name = "opentelemetry-sdk" }, + { name = "requests" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/80/04/2a08fa9c0214ae38880df01e8bfae12b067ec0793446578575e5080d6545/opentelemetry_exporter_otlp_proto_http-1.39.1.tar.gz", hash = "sha256:31bdab9745c709ce90a49a0624c2bd445d31a28ba34275951a6a362d16a0b9cb", size = 17288, upload-time = "2025-12-11T13:32:42.029Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/f1/b27d3e2e003cd9a3592c43d099d2ed8d0a947c15281bf8463a256db0b46c/opentelemetry_exporter_otlp_proto_http-1.39.1-py3-none-any.whl", hash = "sha256:d9f5207183dd752a412c4cd564ca8875ececba13be6e9c6c370ffb752fd59985", size = 19641, upload-time = "2025-12-11T13:32:22.248Z" }, +] + +[[package]] +name = "opentelemetry-proto" +version = "1.39.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/49/1d/f25d76d8260c156c40c97c9ed4511ec0f9ce353f8108ca6e7561f82a06b2/opentelemetry_proto-1.39.1.tar.gz", hash = "sha256:6c8e05144fc0d3ed4d22c2289c6b126e03bcd0e6a7da0f16cedd2e1c2772e2c8", size = 46152, upload-time = "2025-12-11T13:32:48.681Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/95/b40c96a7b5203005a0b03d8ce8cd212ff23f1793d5ba289c87a097571b18/opentelemetry_proto-1.39.1-py3-none-any.whl", hash = "sha256:22cdc78efd3b3765d09e68bfbd010d4fc254c9818afd0b6b423387d9dee46007", size = 72535, upload-time = "2025-12-11T13:32:33.866Z" }, +] + [[package]] name = "opentelemetry-sdk" version = "1.39.1" @@ -1555,6 +1617,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5b/5a/bc7b4a4ef808fa59a816c17b20c4bef6884daebbdf627ff2a161da67da19/propcache-0.4.1-py3-none-any.whl", hash = "sha256:af2a6052aeb6cf17d3e46ee169099044fd8224cbaf75c76a2ef596e8163e2237", size = 13305, upload-time = "2025-10-08T19:49:00.792Z" }, ] +[[package]] +name = "protobuf" +version = "6.33.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/66/70/e908e9c5e52ef7c3a6c7902c9dfbb34c7e29c25d2f81ade3856445fd5c94/protobuf-6.33.6.tar.gz", hash = "sha256:a6768d25248312c297558af96a9f9c929e8c4cee0659cb07e780731095f38135", size = 444531, upload-time = "2026-03-18T19:05:00.988Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/9f/2f509339e89cfa6f6a4c4ff50438db9ca488dec341f7e454adad60150b00/protobuf-6.33.6-cp310-abi3-win32.whl", hash = "sha256:7d29d9b65f8afef196f8334e80d6bc1d5d4adedb449971fefd3723824e6e77d3", size = 425739, upload-time = "2026-03-18T19:04:48.373Z" }, + { url = "https://files.pythonhosted.org/packages/76/5d/683efcd4798e0030c1bab27374fd13a89f7c2515fb1f3123efdfaa5eab57/protobuf-6.33.6-cp310-abi3-win_amd64.whl", hash = "sha256:0cd27b587afca21b7cfa59a74dcbd48a50f0a6400cfb59391340ad729d91d326", size = 437089, upload-time = "2026-03-18T19:04:50.381Z" }, + { url = "https://files.pythonhosted.org/packages/5c/01/a3c3ed5cd186f39e7880f8303cc51385a198a81469d53d0fdecf1f64d929/protobuf-6.33.6-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:9720e6961b251bde64edfdab7d500725a2af5280f3f4c87e57c0208376aa8c3a", size = 427737, upload-time = "2026-03-18T19:04:51.866Z" }, + { url = "https://files.pythonhosted.org/packages/ee/90/b3c01fdec7d2f627b3a6884243ba328c1217ed2d978def5c12dc50d328a3/protobuf-6.33.6-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:e2afbae9b8e1825e3529f88d514754e094278bb95eadc0e199751cdd9a2e82a2", size = 324610, upload-time = "2026-03-18T19:04:53.096Z" }, + { url = "https://files.pythonhosted.org/packages/9b/ca/25afc144934014700c52e05103c2421997482d561f3101ff352e1292fb81/protobuf-6.33.6-cp39-abi3-manylinux2014_s390x.whl", hash = "sha256:c96c37eec15086b79762ed265d59ab204dabc53056e3443e702d2681f4b39ce3", size = 339381, upload-time = "2026-03-18T19:04:54.616Z" }, + { url = "https://files.pythonhosted.org/packages/16/92/d1e32e3e0d894fe00b15ce28ad4944ab692713f2e7f0a99787405e43533a/protobuf-6.33.6-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:e9db7e292e0ab79dd108d7f1a94fe31601ce1ee3f7b79e0692043423020b0593", size = 323436, upload-time = "2026-03-18T19:04:55.768Z" }, + { url = "https://files.pythonhosted.org/packages/c4/72/02445137af02769918a93807b2b7890047c32bfb9f90371cbc12688819eb/protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901", size = 170656, upload-time = "2026-03-18T19:04:59.826Z" }, +] + [[package]] name = "pyasn1" version = "0.6.1"