From e090ffd283470ec4f658b74188b891f2e48eb1cc Mon Sep 17 00:00:00 2001 From: Simon Van de Kerckhove Date: Wed, 1 Jul 2026 17:59:37 +0200 Subject: [PATCH 01/12] feat: first shot --- src/mistralai/extra/observability/__init__.py | 18 + .../extra/observability/redaction.py | 572 ++++++++++++++++++ .../extra/observability/telemetry.py | 68 ++- src/mistralai/extra/tests/test_redaction.py | 269 ++++++++ src/mistralai/extra/tests/test_telemetry.py | 129 ++++ 5 files changed, 1055 insertions(+), 1 deletion(-) create mode 100644 src/mistralai/extra/observability/redaction.py create mode 100644 src/mistralai/extra/tests/test_redaction.py diff --git a/src/mistralai/extra/observability/__init__.py b/src/mistralai/extra/observability/__init__.py index 772d55b7..3524329d 100644 --- a/src/mistralai/extra/observability/__init__.py +++ b/src/mistralai/extra/observability/__init__.py @@ -4,6 +4,16 @@ from opentelemetry import trace as otel_trace from .otel import MISTRAL_SDK_OTEL_TRACER_NAME +from .redaction import ( + AttributeRedactionPolicy, + CallbackRedactionPolicy, + RedactingSpanExporter, + RedactionPolicy, + RegexRedactionPolicy, + default_redaction_policy, + redact_span, + resolve_policy, +) from .telemetry import ( TelemetryConfigurationError, configure_telemetry, @@ -46,9 +56,17 @@ def set_tracer_provider( __all__ = [ + "AttributeRedactionPolicy", + "CallbackRedactionPolicy", + "RedactingSpanExporter", + "RedactionPolicy", + "RegexRedactionPolicy", "TelemetryConfigurationError", "configure_telemetry", + "default_redaction_policy", "get_telemetry_tracer", + "redact_span", + "resolve_policy", "set_tracer_provider", "trace", ] diff --git a/src/mistralai/extra/observability/redaction.py b/src/mistralai/extra/observability/redaction.py new file mode 100644 index 00000000..87c3b55c --- /dev/null +++ b/src/mistralai/extra/observability/redaction.py @@ -0,0 +1,572 @@ +"""Client-side redaction of telemetry spans before they are exported. + +This module provides an export-time masking layer for OpenTelemetry spans so +PII/secrets never leave the client. It is the primary, reusable primitive: +any OTEL application can wrap the exporter it owns with +:class:`RedactingSpanExporter`, and the Mistral SDK installs it automatically +in ``dedicated`` telemetry mode (see ``configure_telemetry``). + +Design notes +------------ +* Redaction happens at *export* time (inside the batch export thread), off the + request hot path, and covers every span that reaches the wrapped exporter -- + including spans the application itself produces (e.g. tool input/output), + not just spans created by the Mistral SDK. +* Because masking is per-pipeline, it only protects the exporter it wraps. In + ``global``/``custom`` provider modes the application owns the pipeline and + must install the wrapper itself (one line); in ``dedicated`` mode the SDK + owns the exporter and wraps it for you. + +Requires the optional OpenTelemetry SDK dependencies +(``pip install 'mistralai[telemetry]'``). The policy classes are importable +without the SDK; only :class:`RedactingSpanExporter` needs it. +""" + +from __future__ import annotations + +import re +from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Callable, Final, Union + +from opentelemetry.util.types import AttributeValue + +if TYPE_CHECKING: + from opentelemetry.sdk.trace import ReadableSpan + from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult + + +__all__ = [ + "AttributeMaskCallback", + "AttributeRedactionPolicy", + "CallbackRedactionPolicy", + "DEFAULT_PII_SECRET_PATTERNS", + "DEFAULT_REDACTED_VALUE", + "DEFAULT_SAFE_ATTRIBUTE_KEYS", + "DEFAULT_SENSITIVE_ATTRIBUTE_FRAGMENTS", + "DEFAULT_SENSITIVE_ATTRIBUTE_KEYS", + "DEFAULT_TOKEN_PATTERNS", + "RedactingSpanExporter", + "RedactionPolicy", + "RedactionPolicyLike", + "RegexRedactionPolicy", + "default_redaction_policy", + "redact_span", + "resolve_policy", +] + + +# --------------------------------------------------------------------------- # +# Types +# --------------------------------------------------------------------------- # + +#: A user-supplied per-attribute masker: given (key, value), return the value +#: to keep. Return the value unchanged to keep it, a redacted value to mask it, +#: or ``None`` to drop the attribute entirely. Mirrors the langfuse/braintrust +#: ``mask`` callback style. +AttributeMaskCallback = Callable[[str, object], object] + +#: Anything accepted where a policy is expected: a full policy or a callback. +RedactionPolicyLike = Union["RedactionPolicy", AttributeMaskCallback] + +DEFAULT_REDACTED_VALUE: Final[str] = "[REDACTED]" + + +# --------------------------------------------------------------------------- # +# Policy interface +# --------------------------------------------------------------------------- # + + +class RedactionPolicy(ABC): + """Base class for redaction policies. + + Subclasses must implement :meth:`redact_attributes`. Span-name and + status-description redaction default to identity so most policies only + override the attribute logic. + + This is an abstract base class: a subclass that omits + :meth:`redact_attributes` fails at instantiation rather than at export + time (redaction runs in the background batch-export thread, where a late + error is easy to miss). Callers that just want a per-attribute masker can + pass a plain ``(key, value)`` callable instead of subclassing; see + :data:`RedactionPolicyLike` and :func:`resolve_policy`. + """ + + @abstractmethod + def redact_attributes( + self, attributes: Mapping[str, AttributeValue] | None + ) -> dict[str, AttributeValue]: + """Return a new attribute mapping with sensitive data removed. + + Implementations may add ``.redacted_*`` metadata attributes to + preserve non-sensitive shape information (length, count, type). + """ + raise NotImplementedError + + def redact_span_name(self, name: str) -> str: + """Return the span name to export. Defaults to unchanged.""" + return name + + def redact_status_description(self, description: str | None) -> str | None: + """Return the status description to export. Defaults to unchanged.""" + return description + + +# --------------------------------------------------------------------------- # +# Default (blunt, hybrid) policy -- lifted & hardened from +# vibe_sdk/.../otel/attribute_redaction.py +# --------------------------------------------------------------------------- # + +#: Attribute keys whose values are always redacted wholesale. +DEFAULT_SENSITIVE_ATTRIBUTE_KEYS: Final[frozenset[str]] = frozenset({ + "client.address", + "db.query.text", + "db.statement", + "exception.message", + "exception.stacktrace", + "gen_ai.input.messages", + "gen_ai.output.messages", + "gen_ai.tool.definitions", + "http.request.body", + "http.request.header.authorization", + "http.request.header.cookie", + "http.response.body", + "http.response.header.set-cookie", + "http.target", + "http.url", + "server.address", + "url.full", + "url.path", + "url.query", +}) + +#: Attribute keys explicitly known to be safe (never redacted). +DEFAULT_SAFE_ATTRIBUTE_KEYS: Final[frozenset[str]] = frozenset({ + "agent.trace.public", + "client.port", + "error.type", + "exception.type", + "gen_ai.agent.name", + "gen_ai.conversation.id", + "gen_ai.operation.name", + "gen_ai.provider.name", + "gen_ai.request.model", + "gen_ai.response.finish_reasons", + "gen_ai.response.id", + "gen_ai.response.model", + "gen_ai.tool.call.id", + "gen_ai.tool.name", + "gen_ai.tool.type", + "http.request.method", + "http.response.status_code", + "network.protocol.name", + "network.protocol.version", + "server.port", + "url.scheme", +}) + +#: Substrings that, when present in a key, mark its value as sensitive +#: (e.g. ``content``, ``prompt``, ``arguments``, ``token``...). +DEFAULT_SENSITIVE_ATTRIBUTE_FRAGMENTS: Final[frozenset[str]] = frozenset({ + "api_key", + "argument", + "arguments", + "authorization", + "body", + "completion", + "content", + "cookie", + "input", + "message", + "messages", + "output", + "password", + "payload", + "prompt", + "secret", + "set_cookie", + "token", +}) + +#: Regexes applied to the *content* of otherwise-kept string values, catching +#: inline secrets (bearer tokens, ``ghp_``/``xox`` tokens, ...). +DEFAULT_TOKEN_PATTERNS: Final[tuple[re.Pattern[str], ...]] = ( + re.compile(r"(?i)bearer\s+[a-z0-9._\-]+"), + re.compile(r"\bgh[pousr]_[A-Za-z0-9_]{20,}\b"), + re.compile(r"\bxox[baprs]-[A-Za-z0-9-]{10,}\b"), + re.compile(r"\bsk-[A-Za-z0-9]{20,}\b"), +) + +#: Attribute key prefixes never redacted (metrics/usage counters). +_SAFE_KEY_PREFIXES: Final[tuple[str, ...]] = ("gen_ai.usage.",) + +_PRIMITIVE_TYPES: Final[tuple[type, ...]] = (str, bool, int, float) + + +def _redact_text( + text: str, patterns: Sequence[re.Pattern[str]], redacted_value: str +) -> str: + """Return *text* with every match of *patterns* replaced by *redacted_value*.""" + redacted = text + for pattern in patterns: + redacted = pattern.sub(redacted_value, redacted) + return redacted + + +class AttributeRedactionPolicy(RedactionPolicy): + """Key-oriented hybrid policy. Redacts whole values for keys judged + sensitive (explicit set, fragment match, or non-primitive value), then + runs :attr:`token_patterns` over the values it keeps. + + This is the out-of-the-box default: high recall / "safe by default", at + the cost of erasing most prompt/response content. Prefer + :class:`RegexRedactionPolicy` when you need to keep trace utility. + """ + + def __init__( + self, + *, + sensitive_keys: frozenset[str] = DEFAULT_SENSITIVE_ATTRIBUTE_KEYS, + safe_keys: frozenset[str] = DEFAULT_SAFE_ATTRIBUTE_KEYS, + sensitive_fragments: frozenset[str] = DEFAULT_SENSITIVE_ATTRIBUTE_FRAGMENTS, + token_patterns: Sequence[re.Pattern[str]] = DEFAULT_TOKEN_PATTERNS, + redact_non_primitive: bool = True, + redacted_value: str = DEFAULT_REDACTED_VALUE, + ) -> None: + self._sensitive_keys = sensitive_keys + self._safe_keys = safe_keys + self._sensitive_fragments = sensitive_fragments + self._token_patterns = tuple(token_patterns) + self._redact_non_primitive = redact_non_primitive + self._redacted_value = redacted_value + + def _should_redact(self, normalized_key: str, value: object) -> bool: + if normalized_key in self._safe_keys: + return False + if normalized_key.startswith(_SAFE_KEY_PREFIXES): + return False + if normalized_key in self._sensitive_keys: + return True + if self._has_sensitive_fragment(normalized_key): + return True + return self._redact_non_primitive and not isinstance(value, _PRIMITIVE_TYPES) + + def _has_sensitive_fragment(self, normalized_key: str) -> bool: + normalized_words = normalized_key.replace("-", "_").replace(".", "_") + key_fragments = {word for word in normalized_words.split("_") if word} + return any( + fragment in key_fragments or fragment in normalized_words + for fragment in self._sensitive_fragments + ) + + def redact_attributes( + self, attributes: Mapping[str, AttributeValue] | None + ) -> dict[str, AttributeValue]: + redacted: dict[str, AttributeValue] = {} + if attributes is None: + return redacted + + for key, value in attributes.items(): + attribute_key = str(key) + normalized_key = attribute_key.lower() + if self._should_redact(normalized_key, value): + redacted[attribute_key] = self._redacted_value + continue + if isinstance(value, str): + redacted[attribute_key] = _redact_text( + value, self._token_patterns, self._redacted_value + ) + continue + redacted[attribute_key] = value + + return redacted + + def redact_status_description(self, description: str | None) -> str | None: + """Redact error descriptions (they often carry request/response text).""" + if description is None: + return None + return self._redacted_value + + +# --------------------------------------------------------------------------- # +# Content-aware policy (lighter, opt-in) +# --------------------------------------------------------------------------- # + +#: PII/secret content patterns: emails, credit cards, IPs, bearer/api tokens... +DEFAULT_PII_SECRET_PATTERNS: Final[tuple[re.Pattern[str], ...]] = ( + # Secrets / tokens + re.compile(r"(?i)bearer\s+[a-z0-9._\-]+"), + re.compile(r"\bgh[pousr]_[A-Za-z0-9_]{20,}\b"), + re.compile(r"\bxox[baprs]-[A-Za-z0-9-]{10,}\b"), + re.compile(r"\bsk-[A-Za-z0-9]{20,}\b"), + # Email addresses + re.compile(r"\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}\b"), + # Credit-card-like sequences (13-16 digits, optional spaces/dashes) + re.compile(r"\b(?:\d[ -]?){13,16}\b"), + # IPv4 addresses + re.compile(r"\b(?:(?:25[0-5]|2[0-4]\d|1?\d?\d)\.){3}(?:25[0-5]|2[0-4]\d|1?\d?\d)\b"), +) + + +class RegexRedactionPolicy(RedactionPolicy): + """Content-only policy: leaves keys and structure intact, scans string + values and redacts matched substrings. Lower false-positive, preserves + observability value; may miss free-form PII (names, addresses). + """ + + def __init__( + self, + patterns: Sequence[re.Pattern[str]] = DEFAULT_PII_SECRET_PATTERNS, + *, + redacted_value: str = DEFAULT_REDACTED_VALUE, + ) -> None: + self._patterns = tuple(patterns) + self._redacted_value = redacted_value + + def redact_attributes( + self, attributes: Mapping[str, AttributeValue] | None + ) -> dict[str, AttributeValue]: + redacted: dict[str, AttributeValue] = {} + if attributes is None: + return redacted + + for key, value in attributes.items(): + attribute_key = str(key) + if isinstance(value, str): + redacted[attribute_key] = _redact_text( + value, self._patterns, self._redacted_value + ) + continue + redacted[attribute_key] = value + + return redacted + + def redact_span_name(self, name: str) -> str: + return _redact_text(name, self._patterns, self._redacted_value) + + def redact_status_description(self, description: str | None) -> str | None: + if description is None: + return None + return _redact_text(description, self._patterns, self._redacted_value) + + +# --------------------------------------------------------------------------- # +# Callback adapter (langfuse-style) +# --------------------------------------------------------------------------- # + + +class CallbackRedactionPolicy(RedactionPolicy): + """Adapt a plain ``Callable[[key, value], value]`` into a policy. + + The callback is invoked per attribute; return the value to keep, a + redacted value, or ``None`` to drop the attribute. Span name and status + description are left unchanged (the callback operates on attributes only). + """ + + def __init__(self, mask: AttributeMaskCallback) -> None: + self._mask = mask + + def redact_attributes( + self, attributes: Mapping[str, AttributeValue] | None + ) -> dict[str, AttributeValue]: + redacted: dict[str, AttributeValue] = {} + if attributes is None: + return redacted + + for key, value in attributes.items(): + attribute_key = str(key) + masked = self._mask(attribute_key, value) + if masked is None: + continue + redacted[attribute_key] = masked + + return redacted + + +# --------------------------------------------------------------------------- # +# Resolution helpers +# --------------------------------------------------------------------------- # + + +def default_redaction_policy() -> RedactionPolicy: + """Return the default policy (an :class:`AttributeRedactionPolicy`).""" + return AttributeRedactionPolicy() + + +def resolve_policy(policy: RedactionPolicyLike | None) -> RedactionPolicy: + """Coerce a policy/callback/None into a concrete :class:`RedactionPolicy`. + + ``None`` -> :func:`default_redaction_policy`; a callable is wrapped in + :class:`CallbackRedactionPolicy`; a policy is returned as-is. + """ + if policy is None: + return default_redaction_policy() + if isinstance(policy, RedactionPolicy): + return policy + if callable(policy): + return CallbackRedactionPolicy(policy) + raise TypeError( + "redaction policy must be a RedactionPolicy, a callable, or None; " + f"got {type(policy).__name__}." + ) + + +# --------------------------------------------------------------------------- # +# Span reconstruction + exporter wrapper +# span rebuild lifted & hardened from vibe_sdk/.../otel/sanitizer.py +# --------------------------------------------------------------------------- # + + +def _load_span_types() -> Any: + """Import the OpenTelemetry SDK span classes needed to rebuild spans. + + Raises a helpful error when the optional ``[telemetry]`` extra is missing. + """ + try: + from opentelemetry.sdk.resources import Resource + from opentelemetry.sdk.trace import Event, ReadableSpan + from opentelemetry.trace import Link, SpanKind, Status, StatusCode + except ImportError as exc: # pragma: no cover - exercised via install matrix + raise ImportError( + "Telemetry redaction requires the optional OpenTelemetry SDK " + "dependencies. Install them with `pip install 'mistralai[telemetry]'` " + "or `uv add 'mistralai[telemetry]'`." + ) from exc + + return _SpanTypes( + Event=Event, + Link=Link, + ReadableSpan=ReadableSpan, + Resource=Resource, + SpanKind=SpanKind, + Status=Status, + StatusCode=StatusCode, + ) + + +class _SpanTypes: + __slots__ = ( + "Event", + "Link", + "ReadableSpan", + "Resource", + "SpanKind", + "Status", + "StatusCode", + ) + + def __init__(self, **types: Any) -> None: + for name, value in types.items(): + setattr(self, name, value) + + +def redact_span(span: ReadableSpan, policy: RedactionPolicy) -> ReadableSpan: + """Return a redacted copy of a readable span. + + Rebuilds a genuine :class:`ReadableSpan` (attributes, events, links, + resource, status, name) so the wrapped exporter's serialization keeps + working. ReadableSpans are frozen at export, hence reconstruction. + """ + types = _load_span_types() + + attributes = policy.redact_attributes(getattr(span, "attributes", None)) + events = _redact_events(getattr(span, "events", ()) or (), policy, types) + links = _redact_links(getattr(span, "links", ()) or (), policy, types) + resource = _redact_resource(getattr(span, "resource", None), policy, types) + status = _redact_status(getattr(span, "status", None), policy, types) + name = policy.redact_span_name(getattr(span, "name", "") or "") + + return types.ReadableSpan( + name=name, + context=getattr(span, "context", None), + parent=getattr(span, "parent", None), + resource=resource, + attributes=attributes, + events=events, + links=links, + kind=getattr(span, "kind", None) or types.SpanKind.INTERNAL, + status=status, + start_time=getattr(span, "start_time", None), + end_time=getattr(span, "end_time", None), + instrumentation_scope=getattr(span, "instrumentation_scope", None), + ) + + +def _redact_events( + events: Sequence[Any], policy: RedactionPolicy, types: Any +) -> list[Any]: + return [ + types.Event( + name=getattr(event, "name", ""), + attributes=policy.redact_attributes(getattr(event, "attributes", None)), + timestamp=getattr(event, "timestamp", None), + ) + for event in events + ] + + +def _redact_links( + links: Sequence[Any], policy: RedactionPolicy, types: Any +) -> list[Any]: + return [ + types.Link( + context=getattr(link, "context", None), + attributes=policy.redact_attributes(getattr(link, "attributes", None)), + ) + for link in links + ] + + +def _redact_resource(resource: Any, policy: RedactionPolicy, types: Any) -> Any: + if resource is None: + return None + return types.Resource( + attributes=policy.redact_attributes(getattr(resource, "attributes", None)), + schema_url=getattr(resource, "schema_url", ""), + ) + + +def _redact_status(status: Any, policy: RedactionPolicy, types: Any) -> Any: + if status is None: + return types.Status() + status_code = getattr(status, "status_code", None) or types.StatusCode.UNSET + description = policy.redact_status_description(getattr(status, "description", None)) + return types.Status(status_code, description) + + +class RedactingSpanExporter: + """Wrap any ``SpanExporter``; redact every span before delegating export. + + Not a subclass of ``SpanExporter``: it duck-types the exporter interface + (``export`` / ``shutdown`` / ``force_flush``), which is all a span + processor calls. This keeps the module importable without the optional + OpenTelemetry SDK; only actual export needs it. + + Example + ------- + >>> exporter = RedactingSpanExporter(OTLPSpanExporter(...)) + >>> provider.add_span_processor(BatchSpanProcessor(exporter)) + """ + + def __init__( + self, + exporter: SpanExporter, + policy: RedactionPolicyLike | None = None, + ) -> None: + """Raises ``ImportError`` if the OpenTelemetry SDK extra is missing.""" + _load_span_types() # fail fast if the SDK is unavailable + self._exporter = exporter + self._policy = resolve_policy(policy) + + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + """Redact each span, then delegate to the wrapped exporter.""" + redacted = [redact_span(span, self._policy) for span in spans] + return self._exporter.export(redacted) + + def shutdown(self) -> None: + """Delegate to the wrapped exporter.""" + self._exporter.shutdown() + + def force_flush(self, timeout_millis: int = 30_000) -> bool: + """Delegate to the wrapped exporter.""" + return self._exporter.force_flush(timeout_millis) diff --git a/src/mistralai/extra/observability/telemetry.py b/src/mistralai/extra/observability/telemetry.py index 6845d2cc..283aa8d7 100644 --- a/src/mistralai/extra/observability/telemetry.py +++ b/src/mistralai/extra/observability/telemetry.py @@ -12,6 +12,12 @@ from mistralai.client.utils import get_security_from_env from .otel import MISTRAL_SDK_OTEL_TRACER_NAME, OTEL_SERVICE_NAME +from .redaction import ( + RedactingSpanExporter, + RedactionPolicyLike, + default_redaction_policy, + resolve_policy, +) if TYPE_CHECKING: from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider @@ -19,6 +25,7 @@ from mistralai.client.sdk import Mistral from mistralai.client.sdkconfiguration import SDKConfiguration from mistralai.client._hooks.tracing import TracingHook + from .redaction import RedactionPolicy MISTRAL_SDK_TELEMETRY_ENV = "MISTRAL_SDK_TELEMETRY" @@ -58,6 +65,43 @@ def _resolve_telemetry_mode(value: bool | str) -> TelemetryProviderMode | None: ) +def _resolve_redaction_policy( + redaction: RedactionPolicyLike | bool | None, +) -> "RedactionPolicy | None": + """Resolve the ``redaction`` argument into a policy (``None`` disables it). + + Redaction is safe-by-default: ``True``/``None`` yield the default policy, + ``False`` disables redaction entirely, and a policy or ``(key, value)`` + callback is used as-is. + """ + if redaction is False: + return None + if redaction is True or redaction is None: + return default_redaction_policy() + return resolve_policy(redaction) + + +def _warn_redaction_ignored( + redaction: RedactionPolicyLike | bool | None, + mode: str, +) -> None: + """Warn when a redaction override cannot take effect for this provider mode. + + Redaction is only applied in ``dedicated`` mode where the SDK owns the + exporter. In ``global``/``custom`` modes the application owns the export + pipeline and must wrap its own exporter with ``RedactingSpanExporter``. + """ + if redaction is True: + return + logger.warning( + "Telemetry redaction is only applied in 'dedicated' provider mode, where " + "the Mistral SDK owns the exporter. In %r mode the application owns the " + "export pipeline; wrap your exporter with RedactingSpanExporter to redact " + "spans. Ignoring the redaction argument.", + mode, + ) + + def _resolve_provider_mode(value: str) -> TelemetryProviderMode: normalized = value.strip().lower() if normalized == TELEMETRY_PROVIDER_DEDICATED: @@ -89,6 +133,7 @@ def _resolve_mistral_telemetry_env() -> TelemetryProviderMode | None: def configure_telemetry( client: "Mistral", provider: str | otel_trace.TracerProvider = TELEMETRY_PROVIDER_DEDICATED, + redaction: RedactionPolicyLike | bool | None = True, ) -> bool: """Configure telemetry provider mode for a Mistral client. @@ -96,6 +141,13 @@ def configure_telemetry( 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. + + In ``dedicated`` mode, spans are redacted before export (safe by default). + Control this with ``redaction``: ``True`` (default) uses the default + policy, ``False`` disables redaction, and a ``RedactionPolicy`` or + ``(key, value)`` callback customizes it. ``redaction`` has no effect in + ``global``/``custom`` modes, where the application owns the export pipeline + and must wrap its own exporter with ``RedactingSpanExporter``. """ hooks = getattr(client.sdk_configuration, "_hooks", None) if hooks is None: @@ -105,6 +157,7 @@ def configure_telemetry( if isinstance(provider, str): provider_mode = _resolve_provider_mode(provider) if provider_mode == TELEMETRY_PROVIDER_GLOBAL: + _warn_redaction_ignored(redaction, provider_mode) return _use_global_tracer_provider(hook, replace_existing=True) return configure_telemetry_for_hook( @@ -113,6 +166,7 @@ def configure_telemetry( telemetry=provider_mode, finalizer_owner=client, replace_existing=True, + redaction=redaction, ) if isinstance(provider, bool): @@ -121,6 +175,7 @@ def configure_telemetry( "or an OpenTelemetry TracerProvider." ) + _warn_redaction_ignored(redaction, "custom") _attach_custom_tracer_provider(hook, provider) return True @@ -170,8 +225,14 @@ def configure_telemetry_for_hook( finalizer_owner: Any | None = None, respect_global_provider: bool = False, replace_existing: bool = False, + redaction: RedactionPolicyLike | bool | None = True, ) -> bool: - """Configure telemetry for a tracing hook when the user has opted in.""" + """Configure telemetry for a tracing hook when the user has opted in. + + In dedicated mode the SDK-owned OTLP exporter is wrapped with a + ``RedactingSpanExporter`` unless ``redaction`` is ``False`` (safe by + default). See :func:`configure_telemetry` for the accepted values. + """ # Fast path: already resolved and no explicit override requested. if telemetry is None and ( hook._auto_telemetry_provider is not None or hook._telemetry_use_global_provider @@ -232,6 +293,7 @@ def configure_telemetry_for_hook( api_key = _resolve_api_key_from_security(getattr(sdk_config, "security", None)) provider = _create_telemetry_tracer_provider( api_key=api_key, + redaction=redaction, ) _attach_telemetry_provider(hook, provider, finalizer_owner or sdk_config) return True @@ -272,6 +334,7 @@ def _resolve_api_key_from_security(security: Any) -> str: def _create_telemetry_tracer_provider( *, api_key: str | None, + redaction: RedactionPolicyLike | bool | None = True, ) -> "SDKTracerProvider": ( batch_span_processor_cls, @@ -289,6 +352,9 @@ def _create_telemetry_tracer_provider( endpoint=_resolve_mistral_telemetry_endpoint(), headers={"Authorization": _as_bearer_token(api_key)}, ) + policy = _resolve_redaction_policy(redaction) + if policy is not None: + exporter = RedactingSpanExporter(exporter, policy) provider = tracer_provider_cls( resource=resource_cls.create({"service.name": OTEL_SERVICE_NAME}) ) diff --git a/src/mistralai/extra/tests/test_redaction.py b/src/mistralai/extra/tests/test_redaction.py new file mode 100644 index 00000000..a060e1d5 --- /dev/null +++ b/src/mistralai/extra/tests/test_redaction.py @@ -0,0 +1,269 @@ +"""Tests for client-side telemetry redaction.""" + +import unittest + +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import SpanKind, Status, StatusCode + +from mistralai.extra.observability.redaction import ( + DEFAULT_REDACTED_VALUE, + AttributeRedactionPolicy, + CallbackRedactionPolicy, + RedactingSpanExporter, + RedactionPolicy, + RegexRedactionPolicy, + default_redaction_policy, + redact_span, + resolve_policy, +) + + +class TestAttributeRedactionPolicy(unittest.TestCase): + def setUp(self): + self.policy = AttributeRedactionPolicy() + + def test_sensitive_key_redacted_wholesale(self): + out = self.policy.redact_attributes({"gen_ai.input.messages": "hello"}) + self.assertEqual(out["gen_ai.input.messages"], DEFAULT_REDACTED_VALUE) + + def test_safe_key_kept(self): + out = self.policy.redact_attributes({"gen_ai.request.model": "mistral-large"}) + self.assertEqual(out["gen_ai.request.model"], "mistral-large") + + def test_usage_prefix_kept(self): + out = self.policy.redact_attributes({"gen_ai.usage.input_tokens": 42}) + self.assertEqual(out["gen_ai.usage.input_tokens"], 42) + + def test_fragment_match_redacted(self): + out = self.policy.redact_attributes({"custom.prompt.text": "secret prompt"}) + self.assertEqual(out["custom.prompt.text"], DEFAULT_REDACTED_VALUE) + + def test_token_pattern_on_kept_string(self): + out = self.policy.redact_attributes( + {"note": "call token ghp_abcdefghijklmnopqrstuvwxyz0123 now"} + ) + self.assertEqual(out["note"], "call token [REDACTED] now") + + def test_non_primitive_redacted(self): + out = self.policy.redact_attributes({"data": ("a", "b")}) + self.assertEqual(out["data"], DEFAULT_REDACTED_VALUE) + + def test_non_primitive_kept_when_disabled(self): + policy = AttributeRedactionPolicy(redact_non_primitive=False) + out = policy.redact_attributes({"safeish.list": ("a", "b")}) + self.assertEqual(out["safeish.list"], ("a", "b")) + + def test_none_attributes_returns_empty(self): + self.assertEqual(self.policy.redact_attributes(None), {}) + + def test_status_description_redacted(self): + self.assertEqual( + self.policy.redact_status_description("boom: user@x.com"), + DEFAULT_REDACTED_VALUE, + ) + self.assertIsNone(self.policy.redact_status_description(None)) + + def test_span_name_unchanged(self): + self.assertEqual(self.policy.redact_span_name("chat mistral-large"), "chat mistral-large") + + def test_custom_redacted_value(self): + policy = AttributeRedactionPolicy(redacted_value="XXX") + out = policy.redact_attributes({"http.url": "https://x"}) + self.assertEqual(out["http.url"], "XXX") + + +class TestRegexRedactionPolicy(unittest.TestCase): + def setUp(self): + self.policy = RegexRedactionPolicy() + + def test_email_redacted_inline_preserving_structure(self): + out = self.policy.redact_attributes( + {"gen_ai.input.messages": '{"content":"reach me at a@b.com"}'} + ) + self.assertEqual( + out["gen_ai.input.messages"], '{"content":"reach me at [REDACTED]"}' + ) + + def test_token_redacted(self): + out = self.policy.redact_attributes({"h": "Bearer abc.def-ghi"}) + self.assertEqual(out["h"], "[REDACTED]") + + def test_non_matching_string_kept(self): + out = self.policy.redact_attributes({"server.address": "prod-host-1"}) + self.assertEqual(out["server.address"], "prod-host-1") + + def test_non_string_untouched(self): + out = self.policy.redact_attributes({"n": 5, "b": True}) + self.assertEqual(out, {"n": 5, "b": True}) + + def test_span_name_scanned(self): + self.assertEqual(self.policy.redact_span_name("op a@b.com"), "op [REDACTED]") + + def test_status_description_scanned(self): + self.assertEqual( + self.policy.redact_status_description("failed for a@b.com"), + "failed for [REDACTED]", + ) + + +class TestCallbackRedactionPolicy(unittest.TestCase): + def test_mask_applied_per_attribute(self): + policy = CallbackRedactionPolicy( + lambda key, value: "[X]" if "message" in key else value + ) + out = policy.redact_attributes( + {"gen_ai.output.messages": "hi", "gen_ai.request.model": "m"} + ) + self.assertEqual(out, {"gen_ai.output.messages": "[X]", "gen_ai.request.model": "m"}) + + def test_returning_none_drops_attribute(self): + policy = CallbackRedactionPolicy( + lambda key, value: None if key == "drop" else value + ) + out = policy.redact_attributes({"drop": "x", "keep": "y"}) + self.assertEqual(out, {"keep": "y"}) + + +class TestRedactionPolicyABC(unittest.TestCase): + def test_base_class_cannot_be_instantiated(self): + with self.assertRaises(TypeError): + RedactionPolicy() # type: ignore[abstract] + + def test_subclass_without_redact_attributes_cannot_be_instantiated(self): + class Incomplete(RedactionPolicy): + pass + + with self.assertRaises(TypeError): + Incomplete() # type: ignore[abstract] + + def test_subclass_implementing_redact_attributes_instantiates(self): + class Minimal(RedactionPolicy): + def redact_attributes(self, attributes): + return dict(attributes or {}) + + policy = Minimal() + # Concrete identity defaults remain available. + self.assertEqual(policy.redact_span_name("span"), "span") + self.assertEqual(policy.redact_status_description("desc"), "desc") + + +class TestResolvePolicy(unittest.TestCase): + def test_none_returns_default(self): + self.assertIsInstance(resolve_policy(None), AttributeRedactionPolicy) + self.assertIsInstance(default_redaction_policy(), AttributeRedactionPolicy) + + def test_policy_passthrough(self): + policy = RegexRedactionPolicy() + self.assertIs(resolve_policy(policy), policy) + + def test_callable_wrapped(self): + resolved = resolve_policy(lambda k, v: v) + self.assertIsInstance(resolved, CallbackRedactionPolicy) + + def test_invalid_raises_type_error(self): + with self.assertRaises(TypeError): + resolve_policy(123) # type: ignore[arg-type] + + +class TestRedactSpan(unittest.TestCase): + def _make_span(self): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer("test") + with tracer.start_as_current_span("parent", kind=SpanKind.CLIENT) as span: + span.set_attribute("gen_ai.input.messages", "secret") + span.set_attribute("gen_ai.request.model", "mistral-large") + span.add_event("exception", {"exception.message": "boom"}) + span.set_status(Status(StatusCode.ERROR, "boom detail")) + provider.force_flush() + return exporter.get_finished_spans()[0] + + def test_rebuilds_genuine_readable_span(self): + from opentelemetry.sdk.trace import ReadableSpan + + original = self._make_span() + redacted = redact_span(original, default_redaction_policy()) + self.assertIsInstance(redacted, ReadableSpan) + + def test_attributes_redacted(self): + redacted = redact_span(self._make_span(), default_redaction_policy()) + self.assertEqual(redacted.attributes["gen_ai.input.messages"], DEFAULT_REDACTED_VALUE) + self.assertEqual(redacted.attributes["gen_ai.request.model"], "mistral-large") + + def test_event_attributes_redacted(self): + redacted = redact_span(self._make_span(), default_redaction_policy()) + event = redacted.events[0] + self.assertEqual(event.name, "exception") + self.assertEqual(event.attributes["exception.message"], DEFAULT_REDACTED_VALUE) + + def test_status_description_redacted(self): + redacted = redact_span(self._make_span(), default_redaction_policy()) + self.assertEqual(redacted.status.status_code, StatusCode.ERROR) + self.assertEqual(redacted.status.description, DEFAULT_REDACTED_VALUE) + + def test_identity_preserved(self): + original = self._make_span() + redacted = redact_span(original, default_redaction_policy()) + self.assertEqual(redacted.context.span_id, original.context.span_id) + self.assertEqual(redacted.context.trace_id, original.context.trace_id) + + +class TestRedactingSpanExporter(unittest.TestCase): + def _export_through(self, policy=None): + wrapped = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor( + SimpleSpanProcessor(RedactingSpanExporter(wrapped, policy)) + ) + tracer = provider.get_tracer("test") + with tracer.start_as_current_span("chat") as span: + span.set_attribute("gen_ai.output.messages", "leak") + span.set_attribute("gen_ai.request.model", "mistral-large") + provider.force_flush() + return wrapped.get_finished_spans() + + def test_wrapped_exporter_receives_redacted_spans(self): + spans = self._export_through() + self.assertEqual(len(spans), 1) + attrs = spans[0].attributes + self.assertEqual(attrs["gen_ai.output.messages"], DEFAULT_REDACTED_VALUE) + self.assertEqual(attrs["gen_ai.request.model"], "mistral-large") + + def test_custom_policy_used(self): + spans = self._export_through(RegexRedactionPolicy()) + # Regex policy keeps structure; "leak" has no PII pattern -> unchanged. + self.assertEqual(spans[0].attributes["gen_ai.output.messages"], "leak") + + def test_export_returns_wrapped_result(self): + wrapped = InMemorySpanExporter() + exporter = RedactingSpanExporter(wrapped) + self.assertEqual(exporter.export([]), SpanExportResult.SUCCESS) + + def test_shutdown_and_force_flush_delegate(self): + class _Recorder(InMemorySpanExporter): + def __init__(self): + super().__init__() + self.shutdown_called = False + self.flush_called = False + + def shutdown(self): + self.shutdown_called = True + super().shutdown() + + def force_flush(self, timeout_millis=30000): + self.flush_called = True + return True + + recorder = _Recorder() + exporter = RedactingSpanExporter(recorder) + self.assertTrue(exporter.force_flush()) + exporter.shutdown() + self.assertTrue(recorder.flush_called) + self.assertTrue(recorder.shutdown_called) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/mistralai/extra/tests/test_telemetry.py b/src/mistralai/extra/tests/test_telemetry.py index feae82dd..48d7ce91 100644 --- a/src/mistralai/extra/tests/test_telemetry.py +++ b/src/mistralai/extra/tests/test_telemetry.py @@ -16,12 +16,19 @@ set_tracer_provider, ) from mistralai.extra.observability.otel import MISTRAL_SDK_OTEL_TRACER_NAME +from mistralai.extra.observability.redaction import ( + AttributeRedactionPolicy, + CallbackRedactionPolicy, + RedactingSpanExporter, + RegexRedactionPolicy, +) from mistralai.extra.observability.telemetry import ( MISTRAL_TELEMETRY_ENDPOINT, MISTRAL_SDK_TELEMETRY_ENV, MISTRAL_OTLP_TRACES_ENDPOINT_ENV, TelemetryConfigurationError, _create_telemetry_tracer_provider, + _resolve_redaction_policy, configure_telemetry_for_hook, ) @@ -144,6 +151,7 @@ def test_env_dedicated_values_attach_provider(self): self.assertTrue(configured) create_provider.assert_called_once_with( api_key="test-key", + redaction=True, ) self.assertIs(_get_tracing_hook(client).tracer_provider, provider) @@ -207,6 +215,7 @@ def test_configure_telemetry_attaches_per_client_provider(self): self.assertTrue(configured) create_provider.assert_called_once_with( api_key="test-key", + redaction=True, ) self.assertIs(_get_tracing_hook(client).tracer_provider, provider) @@ -224,6 +233,7 @@ def test_configure_telemetry_accepts_explicit_dedicated_provider_mode(self): self.assertTrue(configured) create_provider.assert_called_once_with( api_key="test-key", + redaction=True, ) self.assertIs(_get_tracing_hook(client).tracer_provider, provider) @@ -431,6 +441,7 @@ def test_env_dedicated_uses_mistral_api_key_fallback(self): self.assertTrue(configured) create_provider.assert_called_once_with( api_key="env-key", + redaction=True, ) self.assertIs(_get_tracing_hook(client).tracer_provider, provider) @@ -470,6 +481,7 @@ def test_env_dedicated_ignores_standard_otel_endpoint_env(self): self.assertTrue(configured) create_provider.assert_called_once_with( api_key="test-key", + redaction=True, ) def test_sdk_config_global_uses_global_provider_mode(self): @@ -642,5 +654,122 @@ def test_mistral_endpoint_env_overrides_default_endpoint(self): ) +class TestTelemetryRedaction(unittest.TestCase): + def setUp(self): + FakeExporter.instances.clear() + + def _make_provider(self, **kwargs): + with patch( + "mistralai.extra.observability.telemetry._load_otel_sdk", + return_value=( + FakeSpanProcessor, + FakeExporter, + FakeResource, + FakeTracerProvider, + ), + ): + return _create_telemetry_tracer_provider(api_key="test-key", **kwargs) + + def _exporter_of(self, provider): + self.assertEqual(len(provider.span_processors), 1) + return provider.span_processors[0].exporter + + def test_dedicated_wraps_exporter_by_default(self): + provider = self._make_provider() + exporter = self._exporter_of(provider) + self.assertIsInstance(exporter, RedactingSpanExporter) + # The wrapped exporter is still the single OTLP exporter created. + self.assertEqual(len(FakeExporter.instances), 1) + self.assertIs(exporter._exporter, FakeExporter.instances[0]) + self.assertIsInstance(exporter._policy, AttributeRedactionPolicy) + + def test_redaction_true_wraps_with_default_policy(self): + provider = self._make_provider(redaction=True) + exporter = self._exporter_of(provider) + self.assertIsInstance(exporter, RedactingSpanExporter) + self.assertIsInstance(exporter._policy, AttributeRedactionPolicy) + + def test_redaction_none_wraps_with_default_policy(self): + provider = self._make_provider(redaction=None) + exporter = self._exporter_of(provider) + self.assertIsInstance(exporter, RedactingSpanExporter) + self.assertIsInstance(exporter._policy, AttributeRedactionPolicy) + + def test_redaction_false_leaves_exporter_unwrapped(self): + provider = self._make_provider(redaction=False) + exporter = self._exporter_of(provider) + self.assertNotIsInstance(exporter, RedactingSpanExporter) + self.assertIs(exporter, FakeExporter.instances[0]) + + def test_custom_policy_instance_is_used(self): + policy = RegexRedactionPolicy() + provider = self._make_provider(redaction=policy) + exporter = self._exporter_of(provider) + self.assertIsInstance(exporter, RedactingSpanExporter) + self.assertIs(exporter._policy, policy) + + def test_callback_is_wrapped_in_callback_policy(self): + def mask(key, value): + return value + + provider = self._make_provider(redaction=mask) + exporter = self._exporter_of(provider) + self.assertIsInstance(exporter, RedactingSpanExporter) + self.assertIsInstance(exporter._policy, CallbackRedactionPolicy) + + def test_resolve_redaction_policy_semantics(self): + self.assertIsInstance( + _resolve_redaction_policy(True), AttributeRedactionPolicy + ) + self.assertIsInstance( + _resolve_redaction_policy(None), AttributeRedactionPolicy + ) + self.assertIsNone(_resolve_redaction_policy(False)) + + policy = RegexRedactionPolicy() + self.assertIs(_resolve_redaction_policy(policy), policy) + + resolved = _resolve_redaction_policy(lambda k, v: v) + self.assertIsInstance(resolved, CallbackRedactionPolicy) + + def test_configure_dedicated_threads_redaction(self): + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider" + ) as create_provider: + create_provider.return_value = FakeProvider() + client = _make_client(api_key="test-key") + policy = RegexRedactionPolicy() + configure_telemetry(client, provider="dedicated", redaction=policy) + + create_provider.assert_called_once_with(api_key="test-key", redaction=policy) + + def test_global_mode_warns_when_redaction_customized(self): + client = _make_client(api_key="test-key") + with self.assertLogs( + "mistralai.extra.observability.telemetry", level="WARNING" + ) as logs: + configure_telemetry(client, provider="global", redaction=False) + self.assertIn("only applied in 'dedicated'", logs.output[0]) + + def test_global_mode_does_not_warn_by_default(self): + client = _make_client(api_key="test-key") + logger = __import__( + "logging" + ).getLogger("mistralai.extra.observability.telemetry") + with patch.object(logger, "warning") as warn: + configure_telemetry(client, provider="global") + warn.assert_not_called() + + def test_custom_provider_warns_when_redaction_customized(self): + client = _make_client(api_key="test-key") + with self.assertLogs( + "mistralai.extra.observability.telemetry", level="WARNING" + ) as logs: + configure_telemetry( + client, provider=TracerProvider(), redaction=False + ) + self.assertIn("only applied in 'dedicated'", logs.output[0]) + + if __name__ == "__main__": unittest.main() From 7e8290350c7c7563f8bac1974d3fc1683d139cb0 Mon Sep 17 00:00:00 2001 From: Simon Van de Kerckhove Date: Wed, 1 Jul 2026 19:10:19 +0200 Subject: [PATCH 02/12] fix: tweaks --- .../extra/observability/redaction.py | 412 +++++++----------- .../extra/observability/telemetry.py | 85 ++-- 2 files changed, 201 insertions(+), 296 deletions(-) diff --git a/src/mistralai/extra/observability/redaction.py b/src/mistralai/extra/observability/redaction.py index 87c3b55c..b165cd84 100644 --- a/src/mistralai/extra/observability/redaction.py +++ b/src/mistralai/extra/observability/redaction.py @@ -2,24 +2,11 @@ This module provides an export-time masking layer for OpenTelemetry spans so PII/secrets never leave the client. It is the primary, reusable primitive: -any OTEL application can wrap the exporter it owns with -:class:`RedactingSpanExporter`, and the Mistral SDK installs it automatically -in ``dedicated`` telemetry mode (see ``configure_telemetry``). - -Design notes ------------- -* Redaction happens at *export* time (inside the batch export thread), off the - request hot path, and covers every span that reaches the wrapped exporter -- - including spans the application itself produces (e.g. tool input/output), - not just spans created by the Mistral SDK. -* Because masking is per-pipeline, it only protects the exporter it wraps. In - ``global``/``custom`` provider modes the application owns the pipeline and - must install the wrapper itself (one line); in ``dedicated`` mode the SDK - owns the exporter and wraps it for you. - -Requires the optional OpenTelemetry SDK dependencies -(``pip install 'mistralai[telemetry]'``). The policy classes are importable -without the SDK; only :class:`RedactingSpanExporter` needs it. +any OTEL application can wrap the exporter it owns with RedactingSpanExporter, +and the Mistral SDK installs it automatically in dedicated telemetry mode (see +``configure_telemetry``). + +Requires the telemetry dependency extra to run, not to import. """ from __future__ import annotations @@ -36,71 +23,22 @@ from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult -__all__ = [ - "AttributeMaskCallback", - "AttributeRedactionPolicy", - "CallbackRedactionPolicy", - "DEFAULT_PII_SECRET_PATTERNS", - "DEFAULT_REDACTED_VALUE", - "DEFAULT_SAFE_ATTRIBUTE_KEYS", - "DEFAULT_SENSITIVE_ATTRIBUTE_FRAGMENTS", - "DEFAULT_SENSITIVE_ATTRIBUTE_KEYS", - "DEFAULT_TOKEN_PATTERNS", - "RedactingSpanExporter", - "RedactionPolicy", - "RedactionPolicyLike", - "RegexRedactionPolicy", - "default_redaction_policy", - "redact_span", - "resolve_policy", -] - - -# --------------------------------------------------------------------------- # -# Types -# --------------------------------------------------------------------------- # - -#: A user-supplied per-attribute masker: given (key, value), return the value -#: to keep. Return the value unchanged to keep it, a redacted value to mask it, -#: or ``None`` to drop the attribute entirely. Mirrors the langfuse/braintrust -#: ``mask`` callback style. -AttributeMaskCallback = Callable[[str, object], object] - -#: Anything accepted where a policy is expected: a full policy or a callback. +# User-supplied per-attribute masker: given (key, value), return the value +# to keep. Return the value unchanged to keep it, a redacted value to mask it, +# or None to drop the attribute entirely +AttributeMaskCallback = Callable[[str, AttributeValue], AttributeValue | None] RedactionPolicyLike = Union["RedactionPolicy", AttributeMaskCallback] - DEFAULT_REDACTED_VALUE: Final[str] = "[REDACTED]" -# --------------------------------------------------------------------------- # -# Policy interface -# --------------------------------------------------------------------------- # - - class RedactionPolicy(ABC): - """Base class for redaction policies. - - Subclasses must implement :meth:`redact_attributes`. Span-name and - status-description redaction default to identity so most policies only - override the attribute logic. - - This is an abstract base class: a subclass that omits - :meth:`redact_attributes` fails at instantiation rather than at export - time (redaction runs in the background batch-export thread, where a late - error is easy to miss). Callers that just want a per-attribute masker can - pass a plain ``(key, value)`` callable instead of subclassing; see - :data:`RedactionPolicyLike` and :func:`resolve_policy`. - """ + """Base class for redaction policies.""" @abstractmethod def redact_attributes( self, attributes: Mapping[str, AttributeValue] | None ) -> dict[str, AttributeValue]: - """Return a new attribute mapping with sensitive data removed. - - Implementations may add ``.redacted_*`` metadata attributes to - preserve non-sensitive shape information (length, count, type). - """ + """Return a new attribute mapping with sensitive data removed.""" raise NotImplementedError def redact_span_name(self, name: str) -> str: @@ -112,115 +50,96 @@ def redact_status_description(self, description: str | None) -> str | None: return description -# --------------------------------------------------------------------------- # -# Default (blunt, hybrid) policy -- lifted & hardened from -# vibe_sdk/.../otel/attribute_redaction.py -# --------------------------------------------------------------------------- # - -#: Attribute keys whose values are always redacted wholesale. -DEFAULT_SENSITIVE_ATTRIBUTE_KEYS: Final[frozenset[str]] = frozenset({ - "client.address", - "db.query.text", - "db.statement", - "exception.message", - "exception.stacktrace", - "gen_ai.input.messages", - "gen_ai.output.messages", - "gen_ai.tool.definitions", - "http.request.body", - "http.request.header.authorization", - "http.request.header.cookie", - "http.response.body", - "http.response.header.set-cookie", - "http.target", - "http.url", - "server.address", - "url.full", - "url.path", - "url.query", -}) - -#: Attribute keys explicitly known to be safe (never redacted). -DEFAULT_SAFE_ATTRIBUTE_KEYS: Final[frozenset[str]] = frozenset({ - "agent.trace.public", - "client.port", - "error.type", - "exception.type", - "gen_ai.agent.name", - "gen_ai.conversation.id", - "gen_ai.operation.name", - "gen_ai.provider.name", - "gen_ai.request.model", - "gen_ai.response.finish_reasons", - "gen_ai.response.id", - "gen_ai.response.model", - "gen_ai.tool.call.id", - "gen_ai.tool.name", - "gen_ai.tool.type", - "http.request.method", - "http.response.status_code", - "network.protocol.name", - "network.protocol.version", - "server.port", - "url.scheme", -}) - -#: Substrings that, when present in a key, mark its value as sensitive -#: (e.g. ``content``, ``prompt``, ``arguments``, ``token``...). -DEFAULT_SENSITIVE_ATTRIBUTE_FRAGMENTS: Final[frozenset[str]] = frozenset({ - "api_key", - "argument", - "arguments", - "authorization", - "body", - "completion", - "content", - "cookie", - "input", - "message", - "messages", - "output", - "password", - "payload", - "prompt", - "secret", - "set_cookie", - "token", -}) - -#: Regexes applied to the *content* of otherwise-kept string values, catching -#: inline secrets (bearer tokens, ``ghp_``/``xox`` tokens, ...). +DEFAULT_SENSITIVE_ATTRIBUTE_KEYS: Final[frozenset[str]] = frozenset( + { + "client.address", + "db.query.text", + "db.statement", + "exception.message", + "exception.stacktrace", + "gen_ai.input.messages", + "gen_ai.output.messages", + "gen_ai.tool.definitions", + "http.request.body", + "http.request.header.authorization", + "http.request.header.cookie", + "http.response.body", + "http.response.header.set-cookie", + "http.target", + "http.url", + "server.address", + "url.full", + "url.path", + "url.query", + } +) +DEFAULT_SENSITIVE_ATTRIBUTE_FRAGMENTS: Final[frozenset[str]] = frozenset( + { + "api_key", + "argument", + "arguments", + "authorization", + "body", + "completion", + "content", + "cookie", + "input", + "message", + "messages", + "output", + "password", + "payload", + "prompt", + "secret", + "set_cookie", + "token", + } +) +DEFAULT_SAFE_ATTRIBUTE_KEYS: Final[frozenset[str]] = frozenset( + { + "agent.trace.public", + "client.port", + "error.type", + "exception.type", + "gen_ai.agent.name", + "gen_ai.conversation.id", + "gen_ai.operation.name", + "gen_ai.provider.name", + "gen_ai.request.model", + "gen_ai.response.finish_reasons", + "gen_ai.response.id", + "gen_ai.response.model", + "gen_ai.tool.call.id", + "gen_ai.tool.name", + "gen_ai.tool.type", + "http.request.method", + "http.response.status_code", + "network.protocol.name", + "network.protocol.version", + "server.port", + "url.scheme", + } +) DEFAULT_TOKEN_PATTERNS: Final[tuple[re.Pattern[str], ...]] = ( re.compile(r"(?i)bearer\s+[a-z0-9._\-]+"), re.compile(r"\bgh[pousr]_[A-Za-z0-9_]{20,}\b"), re.compile(r"\bxox[baprs]-[A-Za-z0-9-]{10,}\b"), re.compile(r"\bsk-[A-Za-z0-9]{20,}\b"), ) - -#: Attribute key prefixes never redacted (metrics/usage counters). _SAFE_KEY_PREFIXES: Final[tuple[str, ...]] = ("gen_ai.usage.",) - _PRIMITIVE_TYPES: Final[tuple[type, ...]] = (str, bool, int, float) -def _redact_text( - text: str, patterns: Sequence[re.Pattern[str]], redacted_value: str -) -> str: - """Return *text* with every match of *patterns* replaced by *redacted_value*.""" - redacted = text - for pattern in patterns: - redacted = pattern.sub(redacted_value, redacted) - return redacted +# TODO: also redact Sequence[str] ? class AttributeRedactionPolicy(RedactionPolicy): - """Key-oriented hybrid policy. Redacts whole values for keys judged - sensitive (explicit set, fragment match, or non-primitive value), then - runs :attr:`token_patterns` over the values it keeps. + """Key-oriented hybrid policy. - This is the out-of-the-box default: high recall / "safe by default", at - the cost of erasing most prompt/response content. Prefer - :class:`RegexRedactionPolicy` when you need to keep trace utility. + This is the default policy: high recall, "safe by default", at the cost of erasing most + prompt/response content. It redacts whole values for keys judged sensitive (explicit set, + fragment match, or non-primitive value), then runs token_patterns over the values it keeps + to redact values. """ def __init__( @@ -240,7 +159,8 @@ def __init__( self._redact_non_primitive = redact_non_primitive self._redacted_value = redacted_value - def _should_redact(self, normalized_key: str, value: object) -> bool: + def _should_redact(self, key: str, value: object) -> bool: + normalized_key = key.lower() if normalized_key in self._safe_keys: return False if normalized_key.startswith(_SAFE_KEY_PREFIXES): @@ -267,17 +187,15 @@ def redact_attributes( return redacted for key, value in attributes.items(): - attribute_key = str(key) - normalized_key = attribute_key.lower() - if self._should_redact(normalized_key, value): - redacted[attribute_key] = self._redacted_value + if self._should_redact(key, value): + redacted[key] = self._redacted_value continue if isinstance(value, str): - redacted[attribute_key] = _redact_text( + redacted[key] = _redact_text( value, self._token_patterns, self._redacted_value ) continue - redacted[attribute_key] = value + redacted[key] = value return redacted @@ -288,11 +206,7 @@ def redact_status_description(self, description: str | None) -> str | None: return self._redacted_value -# --------------------------------------------------------------------------- # -# Content-aware policy (lighter, opt-in) -# --------------------------------------------------------------------------- # - -#: PII/secret content patterns: emails, credit cards, IPs, bearer/api tokens... +# TODO: add regexes (all from PII/secrets detector ?) These are placeholders for now DEFAULT_PII_SECRET_PATTERNS: Final[tuple[re.Pattern[str], ...]] = ( # Secrets / tokens re.compile(r"(?i)bearer\s+[a-z0-9._\-]+"), @@ -304,14 +218,18 @@ def redact_status_description(self, description: str | None) -> str | None: # Credit-card-like sequences (13-16 digits, optional spaces/dashes) re.compile(r"\b(?:\d[ -]?){13,16}\b"), # IPv4 addresses - re.compile(r"\b(?:(?:25[0-5]|2[0-4]\d|1?\d?\d)\.){3}(?:25[0-5]|2[0-4]\d|1?\d?\d)\b"), + re.compile( + r"\b(?:(?:25[0-5]|2[0-4]\d|1?\d?\d)\.){3}(?:25[0-5]|2[0-4]\d|1?\d?\d)\b" + ), ) class RegexRedactionPolicy(RedactionPolicy): - """Content-only policy: leaves keys and structure intact, scans string - values and redacts matched substrings. Lower false-positive, preserves - observability value; may miss free-form PII (names, addresses). + """Content-oriented policy based on regexes. + + Leaves keys and structure intact, scans string values and redacts matched substrings. + Fewer false positives than default policy and aims to preserve observability value; + may miss free-form PII or secrets not in the default patterns. """ def __init__( @@ -331,13 +249,12 @@ def redact_attributes( return redacted for key, value in attributes.items(): - attribute_key = str(key) if isinstance(value, str): - redacted[attribute_key] = _redact_text( + redacted[key] = _redact_text( value, self._patterns, self._redacted_value ) continue - redacted[attribute_key] = value + redacted[key] = value return redacted @@ -350,21 +267,15 @@ def redact_status_description(self, description: str | None) -> str | None: return _redact_text(description, self._patterns, self._redacted_value) -# --------------------------------------------------------------------------- # -# Callback adapter (langfuse-style) -# --------------------------------------------------------------------------- # - - class CallbackRedactionPolicy(RedactionPolicy): - """Adapt a plain ``Callable[[key, value], value]`` into a policy. + """Callback-based policy for users to provide custom redaction capabilities. - The callback is invoked per attribute; return the value to keep, a - redacted value, or ``None`` to drop the attribute. Span name and status - description are left unchanged (the callback operates on attributes only). + The callback is invoked per attribute and should return the value to keep or None to drop the attribute. + Span name and status description are left unchanged (the callback operates on attributes only). """ - def __init__(self, mask: AttributeMaskCallback) -> None: - self._mask = mask + def __init__(self, mask_function: AttributeMaskCallback) -> None: + self._mask_function = mask_function def redact_attributes( self, attributes: Mapping[str, AttributeValue] | None @@ -374,31 +285,20 @@ def redact_attributes( return redacted for key, value in attributes.items(): - attribute_key = str(key) - masked = self._mask(attribute_key, value) + masked = self._mask_function(key, value) if masked is None: continue - redacted[attribute_key] = masked + redacted[key] = masked return redacted -# --------------------------------------------------------------------------- # -# Resolution helpers -# --------------------------------------------------------------------------- # - - +# Helpers def default_redaction_policy() -> RedactionPolicy: - """Return the default policy (an :class:`AttributeRedactionPolicy`).""" return AttributeRedactionPolicy() def resolve_policy(policy: RedactionPolicyLike | None) -> RedactionPolicy: - """Coerce a policy/callback/None into a concrete :class:`RedactionPolicy`. - - ``None`` -> :func:`default_redaction_policy`; a callable is wrapped in - :class:`CallbackRedactionPolicy`; a policy is returned as-is. - """ if policy is None: return default_redaction_policy() if isinstance(policy, RedactionPolicy): @@ -411,10 +311,37 @@ def resolve_policy(policy: RedactionPolicyLike | None) -> RedactionPolicy: ) -# --------------------------------------------------------------------------- # -# Span reconstruction + exporter wrapper -# span rebuild lifted & hardened from vibe_sdk/.../otel/sanitizer.py -# --------------------------------------------------------------------------- # +# SpanExporter wrapper +# NOTE: in essence this is a subclass of SpanExporter. It's not typed as such because +# the opentelemetry SDK is an optional dependency, so to keep it importable we duck-type it +# TODO: can't we rely on typing to have static linters verify that ? (i.e. don't inherit, but type) +class RedactingSpanExporter: + """Wrap any SpanExporter to redact spans before delegating export. + + Example + ------- + >>> exporter = RedactingSpanExporter(OTLPSpanExporter(...)) + >>> provider.add_span_processor(BatchSpanProcessor(exporter)) + """ + + def __init__( + self, + exporter: SpanExporter, + policy: RedactionPolicyLike | None = None, + ) -> None: + _load_span_types() # fail fast if the SDK is unavailable + self._exporter = exporter + self._policy = resolve_policy(policy) + + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + redacted = [redact_span(span, self._policy) for span in spans] + return self._exporter.export(redacted) + + def shutdown(self) -> None: + self._exporter.shutdown() + + def force_flush(self, timeout_millis: int = 30_000) -> bool: + return self._exporter.force_flush(timeout_millis) def _load_span_types() -> Any: @@ -426,7 +353,7 @@ def _load_span_types() -> Any: from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import Event, ReadableSpan from opentelemetry.trace import Link, SpanKind, Status, StatusCode - except ImportError as exc: # pragma: no cover - exercised via install matrix + except ImportError as exc: # pragma: no cover raise ImportError( "Telemetry redaction requires the optional OpenTelemetry SDK " "dependencies. Install them with `pip install 'mistralai[telemetry]'` " @@ -461,12 +388,6 @@ def __init__(self, **types: Any) -> None: def redact_span(span: ReadableSpan, policy: RedactionPolicy) -> ReadableSpan: - """Return a redacted copy of a readable span. - - Rebuilds a genuine :class:`ReadableSpan` (attributes, events, links, - resource, status, name) so the wrapped exporter's serialization keeps - working. ReadableSpans are frozen at export, hence reconstruction. - """ types = _load_span_types() attributes = policy.redact_attributes(getattr(span, "attributes", None)) @@ -534,39 +455,12 @@ def _redact_status(status: Any, policy: RedactionPolicy, types: Any) -> Any: return types.Status(status_code, description) -class RedactingSpanExporter: - """Wrap any ``SpanExporter``; redact every span before delegating export. - - Not a subclass of ``SpanExporter``: it duck-types the exporter interface - (``export`` / ``shutdown`` / ``force_flush``), which is all a span - processor calls. This keeps the module importable without the optional - OpenTelemetry SDK; only actual export needs it. - - Example - ------- - >>> exporter = RedactingSpanExporter(OTLPSpanExporter(...)) - >>> provider.add_span_processor(BatchSpanProcessor(exporter)) - """ - - def __init__( - self, - exporter: SpanExporter, - policy: RedactionPolicyLike | None = None, - ) -> None: - """Raises ``ImportError`` if the OpenTelemetry SDK extra is missing.""" - _load_span_types() # fail fast if the SDK is unavailable - self._exporter = exporter - self._policy = resolve_policy(policy) - - def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: - """Redact each span, then delegate to the wrapped exporter.""" - redacted = [redact_span(span, self._policy) for span in spans] - return self._exporter.export(redacted) - - def shutdown(self) -> None: - """Delegate to the wrapped exporter.""" - self._exporter.shutdown() - - def force_flush(self, timeout_millis: int = 30_000) -> bool: - """Delegate to the wrapped exporter.""" - return self._exporter.force_flush(timeout_millis) +def _redact_text( + text: str, + patterns: Sequence[re.Pattern[str]], + redacted_value: str = DEFAULT_REDACTED_VALUE, +) -> str: + redacted = text + for pattern in patterns: + redacted = pattern.sub(redacted_value, redacted) + return redacted diff --git a/src/mistralai/extra/observability/telemetry.py b/src/mistralai/extra/observability/telemetry.py index 283aa8d7..a4934286 100644 --- a/src/mistralai/extra/observability/telemetry.py +++ b/src/mistralai/extra/observability/telemetry.py @@ -65,10 +65,43 @@ def _resolve_telemetry_mode(value: bool | str) -> TelemetryProviderMode | None: ) +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(_PROVIDER_VALUES | {_DISABLED_VALUE})) + raise TelemetryConfigurationError( + f"Invalid {MISTRAL_SDK_TELEMETRY_ENV}={env_value!r}. " + f"Expected one of: {accepted_values}." + ) from exc + + +# TODO: Feels like this is redundant with redaction.resolve_policy and we could merge it there. +# Also does the None bring any value here ? +# Currently None -> default, True -> default, False -> no redaction +# Could have RedactionPolicyLike | bool with True -> default and False -> no redaction ? +# RedactingSpanExporter may need some changes accordingly since it consumes resolve_policy def _resolve_redaction_policy( redaction: RedactionPolicyLike | bool | None, ) -> "RedactionPolicy | None": - """Resolve the ``redaction`` argument into a policy (``None`` disables it). + """Resolve the redaction argument into a policy (``None`` disables it). Redaction is safe-by-default: ``True``/``None`` yield the default policy, ``False`` disables redaction entirely, and a policy or ``(key, value)`` @@ -102,34 +135,6 @@ def _warn_redaction_ignored( ) -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(_PROVIDER_VALUES | {_DISABLED_VALUE})) - raise TelemetryConfigurationError( - f"Invalid {MISTRAL_SDK_TELEMETRY_ENV}={env_value!r}. " - f"Expected one of: {accepted_values}." - ) from exc - - def configure_telemetry( client: "Mistral", provider: str | otel_trace.TracerProvider = TELEMETRY_PROVIDER_DEDICATED, @@ -142,12 +147,18 @@ def configure_telemetry( global OpenTelemetry provider. Passing a TracerProvider attaches it to this client without taking ownership of its lifecycle. - In ``dedicated`` mode, spans are redacted before export (safe by default). - Control this with ``redaction``: ``True`` (default) uses the default - policy, ``False`` disables redaction, and a ``RedactionPolicy`` or - ``(key, value)`` callback customizes it. ``redaction`` has no effect in - ``global``/``custom`` modes, where the application owns the export pipeline - and must wrap its own exporter with ``RedactingSpanExporter``. + In dedicated mode, spans are redacted before export (safe by default). + You can control this with the redaction argument: + - True: (default) uses the default policy. It redacts all possibly sensitive attributes + - False: disables redaction + - Some RedactionPolicy classes (e.g. based on regexes) can be found in the redaction + module and provided here + - You can also provide a (key, value) -> value | None callback to customize how attributes + get redacted. Your function should return the modified attribute value or None to drop + the attribute. + Note that redaction has no effect when using the global provider mode or providing your own + telemetry provider, since your application controls the provider then. In that case, wrap + your exporter with redaction.RedactingSpanExporter to redact span before export. """ hooks = getattr(client.sdk_configuration, "_hooks", None) if hooks is None: @@ -230,8 +241,8 @@ def configure_telemetry_for_hook( """Configure telemetry for a tracing hook when the user has opted in. In dedicated mode the SDK-owned OTLP exporter is wrapped with a - ``RedactingSpanExporter`` unless ``redaction`` is ``False`` (safe by - default). See :func:`configure_telemetry` for the accepted values. + RedactingSpanExporter unless redaction is False (safe by + default). See configure_telemetry for the accepted values. """ # Fast path: already resolved and no explicit override requested. if telemetry is None and ( From e0e910c903d82675d938bc8e5e95b9f7d177ce55 Mon Sep 17 00:00:00 2001 From: Simon Van de Kerckhove Date: Thu, 2 Jul 2026 10:56:59 +0200 Subject: [PATCH 03/12] fix: resolve redaction and policy changes --- .../extra/observability/redaction.py | 13 +++++++ .../extra/observability/telemetry.py | 35 ++++--------------- src/mistralai/extra/tests/test_telemetry.py | 23 ++++-------- 3 files changed, 25 insertions(+), 46 deletions(-) diff --git a/src/mistralai/extra/observability/redaction.py b/src/mistralai/extra/observability/redaction.py index b165cd84..affa7560 100644 --- a/src/mistralai/extra/observability/redaction.py +++ b/src/mistralai/extra/observability/redaction.py @@ -311,6 +311,19 @@ def resolve_policy(policy: RedactionPolicyLike | None) -> RedactionPolicy: ) +def resolve_redaction(redaction: RedactionPolicyLike | bool) -> RedactionPolicy | None: + """Resolve redaction setting into a policy or None to disable redaction. + + True yields the default policy, False disables redaction entirely, + and a policy or (key, value)->value | None callback is used as-is. + """ + if redaction is False: + return None + if redaction is True: + return default_redaction_policy() + return resolve_policy(redaction) + + # SpanExporter wrapper # NOTE: in essence this is a subclass of SpanExporter. It's not typed as such because # the opentelemetry SDK is an optional dependency, so to keep it importable we duck-type it diff --git a/src/mistralai/extra/observability/telemetry.py b/src/mistralai/extra/observability/telemetry.py index a4934286..927c5ee6 100644 --- a/src/mistralai/extra/observability/telemetry.py +++ b/src/mistralai/extra/observability/telemetry.py @@ -15,8 +15,7 @@ from .redaction import ( RedactingSpanExporter, RedactionPolicyLike, - default_redaction_policy, - resolve_policy, + resolve_redaction, ) if TYPE_CHECKING: @@ -25,7 +24,6 @@ from mistralai.client.sdk import Mistral from mistralai.client.sdkconfiguration import SDKConfiguration from mistralai.client._hooks.tracing import TracingHook - from .redaction import RedactionPolicy MISTRAL_SDK_TELEMETRY_ENV = "MISTRAL_SDK_TELEMETRY" @@ -93,29 +91,8 @@ def _resolve_mistral_telemetry_env() -> TelemetryProviderMode | None: ) from exc -# TODO: Feels like this is redundant with redaction.resolve_policy and we could merge it there. -# Also does the None bring any value here ? -# Currently None -> default, True -> default, False -> no redaction -# Could have RedactionPolicyLike | bool with True -> default and False -> no redaction ? -# RedactingSpanExporter may need some changes accordingly since it consumes resolve_policy -def _resolve_redaction_policy( - redaction: RedactionPolicyLike | bool | None, -) -> "RedactionPolicy | None": - """Resolve the redaction argument into a policy (``None`` disables it). - - Redaction is safe-by-default: ``True``/``None`` yield the default policy, - ``False`` disables redaction entirely, and a policy or ``(key, value)`` - callback is used as-is. - """ - if redaction is False: - return None - if redaction is True or redaction is None: - return default_redaction_policy() - return resolve_policy(redaction) - - def _warn_redaction_ignored( - redaction: RedactionPolicyLike | bool | None, + redaction: RedactionPolicyLike | bool, mode: str, ) -> None: """Warn when a redaction override cannot take effect for this provider mode. @@ -138,7 +115,7 @@ def _warn_redaction_ignored( def configure_telemetry( client: "Mistral", provider: str | otel_trace.TracerProvider = TELEMETRY_PROVIDER_DEDICATED, - redaction: RedactionPolicyLike | bool | None = True, + redaction: RedactionPolicyLike | bool = True, ) -> bool: """Configure telemetry provider mode for a Mistral client. @@ -236,7 +213,7 @@ def configure_telemetry_for_hook( finalizer_owner: Any | None = None, respect_global_provider: bool = False, replace_existing: bool = False, - redaction: RedactionPolicyLike | bool | None = True, + redaction: RedactionPolicyLike | bool = True, ) -> bool: """Configure telemetry for a tracing hook when the user has opted in. @@ -345,7 +322,7 @@ def _resolve_api_key_from_security(security: Any) -> str: def _create_telemetry_tracer_provider( *, api_key: str | None, - redaction: RedactionPolicyLike | bool | None = True, + redaction: RedactionPolicyLike | bool = True, ) -> "SDKTracerProvider": ( batch_span_processor_cls, @@ -363,7 +340,7 @@ def _create_telemetry_tracer_provider( endpoint=_resolve_mistral_telemetry_endpoint(), headers={"Authorization": _as_bearer_token(api_key)}, ) - policy = _resolve_redaction_policy(redaction) + policy = resolve_redaction(redaction) if policy is not None: exporter = RedactingSpanExporter(exporter, policy) provider = tracer_provider_cls( diff --git a/src/mistralai/extra/tests/test_telemetry.py b/src/mistralai/extra/tests/test_telemetry.py index 48d7ce91..f2ef7f46 100644 --- a/src/mistralai/extra/tests/test_telemetry.py +++ b/src/mistralai/extra/tests/test_telemetry.py @@ -21,6 +21,7 @@ CallbackRedactionPolicy, RedactingSpanExporter, RegexRedactionPolicy, + resolve_redaction, ) from mistralai.extra.observability.telemetry import ( MISTRAL_TELEMETRY_ENDPOINT, @@ -28,7 +29,6 @@ MISTRAL_OTLP_TRACES_ENDPOINT_ENV, TelemetryConfigurationError, _create_telemetry_tracer_provider, - _resolve_redaction_policy, configure_telemetry_for_hook, ) @@ -689,12 +689,6 @@ def test_redaction_true_wraps_with_default_policy(self): self.assertIsInstance(exporter, RedactingSpanExporter) self.assertIsInstance(exporter._policy, AttributeRedactionPolicy) - def test_redaction_none_wraps_with_default_policy(self): - provider = self._make_provider(redaction=None) - exporter = self._exporter_of(provider) - self.assertIsInstance(exporter, RedactingSpanExporter) - self.assertIsInstance(exporter._policy, AttributeRedactionPolicy) - def test_redaction_false_leaves_exporter_unwrapped(self): provider = self._make_provider(redaction=False) exporter = self._exporter_of(provider) @@ -717,19 +711,14 @@ def mask(key, value): self.assertIsInstance(exporter, RedactingSpanExporter) self.assertIsInstance(exporter._policy, CallbackRedactionPolicy) - def test_resolve_redaction_policy_semantics(self): - self.assertIsInstance( - _resolve_redaction_policy(True), AttributeRedactionPolicy - ) - self.assertIsInstance( - _resolve_redaction_policy(None), AttributeRedactionPolicy - ) - self.assertIsNone(_resolve_redaction_policy(False)) + def test_resolve_redaction_semantics(self): + self.assertIsInstance(resolve_redaction(True), AttributeRedactionPolicy) + self.assertIsNone(resolve_redaction(False)) policy = RegexRedactionPolicy() - self.assertIs(_resolve_redaction_policy(policy), policy) + self.assertIs(resolve_redaction(policy), policy) - resolved = _resolve_redaction_policy(lambda k, v: v) + resolved = resolve_redaction(lambda k, v: v) self.assertIsInstance(resolved, CallbackRedactionPolicy) def test_configure_dedicated_threads_redaction(self): From eb46713c6f7e571fa6e4a7b5996b59dc94466a26 Mon Sep 17 00:00:00 2001 From: Simon Van de Kerckhove Date: Thu, 2 Jul 2026 11:05:22 +0200 Subject: [PATCH 04/12] feat: add secret regexes to default redaction patterns Fold AWS, Google, JWT, PEM and Stripe key patterns into DEFAULT_TOKEN_PATTERNS and compose DEFAULT_PII_SECRET_PATTERNS from it, with tests covering each secret and the composition invariant. Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe --- .../extra/observability/redaction.py | 12 ++++----- src/mistralai/extra/tests/test_redaction.py | 25 +++++++++++++++++++ 2 files changed, 31 insertions(+), 6 deletions(-) diff --git a/src/mistralai/extra/observability/redaction.py b/src/mistralai/extra/observability/redaction.py index affa7560..2293dbf8 100644 --- a/src/mistralai/extra/observability/redaction.py +++ b/src/mistralai/extra/observability/redaction.py @@ -125,6 +125,11 @@ def redact_status_description(self, description: str | None) -> str | None: re.compile(r"\bgh[pousr]_[A-Za-z0-9_]{20,}\b"), re.compile(r"\bxox[baprs]-[A-Za-z0-9-]{10,}\b"), re.compile(r"\bsk-[A-Za-z0-9]{20,}\b"), + re.compile(r"\bAKIA[0-9A-Z]{16}\b"), + re.compile(r"\bAIza[0-9A-Za-z_\-]{35}\b"), + re.compile(r"\beyJ[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+\b"), + re.compile(r"-----BEGIN [A-Z ]*PRIVATE KEY-----"), + re.compile(r"\b[sr]k_(?:live|test)_[0-9A-Za-z]{10,}\b"), ) _SAFE_KEY_PREFIXES: Final[tuple[str, ...]] = ("gen_ai.usage.",) _PRIMITIVE_TYPES: Final[tuple[type, ...]] = (str, bool, int, float) @@ -206,13 +211,8 @@ def redact_status_description(self, description: str | None) -> str | None: return self._redacted_value -# TODO: add regexes (all from PII/secrets detector ?) These are placeholders for now DEFAULT_PII_SECRET_PATTERNS: Final[tuple[re.Pattern[str], ...]] = ( - # Secrets / tokens - re.compile(r"(?i)bearer\s+[a-z0-9._\-]+"), - re.compile(r"\bgh[pousr]_[A-Za-z0-9_]{20,}\b"), - re.compile(r"\bxox[baprs]-[A-Za-z0-9-]{10,}\b"), - re.compile(r"\bsk-[A-Za-z0-9]{20,}\b"), + *DEFAULT_TOKEN_PATTERNS, # Email addresses re.compile(r"\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}\b"), # Credit-card-like sequences (13-16 digits, optional spaces/dashes) diff --git a/src/mistralai/extra/tests/test_redaction.py b/src/mistralai/extra/tests/test_redaction.py index a060e1d5..57cada7d 100644 --- a/src/mistralai/extra/tests/test_redaction.py +++ b/src/mistralai/extra/tests/test_redaction.py @@ -8,7 +8,9 @@ from opentelemetry.trace import SpanKind, Status, StatusCode from mistralai.extra.observability.redaction import ( + DEFAULT_PII_SECRET_PATTERNS, DEFAULT_REDACTED_VALUE, + DEFAULT_TOKEN_PATTERNS, AttributeRedactionPolicy, CallbackRedactionPolicy, RedactingSpanExporter, @@ -107,6 +109,29 @@ def test_status_description_scanned(self): "failed for [REDACTED]", ) + def test_secret_patterns_redacted(self): + secrets = { + "aws": "AKIAIOSFODNN7EXAMPLE", + "google": "AIzaabcdefghijklmnopqrstuvwxyz012345678", + "jwt": "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.abc123", + "pem": "-----BEGIN RSA PRIVATE KEY-----", + "stripe": "sk_live_0123456789abcdefghij", + } + for name, secret in secrets.items(): + with self.subTest(secret=name): + out = self.policy.redact_attributes({"v": f"leak {secret} here"}) + self.assertNotIn(secret, out["v"]) + self.assertIn(DEFAULT_REDACTED_VALUE, out["v"]) + + +class TestDefaultPatternComposition(unittest.TestCase): + def test_pii_patterns_extend_token_patterns(self): + prefix = DEFAULT_PII_SECRET_PATTERNS[: len(DEFAULT_TOKEN_PATTERNS)] + self.assertEqual(prefix, DEFAULT_TOKEN_PATTERNS) + self.assertGreater( + len(DEFAULT_PII_SECRET_PATTERNS), len(DEFAULT_TOKEN_PATTERNS) + ) + class TestCallbackRedactionPolicy(unittest.TestCase): def test_mask_applied_per_attribute(self): From 899129c4fb88fd260b330b9251b679b0065b8b47 Mon Sep 17 00:00:00 2001 From: Simon Van de Kerckhove Date: Thu, 2 Jul 2026 11:09:45 +0200 Subject: [PATCH 05/12] refactor: type RedactingSpanExporter against SpanExporter base Use a conditional base (SpanExporter under TYPE_CHECKING, object at runtime) so linters verify the export/shutdown/force_flush overrides while keeping the OpenTelemetry SDK an optional import. Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe --- src/mistralai/extra/observability/redaction.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/mistralai/extra/observability/redaction.py b/src/mistralai/extra/observability/redaction.py index 2293dbf8..2620d94f 100644 --- a/src/mistralai/extra/observability/redaction.py +++ b/src/mistralai/extra/observability/redaction.py @@ -22,6 +22,13 @@ from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult + # Inherit from the real base only for static analysis: linters verify our + # export/shutdown/force_flush signatures. At runtime the base is object so + # the optional OpenTelemetry SDK is not required to import this module. + _SpanExporterBase = SpanExporter +else: + _SpanExporterBase = object + # User-supplied per-attribute masker: given (key, value), return the value # to keep. Return the value unchanged to keep it, a redacted value to mask it, @@ -325,10 +332,7 @@ def resolve_redaction(redaction: RedactionPolicyLike | bool) -> RedactionPolicy # SpanExporter wrapper -# NOTE: in essence this is a subclass of SpanExporter. It's not typed as such because -# the opentelemetry SDK is an optional dependency, so to keep it importable we duck-type it -# TODO: can't we rely on typing to have static linters verify that ? (i.e. don't inherit, but type) -class RedactingSpanExporter: +class RedactingSpanExporter(_SpanExporterBase): """Wrap any SpanExporter to redact spans before delegating export. Example From cd0b13b10b95a3c4a43a54edc21309c0e9a078ce Mon Sep 17 00:00:00 2001 From: Simon Van de Kerckhove Date: Thu, 2 Jul 2026 11:12:00 +0200 Subject: [PATCH 06/12] feat: redact string elements of homogeneous sequences Add a shared _redact_value helper so both policies scan the string elements of list/tuple attribute values instead of passing them through verbatim, preserving the container type and leaving numeric/bool sequences untouched. Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe --- .../extra/observability/redaction.py | 39 +++++++++++-------- src/mistralai/extra/tests/test_redaction.py | 29 ++++++++++++++ 2 files changed, 52 insertions(+), 16 deletions(-) diff --git a/src/mistralai/extra/observability/redaction.py b/src/mistralai/extra/observability/redaction.py index 2620d94f..992626eb 100644 --- a/src/mistralai/extra/observability/redaction.py +++ b/src/mistralai/extra/observability/redaction.py @@ -14,7 +14,7 @@ import re from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Callable, Final, Union +from typing import TYPE_CHECKING, Any, Callable, Final, Union, cast from opentelemetry.util.types import AttributeValue @@ -142,9 +142,6 @@ def redact_status_description(self, description: str | None) -> str | None: _PRIMITIVE_TYPES: Final[tuple[type, ...]] = (str, bool, int, float) -# TODO: also redact Sequence[str] ? - - class AttributeRedactionPolicy(RedactionPolicy): """Key-oriented hybrid policy. @@ -202,12 +199,9 @@ def redact_attributes( if self._should_redact(key, value): redacted[key] = self._redacted_value continue - if isinstance(value, str): - redacted[key] = _redact_text( - value, self._token_patterns, self._redacted_value - ) - continue - redacted[key] = value + redacted[key] = _redact_value( + value, self._token_patterns, self._redacted_value + ) return redacted @@ -256,12 +250,7 @@ def redact_attributes( return redacted for key, value in attributes.items(): - if isinstance(value, str): - redacted[key] = _redact_text( - value, self._patterns, self._redacted_value - ) - continue - redacted[key] = value + redacted[key] = _redact_value(value, self._patterns, self._redacted_value) return redacted @@ -472,6 +461,24 @@ def _redact_status(status: Any, policy: RedactionPolicy, types: Any) -> Any: return types.Status(status_code, description) +def _redact_value( + value: AttributeValue, + patterns: Sequence[re.Pattern[str]], + redacted_value: str = DEFAULT_REDACTED_VALUE, +) -> AttributeValue: + if isinstance(value, str): + return _redact_text(value, patterns, redacted_value) + if isinstance(value, (list, tuple)): + items = [ + _redact_text(item, patterns, redacted_value) + if isinstance(item, str) + else item + for item in value + ] + return cast(AttributeValue, tuple(items) if isinstance(value, tuple) else items) + return value + + def _redact_text( text: str, patterns: Sequence[re.Pattern[str]], diff --git a/src/mistralai/extra/tests/test_redaction.py b/src/mistralai/extra/tests/test_redaction.py index 57cada7d..36ba7ea2 100644 --- a/src/mistralai/extra/tests/test_redaction.py +++ b/src/mistralai/extra/tests/test_redaction.py @@ -57,6 +57,21 @@ def test_non_primitive_kept_when_disabled(self): out = policy.redact_attributes({"safeish.list": ("a", "b")}) self.assertEqual(out["safeish.list"], ("a", "b")) + def test_string_sequence_scanned_element_wise_when_kept(self): + policy = AttributeRedactionPolicy(redact_non_primitive=False) + out = policy.redact_attributes( + {"tags": ["plain", "ghp_abcdefghijklmnopqrstuvwxyz0123"]} + ) + self.assertEqual(out["tags"], ["plain", DEFAULT_REDACTED_VALUE]) + + def test_safe_key_string_sequence_scanned(self): + out = self.policy.redact_attributes( + {"gen_ai.response.finish_reasons": ("stop", "Bearer abc.def")} + ) + self.assertEqual( + out["gen_ai.response.finish_reasons"], ("stop", DEFAULT_REDACTED_VALUE) + ) + def test_none_attributes_returns_empty(self): self.assertEqual(self.policy.redact_attributes(None), {}) @@ -100,6 +115,20 @@ def test_non_string_untouched(self): out = self.policy.redact_attributes({"n": 5, "b": True}) self.assertEqual(out, {"n": 5, "b": True}) + def test_string_sequence_scanned_preserving_container(self): + out = self.policy.redact_attributes( + {"msgs": ["hello", "reach me at a@b.com"]} + ) + self.assertEqual(out["msgs"], ["hello", "reach me at [REDACTED]"]) + + def test_tuple_sequence_stays_tuple(self): + out = self.policy.redact_attributes({"msgs": ("hi", "a@b.com")}) + self.assertEqual(out["msgs"], ("hi", "[REDACTED]")) + + def test_numeric_sequence_untouched(self): + out = self.policy.redact_attributes({"nums": [1, 2, 3]}) + self.assertEqual(out["nums"], [1, 2, 3]) + def test_span_name_scanned(self): self.assertEqual(self.policy.redact_span_name("op a@b.com"), "op [REDACTED]") From 48d636982d6949baea40116fcb25ec395c915066 Mon Sep 17 00:00:00 2001 From: Simon Van de Kerckhove Date: Thu, 2 Jul 2026 11:38:36 +0200 Subject: [PATCH 07/12] test: convert redaction tests to pytest style Rewrite test_redaction.py as pytest classes with fixtures and parametrization, and convert the TestTelemetryRedaction class in test_telemetry.py to a plain pytest class (using caplog for log assertions). Optional span attributes/context are narrowed with asserts instead of file-level pyright suppressions. Older telemetry tests are left as-is. Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe --- src/mistralai/extra/tests/test_redaction.py | 282 ++++++++++---------- src/mistralai/extra/tests/test_telemetry.py | 91 ++++--- 2 files changed, 188 insertions(+), 185 deletions(-) diff --git a/src/mistralai/extra/tests/test_redaction.py b/src/mistralai/extra/tests/test_redaction.py index 36ba7ea2..480f4568 100644 --- a/src/mistralai/extra/tests/test_redaction.py +++ b/src/mistralai/extra/tests/test_redaction.py @@ -1,7 +1,6 @@ """Tests for client-side telemetry redaction.""" -import unittest - +import pytest from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -22,147 +21,146 @@ ) -class TestAttributeRedactionPolicy(unittest.TestCase): - def setUp(self): - self.policy = AttributeRedactionPolicy() +@pytest.fixture +def attribute_policy() -> AttributeRedactionPolicy: + return AttributeRedactionPolicy() + + +@pytest.fixture +def regex_policy() -> RegexRedactionPolicy: + return RegexRedactionPolicy() - def test_sensitive_key_redacted_wholesale(self): - out = self.policy.redact_attributes({"gen_ai.input.messages": "hello"}) - self.assertEqual(out["gen_ai.input.messages"], DEFAULT_REDACTED_VALUE) - def test_safe_key_kept(self): - out = self.policy.redact_attributes({"gen_ai.request.model": "mistral-large"}) - self.assertEqual(out["gen_ai.request.model"], "mistral-large") +class TestAttributeRedactionPolicy: + def test_sensitive_key_redacted_wholesale(self, attribute_policy: AttributeRedactionPolicy): + out = attribute_policy.redact_attributes({"gen_ai.input.messages": "hello"}) + assert out["gen_ai.input.messages"] == DEFAULT_REDACTED_VALUE - def test_usage_prefix_kept(self): - out = self.policy.redact_attributes({"gen_ai.usage.input_tokens": 42}) - self.assertEqual(out["gen_ai.usage.input_tokens"], 42) + def test_safe_key_kept(self, attribute_policy: AttributeRedactionPolicy): + out = attribute_policy.redact_attributes({"gen_ai.request.model": "mistral-large"}) + assert out["gen_ai.request.model"] == "mistral-large" - def test_fragment_match_redacted(self): - out = self.policy.redact_attributes({"custom.prompt.text": "secret prompt"}) - self.assertEqual(out["custom.prompt.text"], DEFAULT_REDACTED_VALUE) + def test_usage_prefix_kept(self, attribute_policy: AttributeRedactionPolicy): + out = attribute_policy.redact_attributes({"gen_ai.usage.input_tokens": 42}) + assert out["gen_ai.usage.input_tokens"] == 42 - def test_token_pattern_on_kept_string(self): - out = self.policy.redact_attributes( + def test_fragment_match_redacted(self, attribute_policy: AttributeRedactionPolicy): + out = attribute_policy.redact_attributes({"custom.prompt.text": "secret prompt"}) + assert out["custom.prompt.text"] == DEFAULT_REDACTED_VALUE + + def test_token_pattern_on_kept_string(self, attribute_policy: AttributeRedactionPolicy): + out = attribute_policy.redact_attributes( {"note": "call token ghp_abcdefghijklmnopqrstuvwxyz0123 now"} ) - self.assertEqual(out["note"], "call token [REDACTED] now") + assert out["note"] == "call token [REDACTED] now" - def test_non_primitive_redacted(self): - out = self.policy.redact_attributes({"data": ("a", "b")}) - self.assertEqual(out["data"], DEFAULT_REDACTED_VALUE) + def test_non_primitive_redacted(self, attribute_policy: AttributeRedactionPolicy): + out = attribute_policy.redact_attributes({"data": ("a", "b")}) + assert out["data"] == DEFAULT_REDACTED_VALUE def test_non_primitive_kept_when_disabled(self): policy = AttributeRedactionPolicy(redact_non_primitive=False) out = policy.redact_attributes({"safeish.list": ("a", "b")}) - self.assertEqual(out["safeish.list"], ("a", "b")) + assert out["safeish.list"] == ("a", "b") def test_string_sequence_scanned_element_wise_when_kept(self): policy = AttributeRedactionPolicy(redact_non_primitive=False) out = policy.redact_attributes( {"tags": ["plain", "ghp_abcdefghijklmnopqrstuvwxyz0123"]} ) - self.assertEqual(out["tags"], ["plain", DEFAULT_REDACTED_VALUE]) + assert out["tags"] == ["plain", DEFAULT_REDACTED_VALUE] - def test_safe_key_string_sequence_scanned(self): - out = self.policy.redact_attributes( + def test_safe_key_string_sequence_scanned(self, attribute_policy: AttributeRedactionPolicy): + out = attribute_policy.redact_attributes( {"gen_ai.response.finish_reasons": ("stop", "Bearer abc.def")} ) - self.assertEqual( - out["gen_ai.response.finish_reasons"], ("stop", DEFAULT_REDACTED_VALUE) - ) + assert out["gen_ai.response.finish_reasons"] == ("stop", DEFAULT_REDACTED_VALUE) - def test_none_attributes_returns_empty(self): - self.assertEqual(self.policy.redact_attributes(None), {}) + def test_none_attributes_returns_empty(self, attribute_policy: AttributeRedactionPolicy): + assert attribute_policy.redact_attributes(None) == {} - def test_status_description_redacted(self): - self.assertEqual( - self.policy.redact_status_description("boom: user@x.com"), - DEFAULT_REDACTED_VALUE, + def test_status_description_redacted(self, attribute_policy: AttributeRedactionPolicy): + assert ( + attribute_policy.redact_status_description("boom: user@x.com") + == DEFAULT_REDACTED_VALUE ) - self.assertIsNone(self.policy.redact_status_description(None)) + assert attribute_policy.redact_status_description(None) is None - def test_span_name_unchanged(self): - self.assertEqual(self.policy.redact_span_name("chat mistral-large"), "chat mistral-large") + def test_span_name_unchanged(self, attribute_policy: AttributeRedactionPolicy): + assert attribute_policy.redact_span_name("chat mistral-large") == "chat mistral-large" def test_custom_redacted_value(self): policy = AttributeRedactionPolicy(redacted_value="XXX") out = policy.redact_attributes({"http.url": "https://x"}) - self.assertEqual(out["http.url"], "XXX") - + assert out["http.url"] == "XXX" -class TestRegexRedactionPolicy(unittest.TestCase): - def setUp(self): - self.policy = RegexRedactionPolicy() - def test_email_redacted_inline_preserving_structure(self): - out = self.policy.redact_attributes( +class TestRegexRedactionPolicy: + def test_email_redacted_inline_preserving_structure(self, regex_policy: RegexRedactionPolicy): + out = regex_policy.redact_attributes( {"gen_ai.input.messages": '{"content":"reach me at a@b.com"}'} ) - self.assertEqual( - out["gen_ai.input.messages"], '{"content":"reach me at [REDACTED]"}' - ) + assert out["gen_ai.input.messages"] == '{"content":"reach me at [REDACTED]"}' - def test_token_redacted(self): - out = self.policy.redact_attributes({"h": "Bearer abc.def-ghi"}) - self.assertEqual(out["h"], "[REDACTED]") + def test_token_redacted(self, regex_policy: RegexRedactionPolicy): + out = regex_policy.redact_attributes({"h": "Bearer abc.def-ghi"}) + assert out["h"] == "[REDACTED]" - def test_non_matching_string_kept(self): - out = self.policy.redact_attributes({"server.address": "prod-host-1"}) - self.assertEqual(out["server.address"], "prod-host-1") + def test_non_matching_string_kept(self, regex_policy: RegexRedactionPolicy): + out = regex_policy.redact_attributes({"server.address": "prod-host-1"}) + assert out["server.address"] == "prod-host-1" - def test_non_string_untouched(self): - out = self.policy.redact_attributes({"n": 5, "b": True}) - self.assertEqual(out, {"n": 5, "b": True}) - - def test_string_sequence_scanned_preserving_container(self): - out = self.policy.redact_attributes( - {"msgs": ["hello", "reach me at a@b.com"]} - ) - self.assertEqual(out["msgs"], ["hello", "reach me at [REDACTED]"]) + def test_non_string_untouched(self, regex_policy: RegexRedactionPolicy): + out = regex_policy.redact_attributes({"n": 5, "b": True}) + assert out == {"n": 5, "b": True} - def test_tuple_sequence_stays_tuple(self): - out = self.policy.redact_attributes({"msgs": ("hi", "a@b.com")}) - self.assertEqual(out["msgs"], ("hi", "[REDACTED]")) + def test_span_name_scanned(self, regex_policy: RegexRedactionPolicy): + assert regex_policy.redact_span_name("op a@b.com") == "op [REDACTED]" - def test_numeric_sequence_untouched(self): - out = self.policy.redact_attributes({"nums": [1, 2, 3]}) - self.assertEqual(out["nums"], [1, 2, 3]) - - def test_span_name_scanned(self): - self.assertEqual(self.policy.redact_span_name("op a@b.com"), "op [REDACTED]") - - def test_status_description_scanned(self): - self.assertEqual( - self.policy.redact_status_description("failed for a@b.com"), - "failed for [REDACTED]", + def test_status_description_scanned(self, regex_policy: RegexRedactionPolicy): + assert ( + regex_policy.redact_status_description("failed for a@b.com") + == "failed for [REDACTED]" ) - def test_secret_patterns_redacted(self): - secrets = { - "aws": "AKIAIOSFODNN7EXAMPLE", - "google": "AIzaabcdefghijklmnopqrstuvwxyz012345678", - "jwt": "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.abc123", - "pem": "-----BEGIN RSA PRIVATE KEY-----", - "stripe": "sk_live_0123456789abcdefghij", - } - for name, secret in secrets.items(): - with self.subTest(secret=name): - out = self.policy.redact_attributes({"v": f"leak {secret} here"}) - self.assertNotIn(secret, out["v"]) - self.assertIn(DEFAULT_REDACTED_VALUE, out["v"]) - - -class TestDefaultPatternComposition(unittest.TestCase): + @pytest.mark.parametrize( + "secret", + [ + "AKIAIOSFODNN7EXAMPLE", + "AIzaabcdefghijklmnopqrstuvwxyz012345678", + "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.abc123", + "-----BEGIN RSA PRIVATE KEY-----", + "sk_live_0123456789abcdefghij", + ], + ) + def test_secret_patterns_redacted(self, regex_policy: RegexRedactionPolicy, secret: str): + out = regex_policy.redact_attributes({"v": f"leak {secret} here"}) + value = out["v"] + assert isinstance(value, str) + assert secret not in value + assert DEFAULT_REDACTED_VALUE in value + + def test_string_sequence_scanned_preserving_container(self, regex_policy: RegexRedactionPolicy): + out = regex_policy.redact_attributes({"msgs": ["hello", "reach me at a@b.com"]}) + assert out["msgs"] == ["hello", "reach me at [REDACTED]"] + + def test_tuple_sequence_stays_tuple(self, regex_policy: RegexRedactionPolicy): + out = regex_policy.redact_attributes({"msgs": ("hi", "a@b.com")}) + assert out["msgs"] == ("hi", "[REDACTED]") + + def test_numeric_sequence_untouched(self, regex_policy: RegexRedactionPolicy): + out = regex_policy.redact_attributes({"nums": [1, 2, 3]}) + assert out["nums"] == [1, 2, 3] + + +class TestDefaultPatternComposition: def test_pii_patterns_extend_token_patterns(self): prefix = DEFAULT_PII_SECRET_PATTERNS[: len(DEFAULT_TOKEN_PATTERNS)] - self.assertEqual(prefix, DEFAULT_TOKEN_PATTERNS) - self.assertGreater( - len(DEFAULT_PII_SECRET_PATTERNS), len(DEFAULT_TOKEN_PATTERNS) - ) + assert prefix == DEFAULT_TOKEN_PATTERNS + assert len(DEFAULT_PII_SECRET_PATTERNS) > len(DEFAULT_TOKEN_PATTERNS) -class TestCallbackRedactionPolicy(unittest.TestCase): +class TestCallbackRedactionPolicy: def test_mask_applied_per_attribute(self): policy = CallbackRedactionPolicy( lambda key, value: "[X]" if "message" in key else value @@ -170,26 +168,26 @@ def test_mask_applied_per_attribute(self): out = policy.redact_attributes( {"gen_ai.output.messages": "hi", "gen_ai.request.model": "m"} ) - self.assertEqual(out, {"gen_ai.output.messages": "[X]", "gen_ai.request.model": "m"}) + assert out == {"gen_ai.output.messages": "[X]", "gen_ai.request.model": "m"} def test_returning_none_drops_attribute(self): policy = CallbackRedactionPolicy( lambda key, value: None if key == "drop" else value ) out = policy.redact_attributes({"drop": "x", "keep": "y"}) - self.assertEqual(out, {"keep": "y"}) + assert out == {"keep": "y"} -class TestRedactionPolicyABC(unittest.TestCase): +class TestRedactionPolicyABC: def test_base_class_cannot_be_instantiated(self): - with self.assertRaises(TypeError): + with pytest.raises(TypeError): RedactionPolicy() # type: ignore[abstract] def test_subclass_without_redact_attributes_cannot_be_instantiated(self): class Incomplete(RedactionPolicy): pass - with self.assertRaises(TypeError): + with pytest.raises(TypeError): Incomplete() # type: ignore[abstract] def test_subclass_implementing_redact_attributes_instantiates(self): @@ -199,30 +197,31 @@ def redact_attributes(self, attributes): policy = Minimal() # Concrete identity defaults remain available. - self.assertEqual(policy.redact_span_name("span"), "span") - self.assertEqual(policy.redact_status_description("desc"), "desc") + assert policy.redact_span_name("span") == "span" + assert policy.redact_status_description("desc") == "desc" -class TestResolvePolicy(unittest.TestCase): +class TestResolvePolicy: def test_none_returns_default(self): - self.assertIsInstance(resolve_policy(None), AttributeRedactionPolicy) - self.assertIsInstance(default_redaction_policy(), AttributeRedactionPolicy) + assert isinstance(resolve_policy(None), AttributeRedactionPolicy) + assert isinstance(default_redaction_policy(), AttributeRedactionPolicy) def test_policy_passthrough(self): policy = RegexRedactionPolicy() - self.assertIs(resolve_policy(policy), policy) + assert resolve_policy(policy) is policy def test_callable_wrapped(self): resolved = resolve_policy(lambda k, v: v) - self.assertIsInstance(resolved, CallbackRedactionPolicy) + assert isinstance(resolved, CallbackRedactionPolicy) def test_invalid_raises_type_error(self): - with self.assertRaises(TypeError): + with pytest.raises(TypeError): resolve_policy(123) # type: ignore[arg-type] -class TestRedactSpan(unittest.TestCase): - def _make_span(self): +class TestRedactSpan: + @staticmethod + def _make_span(): exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) @@ -238,35 +237,41 @@ def _make_span(self): def test_rebuilds_genuine_readable_span(self): from opentelemetry.sdk.trace import ReadableSpan - original = self._make_span() - redacted = redact_span(original, default_redaction_policy()) - self.assertIsInstance(redacted, ReadableSpan) + redacted = redact_span(self._make_span(), default_redaction_policy()) + assert isinstance(redacted, ReadableSpan) def test_attributes_redacted(self): redacted = redact_span(self._make_span(), default_redaction_policy()) - self.assertEqual(redacted.attributes["gen_ai.input.messages"], DEFAULT_REDACTED_VALUE) - self.assertEqual(redacted.attributes["gen_ai.request.model"], "mistral-large") + attrs = redacted.attributes + assert attrs is not None + assert attrs["gen_ai.input.messages"] == DEFAULT_REDACTED_VALUE + assert attrs["gen_ai.request.model"] == "mistral-large" def test_event_attributes_redacted(self): redacted = redact_span(self._make_span(), default_redaction_policy()) event = redacted.events[0] - self.assertEqual(event.name, "exception") - self.assertEqual(event.attributes["exception.message"], DEFAULT_REDACTED_VALUE) + assert event.name == "exception" + attrs = event.attributes + assert attrs is not None + assert attrs["exception.message"] == DEFAULT_REDACTED_VALUE def test_status_description_redacted(self): redacted = redact_span(self._make_span(), default_redaction_policy()) - self.assertEqual(redacted.status.status_code, StatusCode.ERROR) - self.assertEqual(redacted.status.description, DEFAULT_REDACTED_VALUE) + assert redacted.status.status_code == StatusCode.ERROR + assert redacted.status.description == DEFAULT_REDACTED_VALUE def test_identity_preserved(self): original = self._make_span() redacted = redact_span(original, default_redaction_policy()) - self.assertEqual(redacted.context.span_id, original.context.span_id) - self.assertEqual(redacted.context.trace_id, original.context.trace_id) + assert redacted.context is not None + assert original.context is not None + assert redacted.context.span_id == original.context.span_id + assert redacted.context.trace_id == original.context.trace_id -class TestRedactingSpanExporter(unittest.TestCase): - def _export_through(self, policy=None): +class TestRedactingSpanExporter: + @staticmethod + def _export_through(policy=None): wrapped = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor( @@ -281,20 +286,23 @@ def _export_through(self, policy=None): def test_wrapped_exporter_receives_redacted_spans(self): spans = self._export_through() - self.assertEqual(len(spans), 1) + assert len(spans) == 1 attrs = spans[0].attributes - self.assertEqual(attrs["gen_ai.output.messages"], DEFAULT_REDACTED_VALUE) - self.assertEqual(attrs["gen_ai.request.model"], "mistral-large") + assert attrs is not None + assert attrs["gen_ai.output.messages"] == DEFAULT_REDACTED_VALUE + assert attrs["gen_ai.request.model"] == "mistral-large" def test_custom_policy_used(self): spans = self._export_through(RegexRedactionPolicy()) # Regex policy keeps structure; "leak" has no PII pattern -> unchanged. - self.assertEqual(spans[0].attributes["gen_ai.output.messages"], "leak") + attrs = spans[0].attributes + assert attrs is not None + assert attrs["gen_ai.output.messages"] == "leak" def test_export_returns_wrapped_result(self): wrapped = InMemorySpanExporter() exporter = RedactingSpanExporter(wrapped) - self.assertEqual(exporter.export([]), SpanExportResult.SUCCESS) + assert exporter.export([]) == SpanExportResult.SUCCESS def test_shutdown_and_force_flush_delegate(self): class _Recorder(InMemorySpanExporter): @@ -313,11 +321,7 @@ def force_flush(self, timeout_millis=30000): recorder = _Recorder() exporter = RedactingSpanExporter(recorder) - self.assertTrue(exporter.force_flush()) + assert exporter.force_flush() is True exporter.shutdown() - self.assertTrue(recorder.flush_called) - self.assertTrue(recorder.shutdown_called) - - -if __name__ == "__main__": - unittest.main() + assert recorder.flush_called + assert recorder.shutdown_called diff --git a/src/mistralai/extra/tests/test_telemetry.py b/src/mistralai/extra/tests/test_telemetry.py index f2ef7f46..46964e23 100644 --- a/src/mistralai/extra/tests/test_telemetry.py +++ b/src/mistralai/extra/tests/test_telemetry.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, cast from unittest.mock import MagicMock, patch +import pytest from opentelemetry.sdk.trace import TracerProvider from mistralai.client._hooks import SDKHooks @@ -654,11 +655,17 @@ def test_mistral_endpoint_env_overrides_default_endpoint(self): ) -class TestTelemetryRedaction(unittest.TestCase): - def setUp(self): - FakeExporter.instances.clear() +_TELEMETRY_LOGGER = "mistralai.extra.observability.telemetry" + + +@pytest.fixture +def clear_exporters(): + FakeExporter.instances.clear() - def _make_provider(self, **kwargs): + +class TestTelemetryRedaction: + @staticmethod + def _make_provider(**kwargs): with patch( "mistralai.extra.observability.telemetry._load_otel_sdk", return_value=( @@ -670,56 +677,57 @@ def _make_provider(self, **kwargs): ): return _create_telemetry_tracer_provider(api_key="test-key", **kwargs) - def _exporter_of(self, provider): - self.assertEqual(len(provider.span_processors), 1) + @staticmethod + def _exporter_of(provider): + assert len(provider.span_processors) == 1 return provider.span_processors[0].exporter - def test_dedicated_wraps_exporter_by_default(self): + def test_dedicated_wraps_exporter_by_default(self, clear_exporters): provider = self._make_provider() exporter = self._exporter_of(provider) - self.assertIsInstance(exporter, RedactingSpanExporter) + assert isinstance(exporter, RedactingSpanExporter) # The wrapped exporter is still the single OTLP exporter created. - self.assertEqual(len(FakeExporter.instances), 1) - self.assertIs(exporter._exporter, FakeExporter.instances[0]) - self.assertIsInstance(exporter._policy, AttributeRedactionPolicy) + assert len(FakeExporter.instances) == 1 + assert exporter._exporter is FakeExporter.instances[0] + assert isinstance(exporter._policy, AttributeRedactionPolicy) - def test_redaction_true_wraps_with_default_policy(self): + def test_redaction_true_wraps_with_default_policy(self, clear_exporters): provider = self._make_provider(redaction=True) exporter = self._exporter_of(provider) - self.assertIsInstance(exporter, RedactingSpanExporter) - self.assertIsInstance(exporter._policy, AttributeRedactionPolicy) + assert isinstance(exporter, RedactingSpanExporter) + assert isinstance(exporter._policy, AttributeRedactionPolicy) - def test_redaction_false_leaves_exporter_unwrapped(self): + def test_redaction_false_leaves_exporter_unwrapped(self, clear_exporters): provider = self._make_provider(redaction=False) exporter = self._exporter_of(provider) - self.assertNotIsInstance(exporter, RedactingSpanExporter) - self.assertIs(exporter, FakeExporter.instances[0]) + assert not isinstance(exporter, RedactingSpanExporter) + assert exporter is FakeExporter.instances[0] - def test_custom_policy_instance_is_used(self): + def test_custom_policy_instance_is_used(self, clear_exporters): policy = RegexRedactionPolicy() provider = self._make_provider(redaction=policy) exporter = self._exporter_of(provider) - self.assertIsInstance(exporter, RedactingSpanExporter) - self.assertIs(exporter._policy, policy) + assert isinstance(exporter, RedactingSpanExporter) + assert exporter._policy is policy - def test_callback_is_wrapped_in_callback_policy(self): + def test_callback_is_wrapped_in_callback_policy(self, clear_exporters): def mask(key, value): return value provider = self._make_provider(redaction=mask) exporter = self._exporter_of(provider) - self.assertIsInstance(exporter, RedactingSpanExporter) - self.assertIsInstance(exporter._policy, CallbackRedactionPolicy) + assert isinstance(exporter, RedactingSpanExporter) + assert isinstance(exporter._policy, CallbackRedactionPolicy) def test_resolve_redaction_semantics(self): - self.assertIsInstance(resolve_redaction(True), AttributeRedactionPolicy) - self.assertIsNone(resolve_redaction(False)) + assert isinstance(resolve_redaction(True), AttributeRedactionPolicy) + assert resolve_redaction(False) is None policy = RegexRedactionPolicy() - self.assertIs(resolve_redaction(policy), policy) + assert resolve_redaction(policy) is policy resolved = resolve_redaction(lambda k, v: v) - self.assertIsInstance(resolved, CallbackRedactionPolicy) + assert isinstance(resolved, CallbackRedactionPolicy) def test_configure_dedicated_threads_redaction(self): with patch( @@ -732,32 +740,23 @@ def test_configure_dedicated_threads_redaction(self): create_provider.assert_called_once_with(api_key="test-key", redaction=policy) - def test_global_mode_warns_when_redaction_customized(self): + def test_global_mode_warns_when_redaction_customized(self, caplog): client = _make_client(api_key="test-key") - with self.assertLogs( - "mistralai.extra.observability.telemetry", level="WARNING" - ) as logs: + with caplog.at_level("WARNING", logger=_TELEMETRY_LOGGER): configure_telemetry(client, provider="global", redaction=False) - self.assertIn("only applied in 'dedicated'", logs.output[0]) + assert "only applied in 'dedicated'" in caplog.text - def test_global_mode_does_not_warn_by_default(self): + def test_global_mode_does_not_warn_by_default(self, caplog): client = _make_client(api_key="test-key") - logger = __import__( - "logging" - ).getLogger("mistralai.extra.observability.telemetry") - with patch.object(logger, "warning") as warn: + with caplog.at_level("WARNING", logger=_TELEMETRY_LOGGER): configure_telemetry(client, provider="global") - warn.assert_not_called() + assert not [r for r in caplog.records if r.name == _TELEMETRY_LOGGER] - def test_custom_provider_warns_when_redaction_customized(self): + def test_custom_provider_warns_when_redaction_customized(self, caplog): client = _make_client(api_key="test-key") - with self.assertLogs( - "mistralai.extra.observability.telemetry", level="WARNING" - ) as logs: - configure_telemetry( - client, provider=TracerProvider(), redaction=False - ) - self.assertIn("only applied in 'dedicated'", logs.output[0]) + with caplog.at_level("WARNING", logger=_TELEMETRY_LOGGER): + configure_telemetry(client, provider=TracerProvider(), redaction=False) + assert "only applied in 'dedicated'" in caplog.text if __name__ == "__main__": From 2dfee801a85e5b3c77dd98d1183a93ee8cda620e Mon Sep 17 00:00:00 2001 From: Simon Van de Kerckhove Date: Thu, 2 Jul 2026 12:13:54 +0200 Subject: [PATCH 08/12] fix: tweaks --- .../extra/observability/telemetry.py | 11 +- src/mistralai/extra/tests/test_redaction.py | 142 +++++++----------- src/mistralai/extra/tests/test_telemetry.py | 63 ++++---- 3 files changed, 97 insertions(+), 119 deletions(-) diff --git a/src/mistralai/extra/observability/telemetry.py b/src/mistralai/extra/observability/telemetry.py index 927c5ee6..f9e6628e 100644 --- a/src/mistralai/extra/observability/telemetry.py +++ b/src/mistralai/extra/observability/telemetry.py @@ -21,9 +21,9 @@ if TYPE_CHECKING: from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider + from mistralai.client._hooks.tracing import TracingHook 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" @@ -95,13 +95,8 @@ def _warn_redaction_ignored( redaction: RedactionPolicyLike | bool, mode: str, ) -> None: - """Warn when a redaction override cannot take effect for this provider mode. - - Redaction is only applied in ``dedicated`` mode where the SDK owns the - exporter. In ``global``/``custom`` modes the application owns the export - pipeline and must wrap its own exporter with ``RedactingSpanExporter``. - """ - if redaction is True: + """Warn when redaction may not happen when user might expect it.""" + if redaction is False: # Explicitly turned off return logger.warning( "Telemetry redaction is only applied in 'dedicated' provider mode, where " diff --git a/src/mistralai/extra/tests/test_redaction.py b/src/mistralai/extra/tests/test_redaction.py index 480f4568..0b0c0eec 100644 --- a/src/mistralai/extra/tests/test_redaction.py +++ b/src/mistralai/extra/tests/test_redaction.py @@ -1,23 +1,19 @@ -"""Tests for client-side telemetry redaction.""" - import pytest -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import SpanKind, Status, StatusCode from mistralai.extra.observability.redaction import ( - DEFAULT_PII_SECRET_PATTERNS, DEFAULT_REDACTED_VALUE, - DEFAULT_TOKEN_PATTERNS, AttributeRedactionPolicy, CallbackRedactionPolicy, RedactingSpanExporter, - RedactionPolicy, RegexRedactionPolicy, default_redaction_policy, redact_span, resolve_policy, + resolve_redaction, ) @@ -32,12 +28,16 @@ def regex_policy() -> RegexRedactionPolicy: class TestAttributeRedactionPolicy: - def test_sensitive_key_redacted_wholesale(self, attribute_policy: AttributeRedactionPolicy): + def test_sensitive_key_redacted_wholesale( + self, attribute_policy: AttributeRedactionPolicy + ): out = attribute_policy.redact_attributes({"gen_ai.input.messages": "hello"}) assert out["gen_ai.input.messages"] == DEFAULT_REDACTED_VALUE def test_safe_key_kept(self, attribute_policy: AttributeRedactionPolicy): - out = attribute_policy.redact_attributes({"gen_ai.request.model": "mistral-large"}) + out = attribute_policy.redact_attributes( + {"gen_ai.request.model": "mistral-large"} + ) assert out["gen_ai.request.model"] == "mistral-large" def test_usage_prefix_kept(self, attribute_policy: AttributeRedactionPolicy): @@ -45,10 +45,14 @@ def test_usage_prefix_kept(self, attribute_policy: AttributeRedactionPolicy): assert out["gen_ai.usage.input_tokens"] == 42 def test_fragment_match_redacted(self, attribute_policy: AttributeRedactionPolicy): - out = attribute_policy.redact_attributes({"custom.prompt.text": "secret prompt"}) + out = attribute_policy.redact_attributes( + {"custom.prompt.text": "secret prompt"} + ) assert out["custom.prompt.text"] == DEFAULT_REDACTED_VALUE - def test_token_pattern_on_kept_string(self, attribute_policy: AttributeRedactionPolicy): + def test_token_pattern_on_kept_string( + self, attribute_policy: AttributeRedactionPolicy + ): out = attribute_policy.redact_attributes( {"note": "call token ghp_abcdefghijklmnopqrstuvwxyz0123 now"} ) @@ -70,16 +74,22 @@ def test_string_sequence_scanned_element_wise_when_kept(self): ) assert out["tags"] == ["plain", DEFAULT_REDACTED_VALUE] - def test_safe_key_string_sequence_scanned(self, attribute_policy: AttributeRedactionPolicy): + def test_safe_key_string_sequence_scanned( + self, attribute_policy: AttributeRedactionPolicy + ): out = attribute_policy.redact_attributes( {"gen_ai.response.finish_reasons": ("stop", "Bearer abc.def")} ) assert out["gen_ai.response.finish_reasons"] == ("stop", DEFAULT_REDACTED_VALUE) - def test_none_attributes_returns_empty(self, attribute_policy: AttributeRedactionPolicy): + def test_none_attributes_returns_empty( + self, attribute_policy: AttributeRedactionPolicy + ): assert attribute_policy.redact_attributes(None) == {} - def test_status_description_redacted(self, attribute_policy: AttributeRedactionPolicy): + def test_status_description_redacted( + self, attribute_policy: AttributeRedactionPolicy + ): assert ( attribute_policy.redact_status_description("boom: user@x.com") == DEFAULT_REDACTED_VALUE @@ -87,7 +97,10 @@ def test_status_description_redacted(self, attribute_policy: AttributeRedactionP assert attribute_policy.redact_status_description(None) is None def test_span_name_unchanged(self, attribute_policy: AttributeRedactionPolicy): - assert attribute_policy.redact_span_name("chat mistral-large") == "chat mistral-large" + assert ( + attribute_policy.redact_span_name("chat mistral-large") + == "chat mistral-large" + ) def test_custom_redacted_value(self): policy = AttributeRedactionPolicy(redacted_value="XXX") @@ -96,7 +109,9 @@ def test_custom_redacted_value(self): class TestRegexRedactionPolicy: - def test_email_redacted_inline_preserving_structure(self, regex_policy: RegexRedactionPolicy): + def test_email_redacted_inline_preserving_structure( + self, regex_policy: RegexRedactionPolicy + ): out = regex_policy.redact_attributes( {"gen_ai.input.messages": '{"content":"reach me at a@b.com"}'} ) @@ -133,14 +148,18 @@ def test_status_description_scanned(self, regex_policy: RegexRedactionPolicy): "sk_live_0123456789abcdefghij", ], ) - def test_secret_patterns_redacted(self, regex_policy: RegexRedactionPolicy, secret: str): + def test_secret_patterns_redacted( + self, regex_policy: RegexRedactionPolicy, secret: str + ): out = regex_policy.redact_attributes({"v": f"leak {secret} here"}) value = out["v"] assert isinstance(value, str) assert secret not in value assert DEFAULT_REDACTED_VALUE in value - def test_string_sequence_scanned_preserving_container(self, regex_policy: RegexRedactionPolicy): + def test_string_sequence_scanned_preserving_container( + self, regex_policy: RegexRedactionPolicy + ): out = regex_policy.redact_attributes({"msgs": ["hello", "reach me at a@b.com"]}) assert out["msgs"] == ["hello", "reach me at [REDACTED]"] @@ -153,13 +172,6 @@ def test_numeric_sequence_untouched(self, regex_policy: RegexRedactionPolicy): assert out["nums"] == [1, 2, 3] -class TestDefaultPatternComposition: - def test_pii_patterns_extend_token_patterns(self): - prefix = DEFAULT_PII_SECRET_PATTERNS[: len(DEFAULT_TOKEN_PATTERNS)] - assert prefix == DEFAULT_TOKEN_PATTERNS - assert len(DEFAULT_PII_SECRET_PATTERNS) > len(DEFAULT_TOKEN_PATTERNS) - - class TestCallbackRedactionPolicy: def test_mask_applied_per_attribute(self): policy = CallbackRedactionPolicy( @@ -178,33 +190,10 @@ def test_returning_none_drops_attribute(self): assert out == {"keep": "y"} -class TestRedactionPolicyABC: - def test_base_class_cannot_be_instantiated(self): - with pytest.raises(TypeError): - RedactionPolicy() # type: ignore[abstract] - - def test_subclass_without_redact_attributes_cannot_be_instantiated(self): - class Incomplete(RedactionPolicy): - pass - - with pytest.raises(TypeError): - Incomplete() # type: ignore[abstract] - - def test_subclass_implementing_redact_attributes_instantiates(self): - class Minimal(RedactionPolicy): - def redact_attributes(self, attributes): - return dict(attributes or {}) - - policy = Minimal() - # Concrete identity defaults remain available. - assert policy.redact_span_name("span") == "span" - assert policy.redact_status_description("desc") == "desc" - - class TestResolvePolicy: def test_none_returns_default(self): - assert isinstance(resolve_policy(None), AttributeRedactionPolicy) assert isinstance(default_redaction_policy(), AttributeRedactionPolicy) + assert isinstance(resolve_policy(None), AttributeRedactionPolicy) def test_policy_passthrough(self): policy = RegexRedactionPolicy() @@ -219,6 +208,22 @@ def test_invalid_raises_type_error(self): resolve_policy(123) # type: ignore[arg-type] +class TestResolveRedaction: + def test_true_returns_default_policy(self): + assert isinstance(resolve_redaction(True), AttributeRedactionPolicy) + + def test_false_returns_none(self): + assert resolve_redaction(False) is None + + def test_policy_passthrough(self): + policy = RegexRedactionPolicy() + assert resolve_redaction(policy) is policy + + def test_callable_wrapped(self): + resolved = resolve_redaction(lambda k, v: v) + assert isinstance(resolved, CallbackRedactionPolicy) + + class TestRedactSpan: @staticmethod def _make_span(): @@ -234,14 +239,9 @@ def _make_span(): provider.force_flush() return exporter.get_finished_spans()[0] - def test_rebuilds_genuine_readable_span(self): - from opentelemetry.sdk.trace import ReadableSpan - - redacted = redact_span(self._make_span(), default_redaction_policy()) - assert isinstance(redacted, ReadableSpan) - def test_attributes_redacted(self): redacted = redact_span(self._make_span(), default_redaction_policy()) + assert isinstance(redacted, ReadableSpan) attrs = redacted.attributes assert attrs is not None assert attrs["gen_ai.input.messages"] == DEFAULT_REDACTED_VALUE @@ -279,7 +279,7 @@ def _export_through(policy=None): ) tracer = provider.get_tracer("test") with tracer.start_as_current_span("chat") as span: - span.set_attribute("gen_ai.output.messages", "leak") + span.set_attribute("gen_ai.output.messages", "leak Bearer abc.def-ghi") span.set_attribute("gen_ai.request.model", "mistral-large") provider.force_flush() return wrapped.get_finished_spans() @@ -294,34 +294,6 @@ def test_wrapped_exporter_receives_redacted_spans(self): def test_custom_policy_used(self): spans = self._export_through(RegexRedactionPolicy()) - # Regex policy keeps structure; "leak" has no PII pattern -> unchanged. attrs = spans[0].attributes assert attrs is not None - assert attrs["gen_ai.output.messages"] == "leak" - - def test_export_returns_wrapped_result(self): - wrapped = InMemorySpanExporter() - exporter = RedactingSpanExporter(wrapped) - assert exporter.export([]) == SpanExportResult.SUCCESS - - def test_shutdown_and_force_flush_delegate(self): - class _Recorder(InMemorySpanExporter): - def __init__(self): - super().__init__() - self.shutdown_called = False - self.flush_called = False - - def shutdown(self): - self.shutdown_called = True - super().shutdown() - - def force_flush(self, timeout_millis=30000): - self.flush_called = True - return True - - recorder = _Recorder() - exporter = RedactingSpanExporter(recorder) - assert exporter.force_flush() is True - exporter.shutdown() - assert recorder.flush_called - assert recorder.shutdown_called + assert attrs["gen_ai.output.messages"] == f"leak {DEFAULT_REDACTED_VALUE}" diff --git a/src/mistralai/extra/tests/test_telemetry.py b/src/mistralai/extra/tests/test_telemetry.py index 46964e23..32c65608 100644 --- a/src/mistralai/extra/tests/test_telemetry.py +++ b/src/mistralai/extra/tests/test_telemetry.py @@ -22,12 +22,11 @@ CallbackRedactionPolicy, RedactingSpanExporter, RegexRedactionPolicy, - resolve_redaction, ) from mistralai.extra.observability.telemetry import ( - MISTRAL_TELEMETRY_ENDPOINT, - MISTRAL_SDK_TELEMETRY_ENV, MISTRAL_OTLP_TRACES_ENDPOINT_ENV, + MISTRAL_SDK_TELEMETRY_ENV, + MISTRAL_TELEMETRY_ENDPOINT, TelemetryConfigurationError, _create_telemetry_tracer_provider, configure_telemetry_for_hook, @@ -57,7 +56,9 @@ def _make_client(api_key: str | None = "test-key") -> "Mistral": 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)] + tracing_hooks = [ + h for h in hooks.before_request_hooks if isinstance(h, TracingHook) + ] assert len(tracing_hooks) == 1 return tracing_hooks[0] @@ -494,7 +495,9 @@ def test_sdk_config_global_uses_global_provider_mode(self): with patch( "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider" ) as create_provider: - configured = configure_telemetry_for_hook(hook, client.sdk_configuration) + configured = configure_telemetry_for_hook( + hook, client.sdk_configuration + ) self.assertTrue(configured) create_provider.assert_not_called() @@ -544,7 +547,9 @@ def test_configure_telemetry_for_hook_reads_sdk_config_telemetry_flag(self): "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider", return_value=provider, ): - configured = configure_telemetry_for_hook(hook, client.sdk_configuration) + configured = configure_telemetry_for_hook( + hook, client.sdk_configuration + ) self.assertTrue(configured) self.assertIs(hook.tracer_provider, provider) @@ -686,8 +691,6 @@ def test_dedicated_wraps_exporter_by_default(self, clear_exporters): provider = self._make_provider() exporter = self._exporter_of(provider) assert isinstance(exporter, RedactingSpanExporter) - # The wrapped exporter is still the single OTLP exporter created. - assert len(FakeExporter.instances) == 1 assert exporter._exporter is FakeExporter.instances[0] assert isinstance(exporter._policy, AttributeRedactionPolicy) @@ -719,17 +722,8 @@ def mask(key, value): assert isinstance(exporter, RedactingSpanExporter) assert isinstance(exporter._policy, CallbackRedactionPolicy) - def test_resolve_redaction_semantics(self): - assert isinstance(resolve_redaction(True), AttributeRedactionPolicy) - assert resolve_redaction(False) is None - - policy = RegexRedactionPolicy() - assert resolve_redaction(policy) is policy - - resolved = resolve_redaction(lambda k, v: v) - assert isinstance(resolved, CallbackRedactionPolicy) - - def test_configure_dedicated_threads_redaction(self): + def test_dedicated_mode_forwards_custom_redaction_to_provider(self): + # Wiring test with patch( "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider" ) as create_provider: @@ -740,23 +734,40 @@ def test_configure_dedicated_threads_redaction(self): create_provider.assert_called_once_with(api_key="test-key", redaction=policy) - def test_global_mode_warns_when_redaction_customized(self, caplog): + def test_dedicated_mode_forwards_default_redaction_to_provider(self): + # Wiring test + with patch( + "mistralai.extra.observability.telemetry._create_telemetry_tracer_provider" + ) as create_provider: + create_provider.return_value = FakeProvider() + client = _make_client(api_key="test-key") + configure_telemetry(client, provider="dedicated") + + create_provider.assert_called_once_with(api_key="test-key", redaction=True) + + def test_global_mode_warns_by_default(self, caplog): client = _make_client(api_key="test-key") with caplog.at_level("WARNING", logger=_TELEMETRY_LOGGER): - configure_telemetry(client, provider="global", redaction=False) - assert "only applied in 'dedicated'" in caplog.text + configure_telemetry(client, provider="global") + assert "Telemetry redaction is only applied in 'dedicated'" in caplog.text - def test_global_mode_does_not_warn_by_default(self, caplog): + def test_global_mode_does_not_warn_when_redaction_disabled(self, caplog): client = _make_client(api_key="test-key") with caplog.at_level("WARNING", logger=_TELEMETRY_LOGGER): - configure_telemetry(client, provider="global") + configure_telemetry(client, provider="global", redaction=False) assert not [r for r in caplog.records if r.name == _TELEMETRY_LOGGER] - def test_custom_provider_warns_when_redaction_customized(self, caplog): + def test_custom_provider_warns_by_default(self, caplog): + client = _make_client(api_key="test-key") + with caplog.at_level("WARNING", logger=_TELEMETRY_LOGGER): + configure_telemetry(client, provider=TracerProvider()) + assert "Telemetry redaction is only applied in 'dedicated'" in caplog.text + + def test_custom_provider_does_not_warn_when_redaction_disabled(self, caplog): client = _make_client(api_key="test-key") with caplog.at_level("WARNING", logger=_TELEMETRY_LOGGER): configure_telemetry(client, provider=TracerProvider(), redaction=False) - assert "only applied in 'dedicated'" in caplog.text + assert not [r for r in caplog.records if r.name == _TELEMETRY_LOGGER] if __name__ == "__main__": From 5adab3b89928b93e0905c8f41a0cc46525d1490f Mon Sep 17 00:00:00 2001 From: Simon Van de Kerckhove Date: Thu, 2 Jul 2026 16:08:17 +0200 Subject: [PATCH 09/12] fix: additional sensitive keys --- src/mistralai/extra/observability/redaction.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/mistralai/extra/observability/redaction.py b/src/mistralai/extra/observability/redaction.py index 992626eb..61c5333d 100644 --- a/src/mistralai/extra/observability/redaction.py +++ b/src/mistralai/extra/observability/redaction.py @@ -67,6 +67,8 @@ def redact_status_description(self, description: str | None) -> str | None: "gen_ai.input.messages", "gen_ai.output.messages", "gen_ai.tool.definitions", + "gen_ai.tool.call.arguments", + "gen_ai.tool.call.result", "http.request.body", "http.request.header.authorization", "http.request.header.cookie", From 5588e0103afdfe593dfad4bbe44a4deaf4f819ce Mon Sep 17 00:00:00 2001 From: Simon Van de Kerckhove Date: Thu, 2 Jul 2026 16:29:41 +0200 Subject: [PATCH 10/12] docs: add observability docs --- README.md | 115 +++++++++++++++++- .../observability/dedicated_telemetry.py | 31 +++++ .../global_provider_with_redaction.py | 44 +++++++ .../observability/redaction_policies.py | 42 +++++++ 4 files changed, 229 insertions(+), 3 deletions(-) create mode 100644 examples/mistral/observability/dedicated_telemetry.py create mode 100644 examples/mistral/observability/global_provider_with_redaction.py create mode 100644 examples/mistral/observability/redaction_policies.py diff --git a/README.md b/README.md index 03dd017b..011daea5 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,7 @@ Mistral AI API: Our Chat Completion and Embeddings APIs specification. Create yo * [Resource Management](#resource-management) * [Debugging](#debugging) * [IDE Support](#ide-support) + * [Telemetry \& Observability](#telemetry--observability) * [Development](#development) * [Contributions](#contributions) @@ -841,7 +842,7 @@ print(res.choices[0].message.content) operations. These operations will expose the stream as [Generator][generator] that can be consumed using a simple `for` loop. The loop will terminate when the server no longer has any events to send and closes the -underlying connection. +underlying connection. The stream is also a [Context Manager][context-manager] and can be used with the `with` statement and will close the underlying connection when the context is exited. @@ -1282,9 +1283,117 @@ Generally, the SDK will work well with most IDEs out of the box. However, when u + +## Telemetry & Observability + +The SDK can emit [OpenTelemetry](https://opentelemetry.io/) traces for the API calls it makes (chat, agents, embeddings, OCR, …), following the +[GenAI semantic conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/). +Spans capture the operation, model, token usage, and — unless redacted — the input/output messages and tool calls. Telemetry is **opt-in** and lives in the `mistralai.extra.observability` module. + +### Installation + +Install the `telemetry` extra: + +```bash +pip install "mistralai[telemetry]" +# or: uv add "mistralai[telemetry]" +``` + +### Enabling telemetry + +Either set an environment variable before creating the client: + +```bash +export MISTRAL_SDK_TELEMETRY=dedicated # dedicated | global | false +``` + +or configure it in code: + +```python +import os +from mistralai.client import Mistral +from mistralai.extra.observability import configure_telemetry + +with Mistral(api_key=os.environ["MISTRAL_API_KEY"]) as client: + # Dedicated mode (default): the SDK creates and owns an OTLP exporter that + # ships spans to the Mistral telemetry endpoint. Spans are redacted before + # export. + configure_telemetry(client) + + client.chat.complete( + model="mistral-small-latest", + messages=[{"role": "user", "content": "Hello!"}], + ) +``` + +### Provider modes + +`configure_telemetry(client, provider=...)` selects where spans go and who owns the export pipeline: + +| `provider` | Who owns the exporter | Where spans go | Redaction | +| ---------- | --------------------- | -------------- | --------- | +| `"dedicated"` (default) | The SDK | Mistral telemetry endpoint | Applied automatically | +| `"global"` | Your application | Your global OpenTelemetry provider | **Not** applied — you need to wrap your own exporter | +| a `TracerProvider` | Your application | The provider you pass | **Not** applied — you need to wrap your own exporter | + +In `global`/custom modes your application owns the pipeline, so the `redaction` argument is ignored (a warning is logged). Wrap your own exporter with `RedactingSpanExporter` to redact spans there: + +```python +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter +from mistralai.extra.observability import RedactingSpanExporter, configure_telemetry + +provider = TracerProvider() +provider.add_span_processor( + BatchSpanProcessor(RedactingSpanExporter(OTLPSpanExporter())) +) +trace.set_tracer_provider(provider) + +# SDK spans now flow through your global provider (already redacted above). +configure_telemetry(client, provider="global") +``` + +### Redaction + +In dedicated mode, redaction is on by default. Control it with the `redaction` argument, which also accepts any of the reusable policies from `mistralai.extra.observability`: + +```python +from mistralai.extra.observability import RegexRedactionPolicy + +configure_telemetry(client) # default policy +configure_telemetry(client, redaction=RegexRedactionPolicy()) # regex policy +configure_telemetry(client, redaction=False) # disabled - no redaction +configure_telemetry( # custom callback to control how attributes are redacted + client, + redaction=lambda key, value: None if "email" in key else value, +) +``` + +| Policy | Strategy | Trade-off | +| ------ | -------- | --------- | +| `AttributeRedactionPolicy` (default, `redaction=True`) | Key-oriented: redacts whole values for sensitive keys (explicit set, fragment match, or non-primitive value), then scans kept values for secret token patterns. | High recall, "safe by default"; erases most prompt/response content. | +| `RegexRedactionPolicy` | Content-oriented: keeps keys and structure, redacts matched substrings (secret tokens plus PII — emails, card-like sequences, IPv4). | Fewer false positives, preserves observability value; may miss free-form PII not in the pattern set. | +| `CallbackRedactionPolicy` (`redaction=`) | Your `(key, value) -> value \| None` masker per attribute; return `None` to drop the attribute. | Full control; you own the logic. | + +*Note: the `RedactingSpanExporter` primitive is reusable by any OpenTelemetry application, independent of the Mistral client.* + +### Environment variables + +| Variable | Description | Default | +| -------- | ----------- | ------- | +| `MISTRAL_SDK_TELEMETRY` | Auto-enable telemetry: `dedicated`, `global`, or `false`. | unset (disabled) | +| `MISTRAL_OTLP_TRACES_ENDPOINT` | Override the OTLP traces endpoint used in dedicated mode. | `https://api.mistral.ai/telemetry/v1/traces` | +| `MISTRAL_SDK_DEBUG_TRACING` | Set to `true` for verbose tracing logs. | `false` | +| `MISTRAL_API_KEY` | Used as the bearer token for the dedicated-mode exporter. | — | + +Runnable examples live in [`examples/mistral/observability`](/examples/mistral/observability). + + # Development ## Contributions -While we value open-source contributions to this SDK, this library is generated programmatically. Any manual changes added to internal files will be overwritten on the next generation. -We look forward to hearing your feedback. Feel free to open a PR or an issue with a proof of concept and we'll do our best to include it in a future release. +While we value open-source contributions to this SDK, this library is generated programmatically. Any manual changes added to internal files will be overwritten on the next generation. +We look forward to hearing your feedback. Feel free to open a PR or an issue with a proof of concept and we'll do our best to include it in a future release. diff --git a/examples/mistral/observability/dedicated_telemetry.py b/examples/mistral/observability/dedicated_telemetry.py new file mode 100644 index 00000000..fe9107a7 --- /dev/null +++ b/examples/mistral/observability/dedicated_telemetry.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python +"""Dedicated telemetry mode. + +The SDK creates and owns an OTLP exporter that ships spans to the Mistral +telemetry endpoint. Spans are redacted before export. + +Requires the telemetry extra: pip install "mistralai[telemetry]" +""" + +import os + +from mistralai.client import Mistral +from mistralai.extra.observability import configure_telemetry + + +def main() -> None: + api_key = os.environ["MISTRAL_API_KEY"] + + with Mistral(api_key=api_key) as client: + # Dedicated mode is the default; redaction is on by default. + configure_telemetry(client) + + response = client.chat.complete( + model="mistral-small-latest", + messages=[{"role": "user", "content": "What is the best French cheese?"}], + ) + print(response.choices[0].message.content) + + +if __name__ == "__main__": + main() diff --git a/examples/mistral/observability/global_provider_with_redaction.py b/examples/mistral/observability/global_provider_with_redaction.py new file mode 100644 index 00000000..c1fad207 --- /dev/null +++ b/examples/mistral/observability/global_provider_with_redaction.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python +"""Global provider mode with application-owned redaction. + +In global (or custom TracerProvider) mode your application owns the OTEL export +pipeline, so `configure_telemetry`'s redaction argument is ignored. To redact +spans you wrap your own exporter with RedactingSpanExporter. + +Requires the telemetry extra: pip install "mistralai[telemetry]" +""" + +import os + +from opentelemetry import trace +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor + +from mistralai.client import Mistral +from mistralai.extra.observability import RedactingSpanExporter, configure_telemetry + + +def main() -> None: + api_key = os.environ["MISTRAL_API_KEY"] + + # Build your own provider and wrap the exporter with redaction. + provider = TracerProvider() + provider.add_span_processor( + BatchSpanProcessor(RedactingSpanExporter(OTLPSpanExporter())) + ) + trace.set_tracer_provider(provider) + + with Mistral(api_key=api_key) as client: + # SDK spans flow through the global provider configured above. + configure_telemetry(client, provider="global") + + response = client.chat.complete( + model="mistral-small-latest", + messages=[{"role": "user", "content": "Say hello."}], + ) + print(response.choices[0].message.content) + + +if __name__ == "__main__": + main() diff --git a/examples/mistral/observability/redaction_policies.py b/examples/mistral/observability/redaction_policies.py new file mode 100644 index 00000000..8da0d8a9 --- /dev/null +++ b/examples/mistral/observability/redaction_policies.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python +"""Choosing a redaction policy in dedicated telemetry mode. + +The `redaction` argument accepts: + - True (default): the attribute (key-oriented) policy + - False: redaction disabled + - a RedactionPolicy instance (e.g. RegexRedactionPolicy) + - a (key, value) -> value | None callback + +Requires the telemetry extra: pip install "mistralai[telemetry]" +""" + +import os + +from mistralai.client import Mistral +from mistralai.extra.observability import RegexRedactionPolicy, configure_telemetry + + +def main() -> None: + api_key = os.environ["MISTRAL_API_KEY"] + + with Mistral(api_key=api_key) as client: + # Content-oriented policy: keeps keys/structure, redacts matched + # substrings (secret tokens plus PII such as emails, cards, IPv4). + configure_telemetry(client, redaction=RegexRedactionPolicy()) + + # Alternatives: + # configure_telemetry(client, redaction=False) # disable entirely + # configure_telemetry( # custom callback + # client, + # redaction=lambda key, value: None if "email" in key else value, + # ) + + response = client.chat.complete( + model="mistral-small-latest", + messages=[{"role": "user", "content": "Say hello."}], + ) + print(response.choices[0].message.content) + + +if __name__ == "__main__": + main() From e8f0122ef51b68a136beb6910f53d1b86ec7bf2d Mon Sep 17 00:00:00 2001 From: Simon Van de Kerckhove Date: Thu, 2 Jul 2026 17:02:14 +0200 Subject: [PATCH 11/12] feat: more secret patterns --- .../extra/observability/redaction.py | 24 +++++++++++++++++++ src/mistralai/extra/tests/test_redaction.py | 11 +++++++++ 2 files changed, 35 insertions(+) diff --git a/src/mistralai/extra/observability/redaction.py b/src/mistralai/extra/observability/redaction.py index 61c5333d..b4906b28 100644 --- a/src/mistralai/extra/observability/redaction.py +++ b/src/mistralai/extra/observability/redaction.py @@ -139,6 +139,30 @@ def redact_status_description(self, description: str | None) -> str | None: re.compile(r"\beyJ[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+\b"), re.compile(r"-----BEGIN [A-Z ]*PRIVATE KEY-----"), re.compile(r"\b[sr]k_(?:live|test)_[0-9A-Za-z]{10,}\b"), + # AI providers + re.compile(r"\bsk-ant-[A-Za-z0-9\-_]{20,}\b"), + re.compile(r"\bsk-proj-[A-Za-z0-9\-_]{20,}\b"), + re.compile(r"\bhf_[A-Za-z0-9]{30,}\b"), + # Dev / infra tokens + re.compile(r"\bgithub_pat_[A-Za-z0-9_]{22,}\b"), + re.compile(r"\bglpat-[A-Za-z0-9\-=_]{20,22}\b"), + re.compile(r"\bshp(?:at|ca|pa|ss)_[a-fA-F0-9]{32}\b"), + re.compile(r"\bsq0(?:atp|csp|idp)-[0-9A-Za-z\-_]{22,43}\b"), + re.compile(r"\bPMAK-[a-zA-Z0-9]{24,59}\b"), + re.compile(r"\bphc_[a-zA-Z0-9_]{43}\b"), + re.compile(r"\brubygems_[a-f0-9]{48}\b"), + re.compile(r"\blin_api_[0-9A-Za-z]{40}\b"), + re.compile(r"pypi-AgEIcHlwaS5vcmc[A-Za-z0-9\-_]{50,}"), + re.compile(r"\bsecret_[A-Za-z0-9]{43}\b"), + re.compile(r"[A-Za-z0-9]{14}\.atlasv1\.[A-Za-z0-9]{60,}"), + re.compile(r"\bSG\.[A-Za-z0-9_\-]{22}\.[A-Za-z0-9_\-]{43}\b"), + re.compile(r"\bpk_(?:live|test)_[0-9a-zA-Z]{24}\b"), + # Webhook URLs (the whole URL is the secret) + re.compile(r"https://hooks\.slack\.com/services/[A-Za-z0-9/+]{40,}"), + re.compile( + r"https://discord(?:app)?\.com/api/webhooks/[0-9]{17,}/[A-Za-z0-9\-_]{60,}" + ), + re.compile(r"https://hooks\.zapier\.com/hooks/catch/[A-Za-z0-9/]{16,}"), ) _SAFE_KEY_PREFIXES: Final[tuple[str, ...]] = ("gen_ai.usage.",) _PRIMITIVE_TYPES: Final[tuple[type, ...]] = (str, bool, int, float) diff --git a/src/mistralai/extra/tests/test_redaction.py b/src/mistralai/extra/tests/test_redaction.py index 0b0c0eec..865ffb8f 100644 --- a/src/mistralai/extra/tests/test_redaction.py +++ b/src/mistralai/extra/tests/test_redaction.py @@ -146,6 +146,17 @@ def test_status_description_scanned(self, regex_policy: RegexRedactionPolicy): "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.abc123", "-----BEGIN RSA PRIVATE KEY-----", "sk_live_0123456789abcdefghij", + "sk-ant-api03-abcdefghijklmnopqrstuvwxyz012345", + "hf_abcdefghijklmnopqrstuvwxyz0123456789", + "github_pat_abcdefghijklmnopqrstuvwxyz", + "glpat-abcdefghij0123456789ab", + "shpat_0123456789abcdef0123456789abcdef", + "sq0atp-0123456789abcdefghijkl", + "PMAK-0123456789abcdefghijklmn", + "phc_abcdefghijklmnopqrstuvwxyz0123456789abcdefg", + "SG.abcdefghijklmnopqrstuv.abcdefghijklmnopqrstuvwxyz0123456789abcdefg", + "pk_live_0123456789abcdefghijklmn", + "https://hooks.slack.com/services/T00000000/B00000000/abcdefghijklmnopqrstuvwx", ], ) def test_secret_patterns_redacted( From 3ca7da86b535d4e1d8c8809a60cddbffeb9fa57c Mon Sep 17 00:00:00 2001 From: Simon Van de Kerckhove Date: Thu, 2 Jul 2026 18:27:57 +0200 Subject: [PATCH 12/12] feat: make RegexRedactionPolicy the default redaction policy Switch default_redaction_policy() from the key-oriented AttributeRedactionPolicy to the content-oriented RegexRedactionPolicy, which preserves keys/structure and redacts only matched secret/PII substrings. Update docstrings, README policy table, the example, and adapt the behavioural tests accordingly. Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe --- README.md | 14 ++++++------- .../observability/redaction_policies.py | 10 ++++------ .../extra/observability/redaction.py | 17 ++++++++-------- .../extra/observability/telemetry.py | 7 ++++--- src/mistralai/extra/tests/test_redaction.py | 20 +++++++++---------- src/mistralai/extra/tests/test_telemetry.py | 6 +++--- 6 files changed, 37 insertions(+), 37 deletions(-) diff --git a/README.md b/README.md index 011daea5..929e2807 100644 --- a/README.md +++ b/README.md @@ -1360,12 +1360,12 @@ configure_telemetry(client, provider="global") In dedicated mode, redaction is on by default. Control it with the `redaction` argument, which also accepts any of the reusable policies from `mistralai.extra.observability`: ```python -from mistralai.extra.observability import RegexRedactionPolicy +from mistralai.extra.observability import AttributeRedactionPolicy -configure_telemetry(client) # default policy -configure_telemetry(client, redaction=RegexRedactionPolicy()) # regex policy -configure_telemetry(client, redaction=False) # disabled - no redaction -configure_telemetry( # custom callback to control how attributes are redacted +configure_telemetry(client) # default policy (regex) +configure_telemetry(client, redaction=AttributeRedactionPolicy()) # very conservative key-oriented policy +configure_telemetry(client, redaction=False) # disabled - no redaction +configure_telemetry( # custom callback to control how attributes are redacted client, redaction=lambda key, value: None if "email" in key else value, ) @@ -1373,8 +1373,8 @@ configure_telemetry( # custom callback | Policy | Strategy | Trade-off | | ------ | -------- | --------- | -| `AttributeRedactionPolicy` (default, `redaction=True`) | Key-oriented: redacts whole values for sensitive keys (explicit set, fragment match, or non-primitive value), then scans kept values for secret token patterns. | High recall, "safe by default"; erases most prompt/response content. | -| `RegexRedactionPolicy` | Content-oriented: keeps keys and structure, redacts matched substrings (secret tokens plus PII — emails, card-like sequences, IPv4). | Fewer false positives, preserves observability value; may miss free-form PII not in the pattern set. | +| `RegexRedactionPolicy` (default, `redaction=True`) | Content-oriented: keeps keys and structure, redacts matched substrings (secret tokens plus PII — emails, card-like sequences, IPv4). | Redacts most sensitive data while preserving observability value; may miss free-form PII or secrets not in the pattern set. | +| `AttributeRedactionPolicy` | Key-oriented: redacts whole values for sensitive keys (explicit set, fragment match, or non-primitive value), then scans kept values for secret token patterns. | Very conservative, but erases most prompt/response content. | | `CallbackRedactionPolicy` (`redaction=`) | Your `(key, value) -> value \| None` masker per attribute; return `None` to drop the attribute. | Full control; you own the logic. | *Note: the `RedactingSpanExporter` primitive is reusable by any OpenTelemetry application, independent of the Mistral client.* diff --git a/examples/mistral/observability/redaction_policies.py b/examples/mistral/observability/redaction_policies.py index 8da0d8a9..7f6fb024 100644 --- a/examples/mistral/observability/redaction_policies.py +++ b/examples/mistral/observability/redaction_policies.py @@ -2,9 +2,9 @@ """Choosing a redaction policy in dedicated telemetry mode. The `redaction` argument accepts: - - True (default): the attribute (key-oriented) policy + - True (default): the regex (content-oriented) policy - False: redaction disabled - - a RedactionPolicy instance (e.g. RegexRedactionPolicy) + - a RedactionPolicy instance (e.g. AttributeRedactionPolicy) - a (key, value) -> value | None callback Requires the telemetry extra: pip install "mistralai[telemetry]" @@ -13,16 +13,14 @@ import os from mistralai.client import Mistral -from mistralai.extra.observability import RegexRedactionPolicy, configure_telemetry +from mistralai.extra.observability import AttributeRedactionPolicy, configure_telemetry def main() -> None: api_key = os.environ["MISTRAL_API_KEY"] with Mistral(api_key=api_key) as client: - # Content-oriented policy: keeps keys/structure, redacts matched - # substrings (secret tokens plus PII such as emails, cards, IPv4). - configure_telemetry(client, redaction=RegexRedactionPolicy()) + configure_telemetry(client, redaction=AttributeRedactionPolicy()) # Alternatives: # configure_telemetry(client, redaction=False) # disable entirely diff --git a/src/mistralai/extra/observability/redaction.py b/src/mistralai/extra/observability/redaction.py index b4906b28..0b2a7a19 100644 --- a/src/mistralai/extra/observability/redaction.py +++ b/src/mistralai/extra/observability/redaction.py @@ -171,10 +171,10 @@ def redact_status_description(self, description: str | None) -> str | None: class AttributeRedactionPolicy(RedactionPolicy): """Key-oriented hybrid policy. - This is the default policy: high recall, "safe by default", at the cost of erasing most - prompt/response content. It redacts whole values for keys judged sensitive (explicit set, - fragment match, or non-primitive value), then runs token_patterns over the values it keeps - to redact values. + An opt-in, high-recall alternative to the default policy: "safe by default", at the cost + of erasing most prompt/response content. It redacts whole values for keys judged sensitive + (explicit set, fragment match, or non-primitive value), then runs token_patterns over the + values it keeps to redact values. """ def __init__( @@ -254,9 +254,10 @@ def redact_status_description(self, description: str | None) -> str | None: class RegexRedactionPolicy(RedactionPolicy): """Content-oriented policy based on regexes. - Leaves keys and structure intact, scans string values and redacts matched substrings. - Fewer false positives than default policy and aims to preserve observability value; - may miss free-form PII or secrets not in the default patterns. + This is the default policy. Leaves keys and structure intact, scans string values and + redacts matched substrings. Fewer false positives than AttributeRedactionPolicy and aims + to preserve observability value; may miss free-form PII or secrets not in the default + patterns. """ def __init__( @@ -317,7 +318,7 @@ def redact_attributes( # Helpers def default_redaction_policy() -> RedactionPolicy: - return AttributeRedactionPolicy() + return RegexRedactionPolicy() def resolve_policy(policy: RedactionPolicyLike | None) -> RedactionPolicy: diff --git a/src/mistralai/extra/observability/telemetry.py b/src/mistralai/extra/observability/telemetry.py index f9e6628e..b6c77387 100644 --- a/src/mistralai/extra/observability/telemetry.py +++ b/src/mistralai/extra/observability/telemetry.py @@ -121,10 +121,11 @@ def configure_telemetry( In dedicated mode, spans are redacted before export (safe by default). You can control this with the redaction argument: - - True: (default) uses the default policy. It redacts all possibly sensitive attributes + - True: (default) uses the default policy. It scans string values and redacts matched + secrets/PII substrings while preserving keys and surrounding content - False: disables redaction - - Some RedactionPolicy classes (e.g. based on regexes) can be found in the redaction - module and provided here + - Other RedactionPolicy classes (e.g. the conservative but destructive + AttributeRedactionPolicy) can be found in the redaction module and provided here - You can also provide a (key, value) -> value | None callback to customize how attributes get redacted. Your function should return the modified attribute value or None to drop the attribute. diff --git a/src/mistralai/extra/tests/test_redaction.py b/src/mistralai/extra/tests/test_redaction.py index 865ffb8f..16e97e65 100644 --- a/src/mistralai/extra/tests/test_redaction.py +++ b/src/mistralai/extra/tests/test_redaction.py @@ -203,8 +203,8 @@ def test_returning_none_drops_attribute(self): class TestResolvePolicy: def test_none_returns_default(self): - assert isinstance(default_redaction_policy(), AttributeRedactionPolicy) - assert isinstance(resolve_policy(None), AttributeRedactionPolicy) + assert isinstance(default_redaction_policy(), RegexRedactionPolicy) + assert isinstance(resolve_policy(None), RegexRedactionPolicy) def test_policy_passthrough(self): policy = RegexRedactionPolicy() @@ -221,7 +221,7 @@ def test_invalid_raises_type_error(self): class TestResolveRedaction: def test_true_returns_default_policy(self): - assert isinstance(resolve_redaction(True), AttributeRedactionPolicy) + assert isinstance(resolve_redaction(True), RegexRedactionPolicy) def test_false_returns_none(self): assert resolve_redaction(False) is None @@ -251,7 +251,7 @@ def _make_span(): return exporter.get_finished_spans()[0] def test_attributes_redacted(self): - redacted = redact_span(self._make_span(), default_redaction_policy()) + redacted = redact_span(self._make_span(), AttributeRedactionPolicy()) assert isinstance(redacted, ReadableSpan) attrs = redacted.attributes assert attrs is not None @@ -259,7 +259,7 @@ def test_attributes_redacted(self): assert attrs["gen_ai.request.model"] == "mistral-large" def test_event_attributes_redacted(self): - redacted = redact_span(self._make_span(), default_redaction_policy()) + redacted = redact_span(self._make_span(), AttributeRedactionPolicy()) event = redacted.events[0] assert event.name == "exception" attrs = event.attributes @@ -267,13 +267,13 @@ def test_event_attributes_redacted(self): assert attrs["exception.message"] == DEFAULT_REDACTED_VALUE def test_status_description_redacted(self): - redacted = redact_span(self._make_span(), default_redaction_policy()) + redacted = redact_span(self._make_span(), AttributeRedactionPolicy()) assert redacted.status.status_code == StatusCode.ERROR assert redacted.status.description == DEFAULT_REDACTED_VALUE def test_identity_preserved(self): original = self._make_span() - redacted = redact_span(original, default_redaction_policy()) + redacted = redact_span(original, AttributeRedactionPolicy()) assert redacted.context is not None assert original.context is not None assert redacted.context.span_id == original.context.span_id @@ -300,11 +300,11 @@ def test_wrapped_exporter_receives_redacted_spans(self): assert len(spans) == 1 attrs = spans[0].attributes assert attrs is not None - assert attrs["gen_ai.output.messages"] == DEFAULT_REDACTED_VALUE + assert attrs["gen_ai.output.messages"] == f"leak {DEFAULT_REDACTED_VALUE}" assert attrs["gen_ai.request.model"] == "mistral-large" def test_custom_policy_used(self): - spans = self._export_through(RegexRedactionPolicy()) + spans = self._export_through(AttributeRedactionPolicy()) attrs = spans[0].attributes assert attrs is not None - assert attrs["gen_ai.output.messages"] == f"leak {DEFAULT_REDACTED_VALUE}" + assert attrs["gen_ai.output.messages"] == DEFAULT_REDACTED_VALUE diff --git a/src/mistralai/extra/tests/test_telemetry.py b/src/mistralai/extra/tests/test_telemetry.py index 32c65608..5fc48622 100644 --- a/src/mistralai/extra/tests/test_telemetry.py +++ b/src/mistralai/extra/tests/test_telemetry.py @@ -692,13 +692,13 @@ def test_dedicated_wraps_exporter_by_default(self, clear_exporters): exporter = self._exporter_of(provider) assert isinstance(exporter, RedactingSpanExporter) assert exporter._exporter is FakeExporter.instances[0] - assert isinstance(exporter._policy, AttributeRedactionPolicy) + assert isinstance(exporter._policy, RegexRedactionPolicy) def test_redaction_true_wraps_with_default_policy(self, clear_exporters): provider = self._make_provider(redaction=True) exporter = self._exporter_of(provider) assert isinstance(exporter, RedactingSpanExporter) - assert isinstance(exporter._policy, AttributeRedactionPolicy) + assert isinstance(exporter._policy, RegexRedactionPolicy) def test_redaction_false_leaves_exporter_unwrapped(self, clear_exporters): provider = self._make_provider(redaction=False) @@ -707,7 +707,7 @@ def test_redaction_false_leaves_exporter_unwrapped(self, clear_exporters): assert exporter is FakeExporter.instances[0] def test_custom_policy_instance_is_used(self, clear_exporters): - policy = RegexRedactionPolicy() + policy = AttributeRedactionPolicy() provider = self._make_provider(redaction=policy) exporter = self._exporter_of(provider) assert isinstance(exporter, RedactingSpanExporter)