datarobot-moderations 11.2.9__py3-none-any.whl → 11.2.11__py3-none-any.whl

datarobot_dome/guards/base.py

@@ -0,0 +1,259 @@
+#  ---------------------------------------------------------------------------------
+#  Copyright (c) 2026 DataRobot, Inc. and its affiliates. All rights reserved.
+#  Last updated 2025.
+#
+#  DataRobot, Inc. Confidential.
+#  This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+#  This file and its contents are subject to DataRobot Tool and Utility Agreement.
+#  For details, see
+#  https://www.datarobot.com/wp-content/uploads/2021/07/DataRobot-Tool-and-Utility-Agreement.pdf.
+#  ---------------------------------------------------------------------------------
+from abc import ABC
+
+from datarobot.enums import CustomMetricAggregationType
+from datarobot.enums import CustomMetricDirectionality
+
+from datarobot_dome.constants import AGENT_GOAL_ACCURACY_COLUMN_NAME
+from datarobot_dome.constants import COST_COLUMN_NAME
+from datarobot_dome.constants import CUSTOM_METRIC_DESCRIPTION_SUFFIX
+from datarobot_dome.constants import DEFAULT_PROMPT_COLUMN_NAME
+from datarobot_dome.constants import DEFAULT_RESPONSE_COLUMN_NAME
+from datarobot_dome.constants import FAITHFULLNESS_COLUMN_NAME
+from datarobot_dome.constants import GUIDELINE_ADHERENCE_COLUMN_NAME
+from datarobot_dome.constants import NEMO_GUARD_COLUMN_NAME
+from datarobot_dome.constants import ROUGE_1_COLUMN_NAME
+from datarobot_dome.constants import SPAN_PREFIX
+from datarobot_dome.constants import TASK_ADHERENCE_SCORE_COLUMN_NAME
+from datarobot_dome.constants import TOKEN_COUNT_COLUMN_NAME
+from datarobot_dome.constants import GuardAction
+from datarobot_dome.constants import GuardStage
+from datarobot_dome.constants import GuardType
+from datarobot_dome.constants import NemoEvaluatorType
+from datarobot_dome.constants import OOTBType
+
+
+def get_metric_column_name(
+    guard_type: GuardType,
+    ootb_type: OOTBType | None,
+    stage: GuardStage,
+    model_guard_target_name: str | None = None,
+    metric_name: str | None = None,
+    nemo_evaluator_type: str | None = None,
+) -> str:
+    """Gets the metric column name. Note that this function gets used in buzok code. If you update
+    it, please also update the moderation library in the buzok worker image.
+    """
+    if guard_type == GuardType.MODEL:
+        if model_guard_target_name is None:
+            raise ValueError(
+                "For the model guard type, a valid model_guard_target_name has to be provided."
+            )
+        metric_result_key = Guard.get_stage_str(stage) + "_" + model_guard_target_name
+    elif guard_type == GuardType.OOTB:
+        if ootb_type is None:
+            raise ValueError("For the OOTB type, a valid OOTB guard type has to be provided.")
+        elif ootb_type == OOTBType.TOKEN_COUNT:
+            metric_result_key = Guard.get_stage_str(stage) + "_" + TOKEN_COUNT_COLUMN_NAME
+        elif ootb_type == OOTBType.ROUGE_1:
+            metric_result_key = Guard.get_stage_str(stage) + "_" + ROUGE_1_COLUMN_NAME
+        elif ootb_type == OOTBType.FAITHFULNESS:
+            metric_result_key = Guard.get_stage_str(stage) + "_" + FAITHFULLNESS_COLUMN_NAME
+        elif ootb_type == OOTBType.AGENT_GOAL_ACCURACY:
+            metric_result_key = AGENT_GOAL_ACCURACY_COLUMN_NAME
+        elif ootb_type == OOTBType.CUSTOM_METRIC:
+            if metric_name is None:
+                raise ValueError(
+                    "For the custom metric type, a valid metric_name has to be provided."
+                )
+            metric_result_key = Guard.get_stage_str(stage) + "_" + metric_name
+        elif ootb_type == OOTBType.COST:
+            metric_result_key = COST_COLUMN_NAME
+        elif ootb_type == OOTBType.TASK_ADHERENCE:
+            metric_result_key = TASK_ADHERENCE_SCORE_COLUMN_NAME
+        elif ootb_type == OOTBType.GUIDELINE_ADHERENCE:
+            metric_result_key = GUIDELINE_ADHERENCE_COLUMN_NAME
+        else:
+            raise ValueError("The provided OOTB type is not implemented.")
+    elif guard_type == GuardType.NEMO_GUARDRAILS:
+        metric_result_key = Guard.get_stage_str(stage) + "_" + NEMO_GUARD_COLUMN_NAME
+    elif guard_type == GuardType.NEMO_EVALUATOR:
+        if nemo_evaluator_type == NemoEvaluatorType.LLM_JUDGE:
+            metric_result_key = f"{Guard.get_stage_str(stage)}_nemo_{nemo_evaluator_type}"
+        elif nemo_evaluator_type in NemoEvaluatorType.ALL:
+            metric_result_key = f"nemo_{nemo_evaluator_type}"
+        else:
+            raise ValueError("The provided NeMo Evaluator type is not implemented.")
+    else:
+        raise ValueError("The provided guard type is not implemented.")
+    return metric_result_key
+
+
+class GuardIntervention:
+    def __init__(self, intervention_config: dict) -> None:
+        self.action = intervention_config["action"]
+        self.message = intervention_config.get("message")
+        self.threshold = None
+        self.comparator = None
+        if (
+            "conditions" in intervention_config
+            and intervention_config["conditions"] is not None
+            and len(intervention_config["conditions"]) > 0
+        ):
+            self.threshold = intervention_config["conditions"][0].get("comparand")
+            self.comparator = intervention_config["conditions"][0].get("comparator")
+
+
+class Guard(ABC):
+    def __init__(self, config: dict, stage=None):
+        self._name = config["name"]
+        self._description = config.get("description")
+        self._type = config["type"]
+        self._stage = stage if stage else config["stage"]
+        self._pipeline = None
+        self.intervention = None
+        self._deployment_id = config.get("deployment_id")
+        self._dr_cm = None
+        self._faas_url = config.get("faas_url")
+        self._copy_citations = config["copy_citations"]
+        self.is_agentic = config.get("is_agentic", False)
+        self.metric_column_name = get_metric_column_name(
+            config["type"],
+            config.get("ootb_type"),
+            self._stage,
+            config.get("model_info", {}).get("target_name"),
+            config["name"],
+            config.get("nemo_evaluator_type"),
+        )
+        if config.get("intervention"):
+            self.intervention = GuardIntervention(config["intervention"])
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def description(self) -> str:
+        return self._description
+
+    @property
+    def type(self) -> GuardType:
+        return self._type
+
+    @property
+    def stage(self) -> GuardStage:
+        return self._stage
+
+    @property
+    def faas_url(self) -> str:
+        return self._faas_url
+
+    @property
+    def copy_citations(self) -> str:
+        return self._copy_citations
+
+    def set_pipeline(self, pipeline):
+        self._pipeline = pipeline
+
+    @property
+    def llm_type(self):
+        return self._llm_type
+
+    @staticmethod
+    def get_stage_str(stage):
+        return "Prompts" if stage == GuardStage.PROMPT else "Responses"
+
+    def get_input_column_name(self, stage) -> str:
+        match stage:
+            case GuardStage.PROMPT:
+                return DEFAULT_PROMPT_COLUMN_NAME
+            case GuardStage.RESPONSE:
+                return DEFAULT_RESPONSE_COLUMN_NAME
+            case _:
+                raise ValueError(f"Stage ({stage}) is not supported.")
+
+    def has_latency_custom_metric(self) -> bool:
+        """Determines if latency metric is tracked for this guard type. Default is True."""
+        return True
+
+    def get_latency_custom_metric_name(self):
+        return f"{self.name} Guard Latency"
+
+    def get_latency_custom_metric(self):
+        return {
+            "name": self.get_latency_custom_metric_name(),
+            "directionality": CustomMetricDirectionality.LOWER_IS_BETTER,
+            "units": "seconds",
+            "type": CustomMetricAggregationType.AVERAGE,
+            "baselineValue": 0,
+            "isModelSpecific": True,
+            "timeStep": "hour",
+            "description": (
+                f"{self.get_latency_custom_metric_name()}.  {CUSTOM_METRIC_DESCRIPTION_SUFFIX}"
+            ),
+        }
+
+    def has_average_score_custom_metric(self) -> bool:
+        """Determines if an average score metric is tracked for this guard type. Default is True."""
+        return True
+
+    def get_average_score_custom_metric_name(self, stage):
+        return f"{self.name} Guard Average Score for {self.get_stage_str(stage)}"
+
+    def get_average_score_metric(self, stage):
+        return {
+            "name": self.get_average_score_custom_metric_name(stage),
+            "directionality": CustomMetricDirectionality.LOWER_IS_BETTER,
+            "units": "probability",
+            "type": CustomMetricAggregationType.AVERAGE,
+            "baselineValue": 0,
+            "isModelSpecific": True,
+            "timeStep": "hour",
+            "description": (
+                f"{self.get_average_score_custom_metric_name(stage)}. "
+                f" {CUSTOM_METRIC_DESCRIPTION_SUFFIX}"
+            ),
+        }
+
+    def get_guard_enforced_custom_metric_name(self, stage, moderation_method):
+        if moderation_method == GuardAction.REPLACE:
+            return f"{self.name} Guard replaced {self.get_stage_str(stage)}"
+        return f"{self.name} Guard {moderation_method}ed {self.get_stage_str(stage)}"
+
+    def get_enforced_custom_metric(self, stage, moderation_method):
+        return {
+            "name": self.get_guard_enforced_custom_metric_name(stage, moderation_method),
+            "directionality": CustomMetricDirectionality.LOWER_IS_BETTER,
+            "units": "count",
+            "type": CustomMetricAggregationType.SUM,
+            "baselineValue": 0,
+            "isModelSpecific": True,
+            "timeStep": "hour",
+            "description": (
+                f"Number of {self.get_stage_str(stage)} {moderation_method}ed by the "
+                f"{self.name} guard.  {CUSTOM_METRIC_DESCRIPTION_SUFFIX}"
+            ),
+        }
+
+    def get_intervention_action(self):
+        if not self.intervention:
+            return GuardAction.NONE
+        return self.intervention.action
+
+    def get_comparand(self):
+        return self.intervention.threshold
+
+    def get_enforced_span_attribute_name(self, stage):
+        intervention_action = self.get_intervention_action()
+        if intervention_action in [GuardAction.BLOCK, GuardAction.REPORT]:
+            return f"{SPAN_PREFIX}.{stage.lower()}.{intervention_action}ed"
+        elif intervention_action == GuardAction.REPLACE:
+            return f"{SPAN_PREFIX}.{stage.lower()}.replaced"
+        else:
+            raise NotImplementedError
+
+    def get_span_column_name(self, _):
+        raise NotImplementedError
+
+    def get_span_attribute_name(self, _):
+        raise NotImplementedError

datarobot_dome/guards/guard_llm_mixin.py

@@ -1,5 +1,5 @@
 #  ---------------------------------------------------------------------------------
-#  Copyright (c) 2025 DataRobot, Inc. and its affiliates. All rights reserved.
+#  Copyright (c) 2026 DataRobot, Inc. and its affiliates. All rights reserved.
 #  Last updated 2025.
 #
 #  DataRobot, Inc. Confidential.
@@ -94,6 +94,8 @@ class GuardLLMMixin:
                 return f"{secret_env_var_name_prefix}{OOTBType.AGENT_GOAL_ACCURACY}_{llm_type_str}"
             elif config["ootb_type"] == OOTBType.TASK_ADHERENCE:
                 return f"{secret_env_var_name_prefix}{OOTBType.TASK_ADHERENCE}_{llm_type_str}"
+            elif config["ootb_type"] == OOTBType.GUIDELINE_ADHERENCE:
+                return f"{secret_env_var_name_prefix}{OOTBType.GUIDELINE_ADHERENCE}_{llm_type_str}"
             else:
                 raise Exception("Invalid guard config for building env var name")
         else:

datarobot_dome/guards/model_guard.py

@@ -0,0 +1,84 @@
+#  ---------------------------------------------------------------------------------
+#  Copyright (c) 2026 DataRobot, Inc. and its affiliates. All rights reserved.
+#  Last updated 2025.
+#
+#  DataRobot, Inc. Confidential.
+#  This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+#  This file and its contents are subject to DataRobot Tool and Utility Agreement.
+#  For details, see
+#  https://www.datarobot.com/wp-content/uploads/2021/07/DataRobot-Tool-and-Utility-Agreement.pdf.
+#  ---------------------------------------------------------------------------------
+import datarobot as dr
+
+from datarobot_dome.constants import SPAN_PREFIX
+
+from .base import Guard
+
+
+class GuardModelInfo:
+    def __init__(self, model_config: dict):
+        self._model_id = model_config.get("model_id")
+        self._target_name = model_config["target_name"]
+        self._target_type = model_config["target_type"]
+        self._class_names = model_config.get("class_names", [])
+        self._input_column_name = model_config["input_column_name"]
+        self._replacement_text_column_name = model_config.get("replacement_text_column_name")
+
+    @property
+    def model_id(self) -> str:
+        return self._model_id
+
+    @property
+    def target_name(self) -> str:
+        return self._target_name
+
+    @property
+    def target_type(self) -> str:
+        return self._target_type
+
+    @property
+    def class_names(self) -> list[str]:
+        return self._class_names
+
+    @property
+    def input_column_name(self) -> str:
+        return self._input_column_name
+
+    @property
+    def replacement_text_column_name(self) -> str:
+        return self._replacement_text_column_name
+
+
+class ModelGuard(Guard):
+    def __init__(self, config: dict, stage=None):
+        super().__init__(config, stage)
+        self._deployment_id = config["deployment_id"]
+        self._model_info = GuardModelInfo(config["model_info"])
+        # dr.Client is set in the Pipeline init, Lets query the deployment
+        # to get the prediction server information
+        self.deployment = dr.Deployment.get(self._deployment_id)
+
+    @property
+    def deployment_id(self) -> str:
+        return self._deployment_id
+
+    @property
+    def model_info(self):
+        return self._model_info
+
+    def get_input_column_name(self, stage) -> str:
+        return self._model_info.input_column_name
+
+    def get_span_column_name(self, _):
+        if self.model_info is None:
+            raise NotImplementedError("Missing model_info for model guard")
+        # Typically 0th index is the target name
+        return self._model_info.target_name.split("_")[0]
+
+    def get_span_attribute_name(self, stage):
+        return f"{SPAN_PREFIX}.{stage.lower()}.{self.get_span_column_name(stage)}"
+
+    def has_average_score_custom_metric(self) -> bool:
+        """A couple ModelGuard types do not have an average score metric"""
+        return self.model_info.target_type not in ["Multiclass", "TextGeneration"]

datarobot_dome/guards/nemo_evaluator.py

@@ -0,0 +1,73 @@
+#  ---------------------------------------------------------------------------------
+#  Copyright (c) 2026 DataRobot, Inc. and its affiliates. All rights reserved.
+#  Last updated 2025.
+#
+#  DataRobot, Inc. Confidential.
+#  This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+#  This file and its contents are subject to DataRobot Tool and Utility Agreement.
+#  For details, see
+#  https://www.datarobot.com/wp-content/uploads/2021/07/DataRobot-Tool-and-Utility-Agreement.pdf.
+#  ---------------------------------------------------------------------------------
+from functools import cached_property
+
+from nemo_microservices import AsyncNeMoMicroservices
+
+from .base import Guard
+
+
+class NeMoEvaluatorGuard(Guard):
+    def __init__(self, config: dict, stage=None):
+        super().__init__(config, stage)
+        self.nemo_evaluator_type = config["nemo_evaluator_type"]
+        self._llm_type = config["llm_type"]
+        self.llm_deployment_id = config.get("deployment_id")
+
+    @cached_property
+    def _client(self) -> AsyncNeMoMicroservices:
+        """
+        Using localhost for development purpose only.
+        It will be replaced with url to a deployed NeMo evaluator instance later in the PBMP.
+        """
+        return AsyncNeMoMicroservices(base_url="http://localhost:8080")
+
+    def has_average_score_custom_metric(self) -> bool:
+        return False
+
+    async def evaluate(self, prompt: str, response: str) -> float:
+        raise NotImplementedError
+
+
+class NeMoLLMJudgeGuard(NeMoEvaluatorGuard):
+    def __init__(self, config: dict, stage=None):
+        super().__init__(config, stage)
+
+
+class NeMoContextRelevanceGuard(NeMoEvaluatorGuard):
+    def __init__(self, config: dict, stage=None):
+        super().__init__(config, stage)
+
+
+class NeMoResponseGroundednessGuard(NeMoEvaluatorGuard):
+    def __init__(self, config: dict, stage=None):
+        super().__init__(config, stage)
+
+
+class NeMoTopicAdherenceGuard(NeMoEvaluatorGuard):
+    def __init__(self, config: dict, stage=None):
+        super().__init__(config, stage)
+
+
+class NeMoAgentGoalAccuracyGuard(NeMoEvaluatorGuard):
+    def __init__(self, config: dict, stage=None):
+        super().__init__(config, stage)
+
+
+class NeMoResponseRelevancyGuard(NeMoEvaluatorGuard):
+    def __init__(self, config: dict, stage=None):
+        super().__init__(config, stage)
+
+
+class NeMoFaithfulnessGuard(NeMoEvaluatorGuard):
+    def __init__(self, config: dict, stage=None):
+        super().__init__(config, stage)

datarobot_dome/guards/nemo_guard.py

@@ -0,0 +1,146 @@
+#  ---------------------------------------------------------------------------------
+#  Copyright (c) 2026 DataRobot, Inc. and its affiliates. All rights reserved.
+#  Last updated 2025.
+#
+#  DataRobot, Inc. Confidential.
+#  This is proprietary source code of DataRobot, Inc. and its affiliates.
+#
+#  This file and its contents are subject to DataRobot Tool and Utility Agreement.
+#  For details, see
+#  https://www.datarobot.com/wp-content/uploads/2021/07/DataRobot-Tool-and-Utility-Agreement.pdf.
+#  ---------------------------------------------------------------------------------
+import logging
+import os
+
+import datarobot as dr
+from nemoguardrails import LLMRails
+from nemoguardrails import RailsConfig
+
+from datarobot_dome.constants import NEMO_GUARDRAILS_DIR
+from datarobot_dome.constants import GuardLLMType
+from datarobot_dome.constants import GuardOperatorType
+from datarobot_dome.guard_helpers import DEFAULT_OPEN_AI_API_VERSION
+from datarobot_dome.guard_helpers import get_azure_openai_client
+from datarobot_dome.guard_helpers import get_chat_nvidia_llm
+from datarobot_dome.guard_helpers import get_datarobot_endpoint_and_token
+from datarobot_dome.guard_helpers import get_llm_gateway_client
+from datarobot_dome.guard_helpers import use_llm_gateway_inference
+from datarobot_dome.guards.base import Guard
+from datarobot_dome.guards.guard_llm_mixin import GuardLLMMixin
+
+NEMO_THRESHOLD = "TRUE"
+
+
+class NeMoGuard(Guard, GuardLLMMixin):
+    def __init__(self, config: dict, stage=None, model_dir: str = os.getcwd()):
+        super().__init__(config, stage)
+        # NeMo guard only takes a boolean as threshold and equal to as comparator.
+        # Threshold bool == TRUE is defined in the colang file as the output of
+        # `bot should intervene`
+        if self.intervention:
+            if not self.intervention.threshold:
+                self.intervention.threshold = NEMO_THRESHOLD
+            if not self.intervention.comparator:
+                self.intervention.comparator = GuardOperatorType.EQUALS
+
+        # Default LLM Type for NeMo is set to OpenAI
+        self._llm_type = config.get("llm_type", GuardLLMType.OPENAI)
+        self.openai_api_base = config.get("openai_api_base")
+        self.openai_deployment_id = config.get("openai_deployment_id")
+        llm_id = None
+        credentials = None
+        use_llm_gateway = use_llm_gateway_inference(self._llm_type)
+        try:
+            self.openai_api_key = self.get_openai_api_key(config, self._llm_type)
+            if self._llm_type != GuardLLMType.NIM and self.openai_api_key is None:
+                raise ValueError("OpenAI API key is required for NeMo Guardrails")
+
+            if self.llm_type == GuardLLMType.OPENAI:
+                credentials = {
+                    "credential_type": "openai",
+                    "api_key": self.openai_api_key,
+                }
+                os.environ["OPENAI_API_KEY"] = self.openai_api_key
+                llm = None
+            elif self.llm_type == GuardLLMType.AZURE_OPENAI:
+                if self.openai_api_base is None:
+                    raise ValueError("Azure OpenAI API base url is required for LLM Guard")
+                if self.openai_deployment_id is None:
+                    raise ValueError("Azure OpenAI deployment ID is required for LLM Guard")
+                credentials = {
+                    "credential_type": "azure_openai",
+                    "api_base": self.openai_api_base,
+                    "api_version": DEFAULT_OPEN_AI_API_VERSION,
+                    "api_key": self.openai_api_key,
+                }
+                azure_openai_client = get_azure_openai_client(
+                    openai_api_key=self.openai_api_key,
+                    openai_api_base=self.openai_api_base,
+                    openai_deployment_id=self.openai_deployment_id,
+                )
+                llm = azure_openai_client
+            elif self.llm_type == GuardLLMType.GOOGLE:
+                # llm_id = config["google_model"]
+                raise NotImplementedError
+            elif self.llm_type == GuardLLMType.AMAZON:
+                # llm_id = config["aws_model"]
+                raise NotImplementedError
+            elif self.llm_type == GuardLLMType.DATAROBOT:
+                raise NotImplementedError
+            elif self.llm_type == GuardLLMType.NIM:
+                if config.get("deployment_id") is None:
+                    if self.openai_api_base is None:
+                        raise ValueError("NIM DataRobot deployment id is required for NIM LLM Type")
+                    else:
+                        logging.warning(
+                            "Using 'openai_api_base' is being deprecated and will be removed "
+                            "in the next release.  Please configure NIM DataRobot deployment "
+                            "using deployment_id"
+                        )
+                        if self.openai_api_key is None:
+                            raise ValueError("OpenAI API key is required for NeMo Guardrails")
+                else:
+                    self.deployment = dr.Deployment.get(self._deployment_id)
+                    datarobot_endpoint, self.openai_api_key = get_datarobot_endpoint_and_token()
+                    self.openai_api_base = (
+                        f"{datarobot_endpoint}/deployments/{str(self._deployment_id)}"
+                    )
+                llm = get_chat_nvidia_llm(
+                    api_key=self.openai_api_key,
+                    base_url=self.openai_api_base,
+                )
+            else:
+                raise ValueError(f"Invalid LLMType: {self.llm_type}")
+
+        except Exception as e:
+            # no valid user credentials provided, raise if not using LLM Gateway
+            credentials = None
+            if not use_llm_gateway:
+                raise e
+
+        if use_llm_gateway:
+            # Currently only OPENAI and AZURE_OPENAI are supported by NeMoGuard
+            # For Bedrock and Vertex the model in the config is actually the LLM ID
+            # For OpenAI we use the default model defined in get_llm_gateway_client
+            # For Azure we use the deployment ID
+            llm = get_llm_gateway_client(
+                llm_id=llm_id,
+                openai_deployment_id=self.openai_deployment_id,
+                credentials=credentials,
+            )
+
+        # Use guard stage to determine whether to read from prompt/response subdirectory
+        # for nemo configurations.  "nemo_guardrails" folder is at same level of custom.py
+        # So, the config path becomes model_dir + "nemo_guardrails"
+        nemo_config_path = os.path.join(model_dir, NEMO_GUARDRAILS_DIR)
+        self.nemo_rails_config_path = os.path.join(nemo_config_path, self.stage)
+        nemo_rails_config = RailsConfig.from_path(config_path=self.nemo_rails_config_path)
+        self._nemo_llm_rails = LLMRails(nemo_rails_config, llm=llm)
+
+    def has_average_score_custom_metric(self) -> bool:
+        """No average score metrics for NemoGuard's"""
+        return False
+
+    @property
+    def nemo_llm_rails(self):
+        return self._nemo_llm_rails