Add tracing

Author: amenasria

poetry.lock

@@ -1,6 +1,6 @@
 [project]
 name = "mistralai"
-version = "1.9.11"
+version = "1.9.12"
 description = "Python Client SDK for the Mistral AI API."
 authors = [{ name = "Mistral" },]
 readme = "README-PYPI.md"
@@ -13,6 +13,10 @@ dependencies = [
     "typing-inspection >=0.4.0",
     "pyyaml (>=6.0.2,<7.0.0)",
     "invoke (>=2.2.0,<3.0.0)",
+    "opentelemetry-sdk (>=1.33.1,<2.0.0)",
+    "opentelemetry-api (>=1.33.1,<2.0.0)",
+    "opentelemetry-exporter-otlp-proto-http (>=1.37.0,<2.0.0)",
+    "opentelemetry-semantic-conventions (>=0.59b0,<0.60)",
 ]
 
 [tool.poetry]

pyproject.toml

@@ -1,5 +1,6 @@
 from .custom_user_agent import CustomUserAgentHook
 from .deprecation_warning import DeprecationWarningHook
+from .tracing import TracingHook
 from .types import Hooks
 
 # This file is only ever generated once on the first generation and then is free to be modified.
@@ -13,5 +14,9 @@ def init_hooks(hooks: Hooks):
     with an instance of a hook that implements that specific Hook interface
     Hooks are registered per SDK instance, and are valid for the lifetime of the SDK instance
     """
+    tracing_hook = TracingHook()
     hooks.register_before_request_hook(CustomUserAgentHook())
     hooks.register_after_success_hook(DeprecationWarningHook())
+    hooks.register_after_success_hook(tracing_hook)
+    hooks.register_before_request_hook(tracing_hook)
+    hooks.register_after_error_hook(tracing_hook)

src/mistralai/_hooks/registration.py

@@ -0,0 +1,50 @@
+from .types import (
+    AfterErrorContext,
+    AfterErrorHook,
+    AfterSuccessContext,
+    AfterSuccessHook,
+    BeforeRequestContext,
+    BeforeRequestHook,
+)
+from ..extra.observability.otel import (
+    get_traced_request_and_span,
+    get_traced_response,
+    get_response_and_error,
+    get_or_create_otel_tracer,
+)
+
+from opentelemetry.trace import Span
+from typing import Optional, Tuple, Union
+
+import httpx
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class TracingHook(BeforeRequestHook, AfterSuccessHook, AfterErrorHook):
+    def __init__(self) -> None:
+        self.tracing_enabled, self.tracer = get_or_create_otel_tracer()
+        self.request_span: Optional[Span] = None
+
+    def before_request(
+        self, hook_ctx: BeforeRequestContext, request: httpx.Request
+    ) -> Union[httpx.Request, Exception]:
+        request, self.request_span = get_traced_request_and_span(tracing_enabled=self.tracing_enabled, tracer=self.tracer, span=self.request_span, operation_id=hook_ctx.operation_id, request=request)
+        return request
+
+    def after_success(
+        self, hook_ctx: AfterSuccessContext, response: httpx.Response
+    ) -> Union[httpx.Response, Exception]:
+        response = get_traced_response(tracing_enabled=self.tracing_enabled, tracer=self.tracer, span=self.request_span, operation_id=hook_ctx.operation_id, response=response)
+        return response
+
+    def after_error(
+        self,
+        hook_ctx: AfterErrorContext,
+        response: Optional[httpx.Response],
+        error: Optional[Exception],
+    ) -> Union[Tuple[Optional[httpx.Response], Optional[Exception]], Exception]:
+        if response:
+            response, error = get_response_and_error(tracing_enabled=self.tracing_enabled, tracer=self.tracer, span=self.request_span, operation_id=hook_ctx.operation_id, response=response, error=error)
+        return response, error