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)
                _mark_span_error_on_run_error(span, 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] = []
            open_text_ids: set[str] = set()
            try:
                async for chunk in call_next(*args, **kwargs):
                    if isinstance(chunk, DRAgentEventResponse):
                        _emit_tool_call_spans(chunk)
                        _mark_span_error_on_run_error(span, chunk)
                        track_open_text_in_events(open_text_ids, chunk.events)
                        text = _response_text(chunk)
                        if text:
                            parts.append(text)
                    yield chunk
            except Exception as exc:
                # Close open text segments, then end the run with a terminal RUN_ERROR instead of
                # propagating (matches the moderation middleware's failure path).
                logger.exception("Agent stream failed")
                for message_id in open_text_ids:
                    yield DRAgentEventResponse(
                        events=[TextMessageEndEvent(message_id=message_id)],
                        usage_metrics=default_usage_metrics(),
                    )
                error_response = run_error_response(str(exc))
                _mark_span_error_on_run_error(span, error_response)
                yield error_response
            finally:
                # Attach the completion in ``finally`` so it survives early teardown.
                # Downstream moderation may stop consuming and ``aclose()`` this generator
                # (throwing ``GeneratorExit`` at the ``yield``) before the loop exits normally
                # — e.g. when the moderation stream finishes before draining its source. Setting
                # the attribute after the loop would then be skipped, dropping ``gen_ai.completion``
                # even though the deltas were already seen. The span is still open here because the
                # enclosing ``with`` block outlives this ``finally``.
                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)