Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/mistralai/extra/observability/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from .telemetry import (
TelemetryConfigurationError,
configure_telemetry,
get_telemetry_tracer,
)

if TYPE_CHECKING:
Expand Down Expand Up @@ -47,6 +48,7 @@ def set_tracer_provider(
__all__ = [
"TelemetryConfigurationError",
"configure_telemetry",
"get_telemetry_tracer",
"set_tracer_provider",
"trace",
]
40 changes: 39 additions & 1 deletion src/mistralai/extra/observability/telemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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",
Expand Down
48 changes: 48 additions & 0 deletions src/mistralai/extra/tests/test_otel_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
72 changes: 71 additions & 1 deletion src/mistralai/extra/tests/test_telemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading