datarobot_genai.dragent.plugins.datarobot_moderation_middleware

datarobot_moderation_middleware

NAT (NeMo Agent Toolkit) middleware: DataRobot LLM guardrails for DRAgent workflows.

Registered under the nat.plugins distribution entry datarobot_moderation_middleware so NAT loads @register_middleware without a custom recipe mapping (_type: datarobot_moderation).

Expected workflow contracts:

• Input: RunAgentInput (or DRAgentRunAgentInput) for DRAgent workflows, or ChatRequest / ChatRequestOrMessage for native NAT chat agents.
• Output: always DRAgentEventResponse (both single and streaming). Native NAT agents emit str / ChatResponseChunk; the innermost datarobot_dragent_normalization middleware converts that into DRAgentEventResponse before this middleware runs, so moderation only ever handles the canonical type.

Guard configuration (_type: datarobot_moderation):

• Inline (preferred) — nest guards under middleware.<name>.moderation in workflow.yaml. When present, this block is used even if moderation_config.yaml also exists.
• DRUM-style file (fallback) — when moderation is omitted, load moderation_config.yaml from model_dir (defaults to the directory containing workflow.yaml, resolved from DRAGENT_CONFIG_FILE when set, otherwise the process working directory).
• If neither source is present or both are empty, the middleware is a no-op.

ModerationPipeline.stream_response_async only accepts OpenAI ChatCompletionChunk; DRAgent streaming uses convert_dragent_event_response_to_openai_chat_completion_chunk at that boundary, then reverses to AG-UI on the way out.

DataRobotModerationConfig

Bases: FunctionMiddlewareBaseConfig

NAT middleware: DataRobot prescore / postscore guards.

The middleware is a no-op (enabled is False) when no guards are configured in the inline moderation block or in moderation_config.yaml.

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

class DataRobotModerationConfig(
    FunctionMiddlewareBaseConfig,  # type: ignore[misc]
    name="datarobot_moderation",  # type: ignore[call-arg]
):
    """NAT middleware: DataRobot prescore / postscore guards.

    The middleware is a no-op (``enabled`` is ``False``) when no guards are configured in the
    inline ``moderation`` block or in ``moderation_config.yaml``.
    """

    model_dir: str | None = Field(
        default=None,
        description=(
            "Directory containing ``moderation_config.yaml`` and guard assets (DRUM custom model "
            "layout). Used for the inline ``moderation`` block and as the fallback file location. "
            "Defaults to the directory containing ``workflow.yaml`` (``DRAGENT_CONFIG_FILE``), "
            "or the process working directory when that env var is unset."
        ),
    )
    moderation: ModerationConfig | None = Field(
        default=None,
        description=(
            "Inline guard configuration (``ModerationConfig`` from datarobot_dome). When set, "
            "takes priority over ``moderation_config.yaml``."
        ),
    )

DataRobotModerationMiddleware

Bases: FunctionMiddleware

Guardrails middleware for DRAgent NAT workflows and native NAT chat agents.

Prescore/postscore run on AG-UI text of DRAgentEventResponse; streaming moderation uses dragent_event_response_to_dome_chunk at the dome boundary. Native NAT agents (ChatRequest / ChatRequestOrMessage in, str / ChatResponseChunk out) are first converted to DRAgentEventResponse by the innermost datarobot_dragent_normalization middleware, so this middleware only ever sees the canonical output type.

When no guards are configured (missing inline block and YAML file, or empty guard list), load_llm_moderation_pipeline returns None and this middleware is a no-op (enabled is False) without requiring DataRobot credentials.

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

class DataRobotModerationMiddleware(
    FunctionMiddleware,  # type: ignore[misc]
):
    """Guardrails middleware for DRAgent NAT workflows and native NAT chat agents.

    Prescore/postscore run on AG-UI text of ``DRAgentEventResponse``; streaming moderation uses
    ``dragent_event_response_to_dome_chunk`` at the dome boundary. Native NAT agents
    (``ChatRequest`` / ``ChatRequestOrMessage`` in, ``str`` / ``ChatResponseChunk`` out) are first
    converted to ``DRAgentEventResponse`` by the innermost ``datarobot_dragent_normalization``
    middleware, so this middleware only ever sees the canonical output type.

    When no guards are configured (missing inline block and YAML file, or empty guard list),
    ``load_llm_moderation_pipeline`` returns ``None`` and this middleware is a no-op
    (``enabled`` is ``False``) without requiring DataRobot credentials.
    """

    def __init__(self, config: DataRobotModerationConfig, builder: Builder) -> None:  # noqa: ARG002
        super().__init__()
        self._config = config
        self._moderation = load_llm_moderation_pipeline(config)

    def _get_moderation(self) -> ModerationPipeline | None:
        """Return the moderation pipeline, retrying discovery if startup missed ``workflow.yaml``."""  # noqa: E501
        if self._moderation is None:
            publish_dragent_config_file_env()
            self._moderation = load_llm_moderation_pipeline(self._config)
        return self._moderation

    @property
    def enabled(self) -> bool:
        return self._get_moderation() is not None

    async def function_middleware_invoke(
        self,
        *args: Any,
        call_next: CallNext,
        context: FunctionMiddlewareContext,
        **kwargs: Any,
    ) -> Any:
        """Run prescore guards; skip the agent when prescore blocks the prompt.

        The default NAT ``FunctionMiddleware.function_middleware_invoke`` always awaits
        ``call_next`` after ``pre_invoke``. When prescore blocks, ``pre_invoke`` sets
        ``ctx.output`` to the guard response; we return it immediately so ``call_next``
        (and thus the LLM) is never invoked.

        On a moderation failure, ``pre_invoke``/``post_invoke`` set ``ctx.output`` to a terminal
        ``RUN_ERROR``; converters adapt it per route (non-streaming raises, NAT returns 422;
        streaming frames it). Mid-stream: ``_moderated_dragent_stream``.
        """
        ctx = InvocationContext(
            function_context=context,
            original_args=args,
            original_kwargs=dict(kwargs),
            modified_args=args,
            modified_kwargs=dict(kwargs),
            output=None,
        )
        try:
            result = await self.pre_invoke(ctx)
            if result is not None:
                ctx = result
            if ctx.output is not None:
                return ctx.output
            ctx.output = await call_next(*ctx.modified_args, **ctx.modified_kwargs)
            result = await self.post_invoke(ctx)
            if result is not None:
                ctx = result
            return ctx.output
        finally:
            _clear_moderation_invoke_state_if_set()

    async def pre_invoke(self, context: InvocationContext) -> InvocationContext | None:
        """Pre-invocation hook called before the function is invoked.

        Args:
            context: Invocation context containing function metadata and args

        Returns
        -------
            InvocationContext if modified (including when the prompt is blocked and
            ``context.output`` holds the guard message), or None to pass through unchanged.
        """
        workflow_input = _require_workflow_input(context.original_args)
        moderation = self._get_moderation()
        if moderation is None:
            return None

        try:
            pipeline = moderation._pipeline

            prompt_column_name = pipeline.get_input_column(GuardStage.PROMPT)
            prompt = moderation_prompt_from_workflow_input(workflow_input)

            # Step 1: Prescore via ``ModerationPipeline.evaluate_prompt_async`` (non-blocking).
            # Wrap in the NAT workflow trace context: dome emits its ``evaluate_prompt`` span via
            # its own tracer provider, which our ``NatWorkflowTracer`` wrapper never patches.
            # Without an active workflow parent in context here (moderation is the outer middleware,
            # so this runs before the ``datarobot_agent`` span exists) that span would start a
            # disconnected trace.
            with use_nat_workflow_trace_context():
                prompt_eval, prescore_latency, prescore_df = await moderation.evaluate_prompt_async(
                    prompt
                )

            if prompt_eval.blocked:
                # If all prompts in the input are blocked, means history as well as the prompt
                # are not worthy to be sent to LLM. No invoke state: post_invoke / streaming never
                # run.
                context.output = _dragent_event_response_from_blocked_prompt_eval(prompt_eval)
                return context

            prompt_sent, workflow_rewritten = _prompt_sent_after_prescore_replacement(
                workflow_input,
                original_prompt=prompt,
                prompt_eval=prompt_eval,
                prescore_df=prescore_df,
                prompt_column=prompt_column_name,
            )
            _set_moderation_invoke_state(
                prompt=prompt_sent,
                prescore_df=prescore_df,
                latency_so_far=prescore_latency,
            )
            # Return context only when workflow input was rewritten (signals modified_args to NAT).
            return context if workflow_rewritten else None
        except Exception as exc:
            _logger.exception("Moderation prescore failed")
            context.output = run_error_response(f"Moderation failed: {exc}")
            return context

    async def post_invoke(self, context: InvocationContext) -> InvocationContext | None:
        """Post-invocation hook called after the function returns.

        Args:
            context: Invocation context containing function metadata, args, and output

        Returns
        -------
            InvocationContext if modified, or None to pass through unchanged (for example when
            the response has no assistant text to moderate, or prescore did not run). When guards
            are configured, agent output must already be ``DRAgentEventResponse`` (see
            ``datarobot_dragent_normalization``); unexpected types raise ``TypeError``.
        """
        original_output = context.output
        moderation = self._get_moderation()
        if moderation is None:
            return None

        # The ``datarobot_dragent_normalization`` middleware (declared innermost) converts every
        # agent's native output into ``DRAgentEventResponse`` before this hook runs, so moderation
        # only ever handles that canonical type.
        original_output = _require_dragent_event_response(original_output)
        if not _response_has_assistant_text_deltas(original_output):
            return None
        response_text = convert_dragent_event_response_to_str(original_output)

        if not response_text.strip():
            return None

        pipeline = moderation._pipeline
        state = _moderation_invoke_state_ctx.get()
        if state is None:
            return None

        try:
            # ==================================================================
            # Step 3: Postscore via ``ModerationPipeline.evaluate_response`` (same path as
            # ``_run_stage`` in dome) when response text is present.
            prompt_column_name = pipeline.get_input_column(GuardStage.PROMPT)
            # Wrap in the NAT workflow trace context so dome's ``evaluate_response`` span joins the
            # request trace (see the pre_invoke prescore call for the full rationale).
            with use_nat_workflow_trace_context():
                response_eval, _, _postscore_df = await moderation.evaluate_response_async(
                    response_text,
                    prompt=state.prompt,
                )

            prompt_eval = _from_dataframe(state.prescore_df, prompt_column_name)

            if response_eval.blocked:
                response_message = response_eval.blocked_message or ""
                finish_reason = "content_filter"
            elif response_eval.replaced:
                response_message = response_eval.replacement or ""
                finish_reason = "content_filter"
            else:
                response_message = response_text
                finish_reason = "stop"

            moderated_dr = _dragent_event_response_from_postscore_assistant_text(
                response_message,
                finish_reason,
                response_eval,
                upstream_model=_upstream_model_from_dragent_response(original_output),
                prompt_eval=prompt_eval,
            )
            preserve_ag_ui_envelope = (
                len(original_output.events) > 1
                and finish_reason == "stop"
                and not response_eval.blocked
                and not response_eval.replaced
                and response_message == response_text
            )
            if preserve_ag_ui_envelope:
                context.output = _merge_moderations_into_multi_event_response(
                    original_output, moderated_dr
                )
            else:
                context.output = moderated_dr

            return context
        except Exception as exc:
            _logger.exception("Moderation postscore failed")
            context.output = run_error_response(f"Moderation failed: {exc}")
            return context

    async def function_middleware_stream(
        self,
        *args: Any,
        call_next: CallNextStream,
        context: FunctionMiddlewareContext,
        **kwargs: Any,
    ) -> AsyncIterator[DRAgentEventResponse]:
        """Execute middleware hooks around DRAgent streaming (``DRAgentEventResponse`` chunks).

        Pre-invoke runs once before streaming starts.
        Moderation is applied per-chunk as they stream through.

        Note: Framework checks ``enabled`` before calling this method.
        You do NOT need to check ``enabled`` yourself.

        Args:
            args: Positional arguments for the function (first arg is typically the input value).
            call_next: Callable to invoke next middleware or target stream.
            context: Static function metadata.
            kwargs: Keyword arguments for the function.

        Yields
        ------
            Stream chunks with per-delta postscore via ``ModerationPipeline.stream_response_async``.

        When prescore blocks the prompt, ``pre_invoke`` sets ``ctx.output``; we yield that
        response once and return without calling ``call_next``, so the LLM stream is never
        started (NAT's default stream middleware always iterates ``call_next`` after
        ``pre_invoke``).
        """
        ctx = InvocationContext(
            function_context=context,
            original_args=args,
            original_kwargs=dict(kwargs),
            modified_args=args,
            modified_kwargs=dict(kwargs),
            output=None,
        )

        try:
            result = await self.pre_invoke(ctx)
            if result is not None:
                ctx = result
            if ctx.output is not None:
                yield ctx.output
                return

            # Prescore populates invoke state. If ``pre_invoke`` returned early (e.g. no
            # workflow input / prescore skipped), pass the stream through unchanged.
            stream_state = _moderation_invoke_state_ctx.get()
            # Clear at read time, in the context that set it (pre_invoke): the reset token is only
            # valid here. Deferring to the generator ``finally`` runs in a different context.
            _clear_moderation_invoke_state_if_set()
            if stream_state is None:
                async with contextlib.aclosing(
                    cast(
                        AsyncGenerator[DRAgentEventResponse, None],
                        call_next(*ctx.modified_args, **ctx.modified_kwargs),
                    )
                ) as upstream:
                    async for chunk in upstream:
                        yield chunk
                return

            moderation = self._get_moderation()
            assert moderation is not None

            async with contextlib.aclosing(
                _moderated_dragent_stream(
                    # Pin the upstream stream to one Context: dome drains part of it from its
                    # own feed task, and NAT's ``push_active_function`` token cannot be reset
                    # across Contexts. See ``_context_pinned_stream``.
                    _context_pinned_stream(
                        _validated_dragent_stream(
                            cast(
                                AsyncGenerator[Any, None],
                                call_next(*ctx.modified_args, **ctx.modified_kwargs),
                            ),
                        ),
                    ),
                    moderation=moderation,
                    stream_state=stream_state,
                )
            ) as moderated_stream:
                async for response in moderated_stream:
                    yield response
        finally:
            _clear_moderation_invoke_state_if_set()

function_middleware_invoke async

function_middleware_invoke(*args: Any, call_next: CallNext, context: FunctionMiddlewareContext, **kwargs: Any) -> Any

Run prescore guards; skip the agent when prescore blocks the prompt.

The default NAT FunctionMiddleware.function_middleware_invoke always awaits call_next after pre_invoke. When prescore blocks, pre_invoke sets ctx.output to the guard response; we return it immediately so call_next (and thus the LLM) is never invoked.

On a moderation failure, pre_invoke/post_invoke set ctx.output to a terminal RUN_ERROR; converters adapt it per route (non-streaming raises, NAT returns 422; streaming frames it). Mid-stream: _moderated_dragent_stream.

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

async def function_middleware_invoke(
    self,
    *args: Any,
    call_next: CallNext,
    context: FunctionMiddlewareContext,
    **kwargs: Any,
) -> Any:
    """Run prescore guards; skip the agent when prescore blocks the prompt.

    The default NAT ``FunctionMiddleware.function_middleware_invoke`` always awaits
    ``call_next`` after ``pre_invoke``. When prescore blocks, ``pre_invoke`` sets
    ``ctx.output`` to the guard response; we return it immediately so ``call_next``
    (and thus the LLM) is never invoked.

    On a moderation failure, ``pre_invoke``/``post_invoke`` set ``ctx.output`` to a terminal
    ``RUN_ERROR``; converters adapt it per route (non-streaming raises, NAT returns 422;
    streaming frames it). Mid-stream: ``_moderated_dragent_stream``.
    """
    ctx = InvocationContext(
        function_context=context,
        original_args=args,
        original_kwargs=dict(kwargs),
        modified_args=args,
        modified_kwargs=dict(kwargs),
        output=None,
    )
    try:
        result = await self.pre_invoke(ctx)
        if result is not None:
            ctx = result
        if ctx.output is not None:
            return ctx.output
        ctx.output = await call_next(*ctx.modified_args, **ctx.modified_kwargs)
        result = await self.post_invoke(ctx)
        if result is not None:
            ctx = result
        return ctx.output
    finally:
        _clear_moderation_invoke_state_if_set()

pre_invoke async

pre_invoke(context: InvocationContext) -> InvocationContext | None

Pre-invocation hook called before the function is invoked.

Parameters:

Name	Type	Description	Default
context	InvocationContext	Invocation context containing function metadata and args	required

Returns

InvocationContext if modified (including when the prompt is blocked and
``context.output`` holds the guard message), or None to pass through unchanged.

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

async def pre_invoke(self, context: InvocationContext) -> InvocationContext | None:
    """Pre-invocation hook called before the function is invoked.

    Args:
        context: Invocation context containing function metadata and args

    Returns
    -------
        InvocationContext if modified (including when the prompt is blocked and
        ``context.output`` holds the guard message), or None to pass through unchanged.
    """
    workflow_input = _require_workflow_input(context.original_args)
    moderation = self._get_moderation()
    if moderation is None:
        return None

    try:
        pipeline = moderation._pipeline

        prompt_column_name = pipeline.get_input_column(GuardStage.PROMPT)
        prompt = moderation_prompt_from_workflow_input(workflow_input)

        # Step 1: Prescore via ``ModerationPipeline.evaluate_prompt_async`` (non-blocking).
        # Wrap in the NAT workflow trace context: dome emits its ``evaluate_prompt`` span via
        # its own tracer provider, which our ``NatWorkflowTracer`` wrapper never patches.
        # Without an active workflow parent in context here (moderation is the outer middleware,
        # so this runs before the ``datarobot_agent`` span exists) that span would start a
        # disconnected trace.
        with use_nat_workflow_trace_context():
            prompt_eval, prescore_latency, prescore_df = await moderation.evaluate_prompt_async(
                prompt
            )

        if prompt_eval.blocked:
            # If all prompts in the input are blocked, means history as well as the prompt
            # are not worthy to be sent to LLM. No invoke state: post_invoke / streaming never
            # run.
            context.output = _dragent_event_response_from_blocked_prompt_eval(prompt_eval)
            return context

        prompt_sent, workflow_rewritten = _prompt_sent_after_prescore_replacement(
            workflow_input,
            original_prompt=prompt,
            prompt_eval=prompt_eval,
            prescore_df=prescore_df,
            prompt_column=prompt_column_name,
        )
        _set_moderation_invoke_state(
            prompt=prompt_sent,
            prescore_df=prescore_df,
            latency_so_far=prescore_latency,
        )
        # Return context only when workflow input was rewritten (signals modified_args to NAT).
        return context if workflow_rewritten else None
    except Exception as exc:
        _logger.exception("Moderation prescore failed")
        context.output = run_error_response(f"Moderation failed: {exc}")
        return context

post_invoke async

post_invoke(context: InvocationContext) -> InvocationContext | None

Post-invocation hook called after the function returns.

Parameters:

Name	Type	Description	Default
context	InvocationContext	Invocation context containing function metadata, args, and output	required

Returns

InvocationContext if modified, or None to pass through unchanged (for example when
the response has no assistant text to moderate, or prescore did not run). When guards
are configured, agent output must already be ``DRAgentEventResponse`` (see
``datarobot_dragent_normalization``); unexpected types raise ``TypeError``.

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

async def post_invoke(self, context: InvocationContext) -> InvocationContext | None:
    """Post-invocation hook called after the function returns.

    Args:
        context: Invocation context containing function metadata, args, and output

    Returns
    -------
        InvocationContext if modified, or None to pass through unchanged (for example when
        the response has no assistant text to moderate, or prescore did not run). When guards
        are configured, agent output must already be ``DRAgentEventResponse`` (see
        ``datarobot_dragent_normalization``); unexpected types raise ``TypeError``.
    """
    original_output = context.output
    moderation = self._get_moderation()
    if moderation is None:
        return None

    # The ``datarobot_dragent_normalization`` middleware (declared innermost) converts every
    # agent's native output into ``DRAgentEventResponse`` before this hook runs, so moderation
    # only ever handles that canonical type.
    original_output = _require_dragent_event_response(original_output)
    if not _response_has_assistant_text_deltas(original_output):
        return None
    response_text = convert_dragent_event_response_to_str(original_output)

    if not response_text.strip():
        return None

    pipeline = moderation._pipeline
    state = _moderation_invoke_state_ctx.get()
    if state is None:
        return None

    try:
        # ==================================================================
        # Step 3: Postscore via ``ModerationPipeline.evaluate_response`` (same path as
        # ``_run_stage`` in dome) when response text is present.
        prompt_column_name = pipeline.get_input_column(GuardStage.PROMPT)
        # Wrap in the NAT workflow trace context so dome's ``evaluate_response`` span joins the
        # request trace (see the pre_invoke prescore call for the full rationale).
        with use_nat_workflow_trace_context():
            response_eval, _, _postscore_df = await moderation.evaluate_response_async(
                response_text,
                prompt=state.prompt,
            )

        prompt_eval = _from_dataframe(state.prescore_df, prompt_column_name)

        if response_eval.blocked:
            response_message = response_eval.blocked_message or ""
            finish_reason = "content_filter"
        elif response_eval.replaced:
            response_message = response_eval.replacement or ""
            finish_reason = "content_filter"
        else:
            response_message = response_text
            finish_reason = "stop"

        moderated_dr = _dragent_event_response_from_postscore_assistant_text(
            response_message,
            finish_reason,
            response_eval,
            upstream_model=_upstream_model_from_dragent_response(original_output),
            prompt_eval=prompt_eval,
        )
        preserve_ag_ui_envelope = (
            len(original_output.events) > 1
            and finish_reason == "stop"
            and not response_eval.blocked
            and not response_eval.replaced
            and response_message == response_text
        )
        if preserve_ag_ui_envelope:
            context.output = _merge_moderations_into_multi_event_response(
                original_output, moderated_dr
            )
        else:
            context.output = moderated_dr

        return context
    except Exception as exc:
        _logger.exception("Moderation postscore failed")
        context.output = run_error_response(f"Moderation failed: {exc}")
        return context

function_middleware_stream async

function_middleware_stream(*args: Any, call_next: CallNextStream, context: FunctionMiddlewareContext, **kwargs: Any) -> AsyncIterator[DRAgentEventResponse]

Execute middleware hooks around DRAgent streaming (DRAgentEventResponse chunks).

Pre-invoke runs once before streaming starts. Moderation is applied per-chunk as they stream through.

Note: Framework checks enabled before calling this method. You do NOT need to check enabled yourself.

Parameters:

Name	Type	Description	Default
args	Any	Positional arguments for the function (first arg is typically the input value).	()
call_next	CallNextStream	Callable to invoke next middleware or target stream.	required
context	FunctionMiddlewareContext	Static function metadata.	required
kwargs	Any	Keyword arguments for the function.	{}

Yields

Stream chunks with per-delta postscore via ``ModerationPipeline.stream_response_async``.

When prescore blocks the prompt, pre_invoke sets ctx.output; we yield that response once and return without calling call_next, so the LLM stream is never started (NAT's default stream middleware always iterates call_next after pre_invoke).

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

async def function_middleware_stream(
    self,
    *args: Any,
    call_next: CallNextStream,
    context: FunctionMiddlewareContext,
    **kwargs: Any,
) -> AsyncIterator[DRAgentEventResponse]:
    """Execute middleware hooks around DRAgent streaming (``DRAgentEventResponse`` chunks).

    Pre-invoke runs once before streaming starts.
    Moderation is applied per-chunk as they stream through.

    Note: Framework checks ``enabled`` before calling this method.
    You do NOT need to check ``enabled`` yourself.

    Args:
        args: Positional arguments for the function (first arg is typically the input value).
        call_next: Callable to invoke next middleware or target stream.
        context: Static function metadata.
        kwargs: Keyword arguments for the function.

    Yields
    ------
        Stream chunks with per-delta postscore via ``ModerationPipeline.stream_response_async``.

    When prescore blocks the prompt, ``pre_invoke`` sets ``ctx.output``; we yield that
    response once and return without calling ``call_next``, so the LLM stream is never
    started (NAT's default stream middleware always iterates ``call_next`` after
    ``pre_invoke``).
    """
    ctx = InvocationContext(
        function_context=context,
        original_args=args,
        original_kwargs=dict(kwargs),
        modified_args=args,
        modified_kwargs=dict(kwargs),
        output=None,
    )

    try:
        result = await self.pre_invoke(ctx)
        if result is not None:
            ctx = result
        if ctx.output is not None:
            yield ctx.output
            return

        # Prescore populates invoke state. If ``pre_invoke`` returned early (e.g. no
        # workflow input / prescore skipped), pass the stream through unchanged.
        stream_state = _moderation_invoke_state_ctx.get()
        # Clear at read time, in the context that set it (pre_invoke): the reset token is only
        # valid here. Deferring to the generator ``finally`` runs in a different context.
        _clear_moderation_invoke_state_if_set()
        if stream_state is None:
            async with contextlib.aclosing(
                cast(
                    AsyncGenerator[DRAgentEventResponse, None],
                    call_next(*ctx.modified_args, **ctx.modified_kwargs),
                )
            ) as upstream:
                async for chunk in upstream:
                    yield chunk
            return

        moderation = self._get_moderation()
        assert moderation is not None

        async with contextlib.aclosing(
            _moderated_dragent_stream(
                # Pin the upstream stream to one Context: dome drains part of it from its
                # own feed task, and NAT's ``push_active_function`` token cannot be reset
                # across Contexts. See ``_context_pinned_stream``.
                _context_pinned_stream(
                    _validated_dragent_stream(
                        cast(
                            AsyncGenerator[Any, None],
                            call_next(*ctx.modified_args, **ctx.modified_kwargs),
                        ),
                    ),
                ),
                moderation=moderation,
                stream_state=stream_state,
            )
        ) as moderated_stream:
            async for response in moderated_stream:
                yield response
    finally:
        _clear_moderation_invoke_state_if_set()

moderation_config_has_guards

moderation_config_has_guards(moderation: ModerationConfig) -> bool

Return whether moderation defines at least one guard across all targets.

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

def moderation_config_has_guards(moderation: ModerationConfig) -> bool:
    """Return whether ``moderation`` defines at least one guard across all targets."""
    return any(target.guards for target in moderation.targets)

resolve_moderation_model_dir

resolve_moderation_model_dir(model_dir: str | None) -> str

Resolve the base directory for guard assets and moderation_config.yaml.

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

def resolve_moderation_model_dir(model_dir: str | None) -> str:
    """Resolve the base directory for guard assets and ``moderation_config.yaml``."""
    if model_dir is not None:
        return os.path.abspath(model_dir)
    return _default_moderation_model_dir()

moderation_config_file_path

moderation_config_file_path(model_dir: str | None) -> Path

Return the DRUM-style moderation_config.yaml path under model_dir (or workflow dir).

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

def moderation_config_file_path(model_dir: str | None) -> Path:
    """Return the DRUM-style ``moderation_config.yaml`` path under ``model_dir`` (or workflow dir)."""  # noqa: E501
    return Path(resolve_moderation_model_dir(model_dir)) / MODERATION_CONFIG_FILE_NAME

load_moderation_config_from_file

load_moderation_config_from_file(model_dir: str | None) -> ModerationConfig | None

Load and validate moderation_config.yaml from a model directory (DRUM layout).

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

def load_moderation_config_from_file(model_dir: str | None) -> ModerationConfig | None:
    """Load and validate ``moderation_config.yaml`` from a model directory (DRUM layout)."""
    config_path = moderation_config_file_path(model_dir)
    if not config_path.is_file():
        return None
    raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
    if not isinstance(raw, dict):
        return None
    return ModerationConfig.model_validate(raw)

load_llm_moderation_pipeline

load_llm_moderation_pipeline(config: DataRobotModerationConfig) -> ModerationPipeline | None

Build an LLM moderation pipeline from inline moderation or moderation_config.yaml.

Returns None when moderation is disabled, no configuration source is available, or the resolved source has no guards, so the middleware is a no-op and can be listed unconditionally in workflow.yaml. The inline moderation block takes priority over the YAML file.

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

def load_llm_moderation_pipeline(config: DataRobotModerationConfig) -> ModerationPipeline | None:
    """Build an LLM moderation pipeline from inline ``moderation`` or ``moderation_config.yaml``.

    Returns ``None`` when moderation is disabled, no configuration source is available, or the
    resolved source has no guards, so the middleware is a no-op and can be listed unconditionally
    in ``workflow.yaml``. The inline ``moderation`` block takes priority over the YAML file.
    """
    if get_runtime_parameter_value_bool(DISABLE_MODERATION_RUNTIME_PARAM_NAME, default_value=False):
        _logger.warning("Moderation is disabled via runtime parameter on the model")
        return None

    if config.moderation is not None:
        return _load_llm_moderation_pipeline_from_inline(config)
    return _load_llm_moderation_pipeline_from_config_file(config)

dragent_event_response_to_dome_chunk

dragent_event_response_to_dome_chunk(response: DRAgentEventResponse) -> ChatCompletionChunk

Convert one DRAgent stream chunk to an OpenAI chunk for ModerationPipeline streaming.

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

def dragent_event_response_to_dome_chunk(
    response: DRAgentEventResponse,
) -> ChatCompletionChunk:
    """Convert one DRAgent stream chunk to an OpenAI chunk for ``ModerationPipeline`` streaming."""
    completion_chunk = convert_dragent_event_response_to_openai_chat_completion_chunk(response)
    event_tool_calls = _tool_calls_from_ag_ui_events(response.events)
    if not event_tool_calls:
        return completion_chunk
    stream_choice = completion_chunk.choices[0]
    new_delta = stream_choice.delta.model_copy(
        update={"tool_calls": _message_tool_calls_to_openai_delta_tool_calls(event_tool_calls)}
    )
    return ChatCompletionChunk(
        id=completion_chunk.id,
        choices=[
            OpenAIChunkChoice(
                index=0,
                delta=new_delta,
                finish_reason="tool_calls",
            )
        ],
        created=completion_chunk.created,
        model=completion_chunk.model,
        object="chat.completion.chunk",
        usage=completion_chunk.usage,
    )

dome_chunk_to_dragent_event_response

dome_chunk_to_dragent_event_response(completion: ChatCompletionChunk, *, response_eval: EvaluationResult | None = None, source_ag_ui_events: list[Any] | None = None, stream_tool_index_map: dict[int, str] | None = None) -> DRAgentEventResponse

Convert a moderated OpenAI streaming chunk back to DRAgentEventResponse.

When response_eval is set (non-streaming postscore path), datarobot_moderations is taken from response_eval.metrics; otherwise it is read from the completion's moderation sidecar attribute (streaming chunks from ModerationIterator).

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

def dome_chunk_to_dragent_event_response(
    completion: ChatCompletionChunk,
    *,
    response_eval: EvaluationResult | None = None,
    source_ag_ui_events: list[Any] | None = None,
    stream_tool_index_map: dict[int, str] | None = None,
) -> DRAgentEventResponse:
    """Convert a moderated OpenAI streaming chunk back to ``DRAgentEventResponse``.

    When ``response_eval`` is set (non-streaming postscore path), ``datarobot_moderations`` is
    taken from ``response_eval.metrics``; otherwise it is read from the completion's moderation
    sidecar attribute (streaming chunks from ``ModerationIterator``).
    """
    chunk = _openai_chat_completion_chunk_to_nat_chat_response_chunk(completion)
    d = completion.choices[0].delta
    idx_map = stream_tool_index_map if stream_tool_index_map is not None else {}
    events: list[Any] = []
    if d.content:
        idx_map.clear()
        events.extend(
            _streaming_text_events_from_openai_chunk(
                completion, source_ag_ui_events=source_ag_ui_events
            )
        )
    elif not d.tool_calls:
        events.extend(
            _streaming_text_events_from_openai_chunk(
                completion, source_ag_ui_events=source_ag_ui_events
            )
        )
    if d.tool_calls:
        if d.content and events:
            last = events[-1]
            if isinstance(last, (TextMessageContentEvent, TextMessageChunkEvent)):
                events.append(TextMessageEndEvent(message_id=last.message_id))
        parent_id = _infer_parent_message_id_for_tool_calls(source_ag_ui_events, events)
        events.extend(
            _agui_tool_events_from_openai_delta_tool_calls(
                d.tool_calls,
                parent_message_id=parent_id,
                tool_index_map=idx_map,
            )
        )
    usage_metrics = _openai_usage_to_usage_metrics(completion.usage) or default_usage_metrics()
    datarobot_moderations = (
        _datarobot_moderations_from_evaluation_result(response_eval)
        if response_eval is not None
        else _datarobot_moderations_from_completion(completion)
    )
    return DRAgentEventResponse(
        events=events,
        usage_metrics=usage_metrics,
        original_chunk=chunk,
        model=completion.model,
        datarobot_moderations=datarobot_moderations,
    )

moderation_prompt_from_workflow_input

moderation_prompt_from_workflow_input(workflow_input: WorkflowInput) -> str

Extract the prescore prompt string from AG-UI or NAT chat input.

Delegates to get_chat_prompt on completion-params built from workflow_input. ChatRequestOrMessage with only input_message (no messages) is handled directly because get_chat_prompt requires a non-empty messages list.

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

def moderation_prompt_from_workflow_input(workflow_input: WorkflowInput) -> str:
    """Extract the prescore prompt string from AG-UI or NAT chat input.

    Delegates to ``get_chat_prompt`` on completion-params built from ``workflow_input``.
    ``ChatRequestOrMessage`` with only ``input_message`` (no ``messages``) is handled directly
    because ``get_chat_prompt`` requires a non-empty ``messages`` list.
    """
    if (
        isinstance(workflow_input, ChatRequestOrMessage)
        and workflow_input.input_message is not None
    ):
        return workflow_input.input_message
    return get_chat_prompt(workflow_input_to_completion_dict(workflow_input))

nat_chat_request_like_to_completion_dict

nat_chat_request_like_to_completion_dict(request: ChatRequest | ChatRequestOrMessage) -> dict[str, Any]

Build completion-style params from NAT chat request models (LLM Gateway path).

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

def nat_chat_request_like_to_completion_dict(
    request: ChatRequest | ChatRequestOrMessage,
) -> dict[str, Any]:
    """Build completion-style params from NAT chat request models (LLM Gateway path)."""
    return _normalize_nat_chat_request_completion_dict(request.model_dump(mode="json"))

workflow_input_to_completion_dict

workflow_input_to_completion_dict(workflow_input: WorkflowInput) -> dict[str, Any]

Build OpenAI-style completion params for prescore from AG-UI or NAT chat inputs.

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

def workflow_input_to_completion_dict(workflow_input: WorkflowInput) -> dict[str, Any]:
    """Build OpenAI-style completion params for prescore from AG-UI or NAT chat inputs."""
    if isinstance(workflow_input, RunAgentInput):
        return run_agent_input_to_completion_dict(workflow_input)
    if isinstance(workflow_input, (ChatRequest, ChatRequestOrMessage)):
        return nat_chat_request_like_to_completion_dict(workflow_input)
    raise TypeError(
        f"Unsupported workflow input type for moderation: {type(workflow_input).__name__}"
    )

datarobot_moderation_middleware async

datarobot_moderation_middleware(config: DataRobotModerationConfig, builder: Builder) -> AsyncIterator[DataRobotModerationMiddleware]

Register DataRobot LLM guard middleware for NAT/DRAgent workflows.

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

@register_middleware(  # type: ignore[untyped-decorator]
    config_type=DataRobotModerationConfig
)
async def datarobot_moderation_middleware(
    config: DataRobotModerationConfig,
    builder: Builder,  # noqa: ARG001
) -> AsyncIterator[DataRobotModerationMiddleware]:
    """Register DataRobot LLM guard middleware for NAT/DRAgent workflows."""
    yield DataRobotModerationMiddleware(config, builder)