datarobot_genai.dragent.plugins.datarobot_otel_conventions_middleware

datarobot_otel_conventions_middleware

DataRobotOtelConventionsMiddlewareConfig

Bases: FunctionMiddlewareBaseConfig

DataRobot Open Telemetry Conventions: https://docs.datarobot.com/en/docs/agentic-ai/agentic-develop/agentic-tracing-code.html#map-spans-and-attributes-to-the-tracing-table.

Source code in datarobot_genai/dragent/plugins/datarobot_otel_conventions_middleware.py

class DataRobotOtelConventionsMiddlewareConfig(
    FunctionMiddlewareBaseConfig,  # type: ignore[misc]
    name="datarobot_otel_conventions",  # type: ignore[call-arg]
):
    """DataRobot Open Telemetry Conventions:
    https://docs.datarobot.com/en/docs/agentic-ai/agentic-develop/agentic-tracing-code.html#map-spans-and-attributes-to-the-tracing-table.
    """

DataRobotOtelConventionsMiddleware

Bases: FunctionMiddleware

DataRobot Open Telemetry Conventions middleware for DRAgent NAT workflows.

Each invocation is wrapped in a dedicated datarobot_agent SDK span that carries the Tracing table attributes: the last user message becomes gen_ai.prompt and the workflow output becomes gen_ai.completion. NAT builds its own (non-SDK) spans and may not open an OTel parent span, so we create our own to guarantee the attributes have a recording span to live on. Tool-call spans are emitted as children. The streaming path is reimplemented so text deltas can be aggregated across chunks within a single invocation.

Source code in datarobot_genai/dragent/plugins/datarobot_otel_conventions_middleware.py

class DataRobotOtelConventionsMiddleware(
    FunctionMiddleware,  # type: ignore[misc]
):
    """DataRobot Open Telemetry Conventions middleware for DRAgent NAT workflows.

    Each invocation is wrapped in a dedicated ``datarobot_agent`` SDK span that
    carries the Tracing table attributes: the last user message becomes
    ``gen_ai.prompt`` and the workflow output becomes ``gen_ai.completion``.
    NAT builds its own (non-SDK) spans and may not open an OTel parent span, so
    we create our own to guarantee the attributes have a recording span to live
    on. Tool-call spans are emitted as children. The streaming path is
    reimplemented so text deltas can be aggregated across chunks within a single
    invocation.
    """

    def __init__(self, config: DataRobotOtelConventionsMiddlewareConfig, builder: Builder) -> None:  # noqa: ARG002
        super().__init__()

    @staticmethod
    def _prompt_from_args(args: tuple[Any, ...]) -> str | None:
        return _last_user_message_content(args[0]) if args else None

    @staticmethod
    def _completion_from_output(output: Any) -> str | None:
        # NAT non-streaming returns a plain str; every other path returns a
        # single aggregated DRAgentEventResponse.
        if isinstance(output, str):
            return output
        if isinstance(output, DRAgentEventResponse):
            return _response_text(output)
        return None

    async def function_middleware_invoke(
        self,
        *args: Any,
        call_next: CallNext,
        context: FunctionMiddlewareContext,  # noqa: ARG002
        **kwargs: Any,
    ) -> Any:
        with (
            use_nat_workflow_trace_context(),
            tracer.start_as_current_span(AGENT_SPAN_NAME) as span,
        ):
            prompt = self._prompt_from_args(args)
            if prompt is not None:
                span.set_attribute(GEN_AI_PROMPT, prompt)
            output = await call_next(*args, **kwargs)
            if isinstance(output, DRAgentEventResponse):
                _emit_tool_call_spans(output)
            completion = self._completion_from_output(output)
            if completion is not None:
                span.set_attribute(GEN_AI_COMPLETION, completion)
            return output

    async def function_middleware_stream(
        self,
        *args: Any,
        call_next: CallNextStream,
        context: FunctionMiddlewareContext,  # noqa: ARG002
        **kwargs: Any,
    ) -> AsyncIterator[Any]:
        with (
            use_nat_workflow_trace_context(),
            tracer.start_as_current_span(AGENT_SPAN_NAME) as span,
        ):
            prompt = self._prompt_from_args(args)
            if prompt is not None:
                span.set_attribute(GEN_AI_PROMPT, prompt)
            # Per-invocation accumulator; no cross-session state to manage.
            parts: list[str] = []
            async for chunk in call_next(*args, **kwargs):
                if isinstance(chunk, DRAgentEventResponse):
                    _emit_tool_call_spans(chunk)
                    text = _response_text(chunk)
                    if text:
                        parts.append(text)
                yield chunk
            if parts:
                span.set_attribute(GEN_AI_COMPLETION, "".join(parts))

datarobot_otel_conventions_middleware async

datarobot_otel_conventions_middleware(config: DataRobotOtelConventionsMiddlewareConfig, builder: Builder) -> AsyncIterator[DataRobotOtelConventionsMiddleware]

Register DataRobot Open Telemetry Conventions middleware for NAT/DRAgent workflows.

Source code in datarobot_genai/dragent/plugins/datarobot_otel_conventions_middleware.py

@register_middleware(  # type: ignore[untyped-decorator]
    config_type=DataRobotOtelConventionsMiddlewareConfig
)
async def datarobot_otel_conventions_middleware(
    config: DataRobotOtelConventionsMiddlewareConfig,
    builder: Builder,  # noqa: ARG001
) -> AsyncIterator[DataRobotOtelConventionsMiddleware]:
    """Register DataRobot Open Telemetry Conventions middleware for NAT/DRAgent workflows."""
    yield DataRobotOtelConventionsMiddleware(config, builder)