datarobot_genai.drmcp.core.telemetry

telemetry

OtelASGIMiddleware

Bases: BaseHTTPMiddleware

ASGI middleware extracts trace_id and parent span id from raw http request and set them in the context, so downsream trace can link to them if avaialble.

Source code in datarobot_genai/drmcp/core/telemetry.py

class OtelASGIMiddleware(BaseHTTPMiddleware):
    """ASGI middleware extracts trace_id and parent span id from raw http request
    and set them in the context, so downsream trace can link to them
    if avaialble.
    """

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        with with_otel_context(request.headers):
            response = await call_next(request)
            return response

initialize_telemetry

initialize_telemetry(mcp: FastMCP) -> None

Initialize OpenTelemetry for the FastMCP application.

Source code in datarobot_genai/drmcp/core/telemetry.py

def initialize_telemetry(mcp: FastMCP) -> None:
    """Initialize OpenTelemetry for the FastMCP application."""
    config = get_config()

    # If OpenTelemetry is disabled, return None
    if not config.otel_enabled:
        root_logger.info("OpenTelemetry is disabled")
        return None

    # If OTEL_ENTITY_ID is not set, skip telemetry
    if (
        not config.otel_entity_id
        and not config.otel_exporter_otlp_headers
        and not os.environ.get("OTEL_EXPORTER_OTLP_HEADERS")
    ):
        root_logger.info(
            "Neither OTEL_ENTITY_ID nor OTEL_EXPORTER_OTLP_HEADERS is set, skipping telemetry"
        )
        return None

    resource_attrs = {"datarobot.service.name": config.mcp_server_name}
    if config.otel_attributes:
        resource_attrs.update(_prepare_shared_attributes(config.otel_attributes))
    resource = Resource.create(resource_attrs)

    # Set up tracer provider with service name from config
    provider = TracerProvider(resource=resource)
    trace.set_tracer_provider(provider)

    # Setup environment
    _setup_otel_env_variables()

    # Setup OTEL exporter
    _setup_otel_exporter()
    # Metrics leg: the OTLP exporter resolves the endpoint/headers from the same
    # OTEL_EXPORTER_OTLP_* env vars set above (no-op when they are unset).
    bootstrap_metrics_provider(resource_attributes=resource_attrs)
    _setup_otel_logging(resource)

    # Setup HTTP client instrumentation
    if config.otel_enabled_http_instrumentors:
        _setup_http_instrumentors()

    mcp.add_middleware(OpenTelemetryMiddleware(__name__))

get_trace_id

get_trace_id() -> str | None

Get the current trace ID if available.

Source code in datarobot_genai/drmcp/core/telemetry.py

def get_trace_id() -> str | None:
    """Get the current trace ID if available."""
    current_span = trace.get_current_span()
    if not current_span:
        return None

    context: SpanContext = current_span.get_span_context()
    if not context.is_valid:
        return None

    return str(format_trace_id(context.trace_id))

trace_execution

trace_execution(trace_name: str | None = None, trace_type: str = 'tool') -> Callable[[T], T]

Trace tool execution.

Parameters:

Name	Type	Description	Default
trace_name	str | None	Optional name for the span. If not provided, uses the function name.	None
trace_type	str	Optional type for the span. If not provided, uses "tool".	'tool'

Example

@trace_execution() async def my_tool(self, param1: str) -> str: return "result"

@trace_execution("custom_name") async def another_tool(self, param1: str) -> str: return "result"

Source code in datarobot_genai/drmcp/core/telemetry.py

def trace_execution(trace_name: str | None = None, trace_type: str = "tool") -> Callable[[T], T]:
    """Trace tool execution.

    Args:
        trace_name: Optional name for the span. If not provided, uses the function name.
        trace_type: Optional type for the span. If not provided, uses "tool".

    Example:
        @trace_execution()
        async def my_tool(self, param1: str) -> str:
            return "result"

        @trace_execution("custom_name")
        async def another_tool(self, param1: str) -> str:
            return "result"
    """

    def decorator(func: T) -> T:
        def _create_span_for_tool(
            trace_name: str | None,
            trace_type: str,
            args: tuple[Any, ...],
            kwargs: dict[str, Any],
        ) -> Span:
            # Get span name from decorator arg, function name, or class method name
            span_name = trace_name
            if not span_name:
                if (
                    args
                    and hasattr(args[0], "__class__")
                    and not isinstance(args[0], (str, int, float, bool))
                ):
                    # If it's a method, include the class name
                    span_name = f"{args[0].__class__.__name__}.{func.__name__}"
                else:
                    # Just use the function name without any prefix
                    span_name = func.__name__

            # Start a new span
            tracer = trace.get_tracer(__name__)
            span = tracer.start_span(f"{trace_type}.{span_name}")

            # Add standard attributes: tool-specific only when trace_type is "tool"
            if trace_type == "tool":
                span.set_attribute("gen_ai.tool.name", span_name)
                span.set_attribute("gen_ai.operation.name", "execute_tool")
            else:
                span.set_attribute("gen_ai.operation.name", trace_type)
                span.set_attribute(f"{trace_type}.name", span_name)

            # Add tool parameters as span attributes
            _add_parameters_to_span(span, func, args, kwargs, trace_type)

            # Add configured attributes from config
            config = get_config()
            if config.otel_attributes:
                _set_otel_attributes(span, config.otel_attributes)

            return span

        @functools.wraps(func)
        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
            span = _create_span_for_tool(trace_name, trace_type, args, kwargs)
            try:
                result = await func(*args, **kwargs)
                span.set_status(Status(StatusCode.OK))
                return result
            except Exception as e:
                span.set_attribute("error.type", type(e).__name__)
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.record_exception(e)
                raise e
            finally:
                span.end()

        @functools.wraps(func)
        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
            span = _create_span_for_tool(trace_name, trace_type, args, kwargs)
            try:
                result = func(*args, **kwargs)
                span.set_status(Status(StatusCode.OK))
                return result
            except Exception as e:
                span.set_attribute("error.type", type(e).__name__)
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.record_exception(e)
                raise e
            finally:
                span.end()

        # Use appropriate wrapper based on whether the function is async
        return async_wrapper if inspect.iscoroutinefunction(func) else sync_wrapper  # type: ignore[return-value]

    return decorator