datarobot_genai.nat.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:

• DRAgent (dragent_fastapi + LangGraph / LlamaIndex / CrewAI per-user functions): input RunAgentInput (or DRAgentRunAgentInput), streaming output DRAgentEventResponse.
• Native NAT chat (LLM Gateway agents): input ChatRequest / ChatRequestOrMessage, non-streaming output ChatResponse.

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/nat/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.

• DRAgent (RunAgentInput in, DRAgentEventResponse stream out): prescore/postscore on AG-UI text; streaming moderation uses dragent_event_response_to_dome_chunk at the dome boundary.
• NAT chat (ChatRequest / ChatRequestOrMessage in, ChatResponse out): same guard pipeline with NAT message models instead of AG-UI.

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/nat/datarobot_moderation_middleware.py

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

    * **DRAgent** (``RunAgentInput`` in, ``DRAgentEventResponse`` stream out): prescore/postscore
      on AG-UI text; streaming moderation uses ``dragent_event_response_to_dome_chunk`` at the
      dome boundary.
    * **NAT chat** (``ChatRequest`` / ``ChatRequestOrMessage`` in, ``ChatResponse`` out): same
      guard pipeline with NAT message models instead of AG-UI.

    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``."""
        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.
        """
        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 = _workflow_input_from_args(context.original_args)
        if workflow_input is None:
            return None
        moderation = self._get_moderation()
        if moderation is None:
            return None

        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).
        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

    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.
        """
        original_output = context.output
        moderation = self._get_moderation()
        if moderation is None:
            return None

        if isinstance(original_output, DRAgentEventResponse):
            if not _response_has_assistant_text_deltas(original_output):
                return None
            response_text = convert_dragent_event_response_to_str(original_output)
        elif isinstance(original_output, ChatResponse):
            if not original_output.choices:
                return None
            response_text = original_output.choices[0].message.content or ""
        elif isinstance(original_output, str):
            response_text = original_output
        else:
            return None

        if not response_text.strip():
            return None

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

        # ==================================================================
        # 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)
        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"

        if isinstance(original_output, DRAgentEventResponse):
            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
        elif isinstance(original_output, ChatResponse):
            context.output = _nat_chat_response_from_postscore_assistant_text(
                response_message,
                finish_reason,
                original_output,
                response_eval,
                prompt_eval=prompt_eval,
            )
        else:
            context.output = response_message

        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()
            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(
                    cast(
                        AsyncGenerator[DRAgentEventResponse, 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.

Source code in datarobot_genai/nat/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.
    """
    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:

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

Source code in datarobot_genai/nat/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 = _workflow_input_from_args(context.original_args)
    if workflow_input is None:
        return None
    moderation = self._get_moderation()
    if moderation is None:
        return None

    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).
    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

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:

Type	Description
InvocationContext | None	InvocationContext if modified, or None to pass through unchanged.

Source code in datarobot_genai/nat/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.
    """
    original_output = context.output
    moderation = self._get_moderation()
    if moderation is None:
        return None

    if isinstance(original_output, DRAgentEventResponse):
        if not _response_has_assistant_text_deltas(original_output):
            return None
        response_text = convert_dragent_event_response_to_str(original_output)
    elif isinstance(original_output, ChatResponse):
        if not original_output.choices:
            return None
        response_text = original_output.choices[0].message.content or ""
    elif isinstance(original_output, str):
        response_text = original_output
    else:
        return None

    if not response_text.strip():
        return None

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

    # ==================================================================
    # 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)
    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"

    if isinstance(original_output, DRAgentEventResponse):
        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
    elif isinstance(original_output, ChatResponse):
        context.output = _nat_chat_response_from_postscore_assistant_text(
            response_message,
            finish_reason,
            original_output,
            response_eval,
            prompt_eval=prompt_eval,
        )
    else:
        context.output = response_message

    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:

Type	Description
AsyncIterator[DRAgentEventResponse]	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/nat/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()
        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(
                cast(
                    AsyncGenerator[DRAgentEventResponse, 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/nat/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/nat/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/nat/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)."""
    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/nat/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/nat/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/nat/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/nat/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/nat/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/nat/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/nat/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/nat/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)