datarobot-moderations 11.1.12__py3-none-any.whl

datarobot_dome/guards/guard_llm_mixin.py

@@ -0,0 +1,232 @@
+#  ---------------------------------------------------------------------------------
+#  Copyright (c) 2025 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 json
+import os
+
+import datarobot as dr
+import trafaret as t
+
+from datarobot_dome.constants import AWS_ACCOUNT_SECRET_DEFINITION_SUFFIX
+from datarobot_dome.constants import GOOGLE_SERVICE_ACCOUNT_SECRET_DEFINITION_SUFFIX
+from datarobot_dome.constants import OPENAI_SECRET_DEFINITION_SUFFIX
+from datarobot_dome.constants import SECRET_DEFINITION_PREFIX
+from datarobot_dome.constants import GuardLLMType
+from datarobot_dome.constants import GuardType
+from datarobot_dome.constants import OOTBType
+from datarobot_dome.guard_helpers import get_azure_openai_client
+from datarobot_dome.guard_helpers import get_bedrock_client
+from datarobot_dome.guard_helpers import get_datarobot_llm
+from datarobot_dome.guard_helpers import get_vertex_client
+from datarobot_dome.guard_helpers import try_to_fallback_to_llm_gateway
+
+basic_credential_trafaret = t.Dict(
+    {
+        t.Key("credentialType", to_name="credential_type", optional=False): t.Enum("basic"),
+        t.Key("password", to_name="password", optional=False): t.String,
+    },
+    allow_extra=["*"],
+)
+
+api_token_credential_trafaret = t.Dict(
+    {
+        t.Key("credentialType", to_name="credential_type", optional=False): t.Enum("api_token"),
+        t.Key("apiToken", to_name="api_token", optional=False): t.String,
+    },
+    allow_extra=["*"],
+)
+
+google_service_account_trafaret = t.Dict(
+    {
+        t.Key("credentialType", to_name="credential_type", optional=False): t.Enum("gcp"),
+        t.Key("gcpKey", to_name="gcp_key", optional=False): t.Dict(allow_extra=["*"]),
+    },
+    allow_extra=["*"],
+)
+
+aws_account_trafaret = t.Dict(
+    {
+        t.Key("credentialType", to_name="credential_type", optional=False): t.Enum("s3"),
+        t.Key("awsAccessKeyId", to_name="aws_access_key_id", optional=False): t.String,
+        t.Key("awsSecretAccessKey", to_name="aws_secret_access_key", optional=False): t.String,
+        t.Key("awsSessionToken", to_name="aws_session_token", optional=True, default=None): t.String
+        | t.Null,
+    },
+    allow_extra=["*"],
+)
+
+
+credential_trafaret = t.Dict(
+    {
+        t.Key("type", optional=False): t.Enum("credential"),
+        t.Key("payload", optional=False): t.Or(
+            basic_credential_trafaret,
+            api_token_credential_trafaret,
+            google_service_account_trafaret,
+            aws_account_trafaret,
+        ),
+    }
+)
+
+
+class GuardLLMMixin:
+    def get_secret_env_var_base(self, config, llm_type_str):
+        guard_type = config["type"]
+        guard_stage = config["stage"]
+        secret_env_var_name_prefix = f"{SECRET_DEFINITION_PREFIX}_{guard_type}_{guard_stage}_"
+        if guard_type == GuardType.NEMO_GUARDRAILS:
+            return f"{secret_env_var_name_prefix}{llm_type_str}"
+        elif guard_type == GuardType.OOTB:
+            if config["ootb_type"] == OOTBType.FAITHFULNESS:
+                return f"{secret_env_var_name_prefix}{OOTBType.FAITHFULNESS}_{llm_type_str}"
+            elif config["ootb_type"] == OOTBType.AGENT_GOAL_ACCURACY:
+                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}"
+            else:
+                raise Exception("Invalid guard config for building env var name")
+        else:
+            raise Exception("Invalid guard config for building env var name")
+
+    def build_open_ai_api_key_env_var_name(self, config, llm_type):
+        llm_type_str = ""
+        if llm_type == GuardLLMType.AZURE_OPENAI:
+            llm_type_str = "AZURE_"
+        elif llm_type == GuardLLMType.NIM:
+            llm_type_str = "NIM_"
+        var_name = self.get_secret_env_var_base(config, llm_type_str)
+        var_name += OPENAI_SECRET_DEFINITION_SUFFIX
+        return var_name.upper()
+
+    def get_openai_api_key(self, config, llm_type):
+        api_key_env_var_name = self.build_open_ai_api_key_env_var_name(config, llm_type)
+        if api_key_env_var_name not in os.environ:
+            if llm_type == GuardLLMType.NIM:
+                return None
+            raise Exception(f"Expected environment variable '{api_key_env_var_name}' not found")
+
+        env_var_value = json.loads(os.environ[api_key_env_var_name])
+        credential_config = credential_trafaret.check(env_var_value)
+        if credential_config["payload"]["credential_type"] == "basic":
+            return credential_config["payload"]["password"]
+        else:
+            return credential_config["payload"]["api_token"]
+
+    def get_google_service_account(self, config):
+        service_account_env_var_name = self.get_secret_env_var_base(
+            config, GOOGLE_SERVICE_ACCOUNT_SECRET_DEFINITION_SUFFIX
+        ).upper()
+        if service_account_env_var_name not in os.environ:
+            raise Exception(
+                f"Expected environment variable '{service_account_env_var_name}' not found"
+            )
+
+        env_var_value = json.loads(os.environ[service_account_env_var_name])
+        credential_config = credential_trafaret.check(env_var_value)
+        if credential_config["payload"]["credential_type"] == "gcp":
+            return credential_config["payload"]["gcp_key"]
+        else:
+            raise Exception("Google model requires a credential of type 'gcp'")
+
+    def get_aws_account(self, config):
+        service_account_env_var_name = self.get_secret_env_var_base(
+            config, AWS_ACCOUNT_SECRET_DEFINITION_SUFFIX
+        ).upper()
+        if service_account_env_var_name not in os.environ:
+            raise Exception(
+                f"Expected environment variable '{service_account_env_var_name}' not found"
+            )
+
+        env_var_value = json.loads(os.environ[service_account_env_var_name])
+        credential_config = credential_trafaret.check(env_var_value)
+        if credential_config["payload"]["credential_type"] == "s3":
+            return credential_config["payload"]
+        else:
+            raise Exception("Amazon model requires a credential of type 's3'")
+
+    def get_llm(self, config, llm_type):
+        openai_api_base = config.get("openai_api_base")
+        openai_deployment_id = config.get("openai_deployment_id")
+        llm_id = None
+        try:
+            if llm_type in [GuardLLMType.OPENAI, GuardLLMType.AZURE_OPENAI]:
+                openai_api_key = self.get_openai_api_key(config, llm_type)
+                if openai_api_key is None:
+                    raise ValueError("OpenAI API key is required for Faithfulness guard")
+
+                if llm_type == GuardLLMType.OPENAI:
+                    os.environ["OPENAI_API_KEY"] = openai_api_key
+                    llm = "default"
+                elif llm_type == GuardLLMType.AZURE_OPENAI:
+                    if openai_api_base is None:
+                        raise ValueError("OpenAI API base url is required for LLM Guard")
+                    if openai_deployment_id is None:
+                        raise ValueError("OpenAI deployment ID is required for LLM Guard")
+                    azure_openai_client = get_azure_openai_client(
+                        openai_api_key=openai_api_key,
+                        openai_api_base=openai_api_base,
+                        openai_deployment_id=openai_deployment_id,
+                    )
+                    llm = azure_openai_client
+            elif llm_type == GuardLLMType.GOOGLE:
+                llm_id = config["google_model"]
+                if llm_id is None:
+                    raise ValueError("Google model is required for LLM Guard")
+                if config.get("google_region") is None:
+                    raise ValueError("Google region is required for LLM Guard")
+                llm = get_vertex_client(
+                    google_model=llm_id,
+                    google_service_account=self.get_google_service_account(config),
+                    google_region=config["google_region"],
+                )
+            elif llm_type == GuardLLMType.AMAZON:
+                llm_id = config["aws_model"]
+                if llm_id is None:
+                    raise ValueError("AWS model is required for LLM Guard")
+                if config.get("aws_region") is None:
+                    raise ValueError("AWS region is required for LLM Guard")
+                credential_config = self.get_aws_account(config)
+                llm = get_bedrock_client(
+                    aws_model=llm_id,
+                    aws_access_key_id=credential_config["aws_access_key_id"],
+                    aws_secret_access_key=credential_config["aws_secret_access_key"],
+                    aws_session_token=credential_config["aws_session_token"],
+                    aws_region=config["aws_region"],
+                )
+            elif llm_type == GuardLLMType.DATAROBOT:
+                if config["type"] == GuardType.OOTB and config["ootb_type"] in [
+                    OOTBType.AGENT_GOAL_ACCURACY,
+                    OOTBType.TASK_ADHERENCE,
+                ]:
+                    # DataRobot LLM does not implement generate / agenerate yet
+                    # so can't support it for Agent Goal Accuracy
+                    raise NotImplementedError
+                if config.get("deployment_id") is None:
+                    raise ValueError("Deployment ID is required for LLM Guard")
+                deployment = dr.Deployment.get(config["deployment_id"])
+                llm = get_datarobot_llm(deployment)
+            elif llm_type == GuardLLMType.NIM:
+                raise NotImplementedError
+            else:
+                raise ValueError(f"Invalid LLMType: {llm_type}")
+
+        except Exception as e:
+            llm = try_to_fallback_to_llm_gateway(
+                # 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_id=llm_id,
+                openai_deployment_id=openai_deployment_id,
+                llm_type=llm_type,
+                e=e,
+            )
+
+        return llm

datarobot_dome/llm.py

@@ -0,0 +1,148 @@
+#  ---------------------------------------------------------------------------------
+#  Copyright (c) 2025 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 typing import Any
+
+import datarobot as dr
+import pandas as pd
+from datarobot_predict.deployment import predict
+from llama_index.core.base.llms.types import CompletionResponse
+from llama_index.core.base.llms.types import LLMMetadata
+from llama_index.core.bridge.pydantic import PrivateAttr
+from llama_index.core.llms.callbacks import llm_chat_callback
+from llama_index.core.llms.callbacks import llm_completion_callback
+from llama_index.core.llms.llm import LLM
+
+from datarobot_dome.async_http_client import AsyncHTTPClient
+
+DEFAULT_TEMPERATURE = 1.0
+MAX_TOKENS = 512
+DEFAULT_TIMEOUT = 30
+MAX_RETRIES = 5
+
+
+class DataRobotLLM(LLM):
+    # DataRobot deployment object.  Only one of `deployment` or `deployment_id` is required
+    _deployment: dr.Deployment = PrivateAttr()
+    # DataRobot endpoint URL, Only one of `dr_client` or the pair
+    # (datarobot_endpoint, datarobot_api_token) is required
+    _datarobot_endpoint: str = PrivateAttr()
+    # DataRobot API Token to use, Only one of `dr_client` or the pair
+    # (datarobot_endpoint, datarobot_api_token) is required
+    _datarobot_api_token: str = PrivateAttr()
+    # Async HTTP Client for all async prediction requests with DataRobot Deployment
+    _async_http_client: Any = PrivateAttr()
+
+    _prompt_column_name: str = PrivateAttr()
+    _target_column_name: str = PrivateAttr()
+
+    def __init__(
+        self,
+        deployment,
+        datarobot_endpoint=None,
+        datarobot_api_token=None,
+        callback_manager=None,
+    ):
+        super().__init__(
+            model="DataRobot LLM",
+            temperature=DEFAULT_TEMPERATURE,
+            max_tokens=MAX_TOKENS,
+            timeout=DEFAULT_TIMEOUT,
+            max_retries=MAX_RETRIES,
+            callback_manager=callback_manager,
+        )
+        if deployment is None:
+            raise ValueError("DataRobot deployment is required")
+
+        if datarobot_api_token is None and datarobot_endpoint is None:
+            raise ValueError(
+                "Connection parameters 'datarobot_endpoint' and 'datarobot_api_token' "
+                "needs to be provided"
+            )
+        self._deployment = deployment
+        self._datarobot_endpoint = datarobot_endpoint
+        self._datarobot_api_token = datarobot_api_token
+
+        if self._deployment.model["target_type"] != "TextGeneration":
+            raise ValueError(
+                f"Invalid deployment '{self._deployment.label}' for LLM.  Expecting an LLM "
+                f"deployment, but is a '{self._deployment.model['target_type']}' deployment"
+            )
+
+        self._prompt_column_name = self._deployment.model.get("prompt")
+        if self._prompt_column_name is None:
+            raise ValueError("Prompt column name 'prompt' is not set on the deployment / model")
+
+        self._target_column_name = self._deployment.model["target_name"] + "_PREDICTION"
+        self._async_http_client = AsyncHTTPClient(DEFAULT_TIMEOUT)
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return "datarobot-llm"
+
+    @property
+    def _identifying_params(self):
+        """Get the identifying parameters."""
+        _model_kwargs = self.model_kwargs or {}
+        return {
+            **{"endpoint_url": self._datarobot_endpoint, "deployment_id": str(self._deployment.id)},
+            **{"model_kwargs": _model_kwargs},
+        }
+
+    def _call(self, prompt, stop=None, run_manager=None, **kwargs):
+        df = pd.DataFrame({self._prompt_column_name: [prompt]})
+        result_df, _ = predict(self._deployment, df)
+        return result_df[self._target_column_name].iloc[0]
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "DataRobotLLM"
+
+    @property
+    def metadata(self):
+        return LLMMetadata(is_chat_model=False)
+
+    @llm_chat_callback()
+    def chat(self, messages, **kwargs: Any):
+        raise NotImplementedError
+
+    @llm_completion_callback()
+    def complete(self, prompt, formatted, **kwargs):
+        df = pd.DataFrame({self._prompt_column_name: [prompt]})
+        result_df, _ = predict(self._deployment, df)
+        return CompletionResponse(text=result_df[self._target_column_name].iloc[0], raw={})
+
+    @llm_chat_callback()
+    def stream_chat(self, messages, **kwargs):
+        raise NotImplementedError
+
+    @llm_completion_callback()
+    def stream_complete(self, prompt, formatted=False, **kwargs):
+        raise NotImplementedError
+
+    @llm_chat_callback()
+    async def achat(self, messages, **kwargs):
+        raise NotImplementedError
+
+    @llm_completion_callback()
+    async def acomplete(self, prompt, formatted=False, **kwargs):
+        input_df_to_guard = pd.DataFrame({self._prompt_column_name: [prompt]})
+        result_df = await self._async_http_client.predict(self._deployment, input_df_to_guard)
+        return CompletionResponse(text=result_df[self._target_column_name].iloc[0], raw={})
+
+    @llm_chat_callback()
+    async def astream_chat(self, messages, **kwargs):
+        raise NotImplementedError
+
+    @llm_completion_callback()
+    async def astream_complete(self, prompt, formatted=False, **kwargs):
+        raise NotImplementedError

datarobot_dome/metrics/__init__.py

@@ -0,0 +1,11 @@
+#  ---------------------------------------------------------------------------------
+#  Copyright (c) 2025 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.
+#  ---------------------------------------------------------------------------------

datarobot_dome/metrics/citation_metrics.py

@@ -0,0 +1,98 @@
+#  ---------------------------------------------------------------------------------
+#  Copyright (c) 2025 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 pandas as pd
+from datarobot.enums import CustomMetricAggregationType
+from datarobot.enums import CustomMetricDirectionality
+
+from datarobot_dome.constants import CUSTOM_METRIC_DESCRIPTION_SUFFIX
+from datarobot_dome.guard_helpers import get_token_count
+from datarobot_dome.metrics.metric_scorer import MetricScorer
+
+CITATION_COLUMN = "response.citations"
+
+
+class CitationTokenCountScorer(MetricScorer):
+    NAME = "Total Citation Tokens"
+    DESCRIPTION = f"Total number of citation tokens. {CUSTOM_METRIC_DESCRIPTION_SUFFIX}"
+    DIRECTIONALITY = CustomMetricDirectionality.LOWER_IS_BETTER
+    UNITS = "count"
+    AGGREGATION_TYPE = CustomMetricAggregationType.SUM
+    BASELINE_VALUE = 0
+    INPUT_COLUMN = CITATION_COLUMN
+
+    def score(self, df: pd.DataFrame) -> float:
+        column = self.input_column
+        if column not in df.columns:
+            return 0.0
+
+        return sum(
+            sum(get_token_count(v, self.encoding) for v in cell) for cell in df[column].values
+        )
+
+
+class CitationTokenAverageScorer(MetricScorer):
+    NAME = "Average Citation Tokens"
+    DESCRIPTION = f"Average number of citation tokens. {CUSTOM_METRIC_DESCRIPTION_SUFFIX}"
+    DIRECTIONALITY = CustomMetricDirectionality.LOWER_IS_BETTER
+    UNITS = "count"
+    AGGREGATION_TYPE = CustomMetricAggregationType.AVERAGE
+    BASELINE_VALUE = 0
+    INPUT_COLUMN = CITATION_COLUMN
+
+    def score(self, df: pd.DataFrame) -> float:
+        average = 0.0
+        total = 0
+        count = 0
+        column = self.input_column
+        if column not in df.columns:
+            return 0.0
+
+        for cell in df[column].values:
+            total += sum(get_token_count(v, self.encoding) for v in cell)
+            count += sum(v != "" for v in cell)
+            average = total / count
+
+        return average
+
+
+class DocumentCountScorer(MetricScorer):
+    NAME = "Total Documents"
+    DESCRIPTION = f"Total number of documents. {CUSTOM_METRIC_DESCRIPTION_SUFFIX}"
+    DIRECTIONALITY = CustomMetricDirectionality.LOWER_IS_BETTER
+    UNITS = "count"
+    AGGREGATION_TYPE = CustomMetricAggregationType.SUM
+    BASELINE_VALUE = 0
+    INPUT_COLUMN = CITATION_COLUMN
+
+    def score(self, df: pd.DataFrame) -> float:
+        column = self.input_column
+        if column not in df.columns:
+            return 0.0
+
+        return sum(sum(bool(v) for v in cell) for cell in df[column].values)
+
+
+class DocumentAverageScorer(MetricScorer):
+    NAME = "Average Documents"
+    DESCRIPTION = f"Average number of documents. {CUSTOM_METRIC_DESCRIPTION_SUFFIX}"
+    DIRECTIONALITY = CustomMetricDirectionality.LOWER_IS_BETTER
+    UNITS = "count"
+    AGGREGATION_TYPE = CustomMetricAggregationType.AVERAGE
+    BASELINE_VALUE = 0
+    INPUT_COLUMN = CITATION_COLUMN
+
+    def score(self, df: pd.DataFrame) -> float:
+        column = self.input_column
+        if column not in df.columns:
+            return 0.0
+
+        return sum(sum(bool(v) for v in cell) for cell in df[column].values)

datarobot_dome/metrics/factory.py

@@ -0,0 +1,52 @@
+#  ---------------------------------------------------------------------------------
+#  Copyright (c) 2025 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 typing import Any
+from typing import ClassVar
+from typing import Optional
+
+from datarobot_dome.metrics.citation_metrics import CitationTokenAverageScorer
+from datarobot_dome.metrics.citation_metrics import CitationTokenCountScorer
+from datarobot_dome.metrics.citation_metrics import DocumentAverageScorer
+from datarobot_dome.metrics.citation_metrics import DocumentCountScorer
+from datarobot_dome.metrics.metric_scorer import MetricScorer
+from datarobot_dome.metrics.metric_scorer import ScorerType
+
+METRIC_SCORE_CLASS_MAP: dict[ScorerType, ClassVar] = {
+    ScorerType.CITATION_TOKEN_AVERAGE: CitationTokenAverageScorer,
+    ScorerType.CITATION_TOKEN_COUNT: CitationTokenCountScorer,
+    ScorerType.DOCUMENT_AVERAGE: DocumentAverageScorer,
+    ScorerType.DOCUMENT_COUNT: DocumentCountScorer,
+}
+
+
+class MetricScorerFactory:
+    @staticmethod
+    def get_class(metric_type: ScorerType) -> ClassVar:
+        clazz = METRIC_SCORE_CLASS_MAP.get(metric_type)
+        if clazz is None:
+            raise ValueError(f"Unknown metric type: {metric_type}")
+
+        return clazz
+
+    @staticmethod
+    def create(metric_type: ScorerType, config: Optional[dict[str, Any]] = None) -> MetricScorer:
+        _config = config or {}
+        clazz = MetricScorerFactory.get_class(metric_type)
+        return clazz(_config)
+
+    @staticmethod
+    def custom_metric_config(
+        metric_type: ScorerType, config: Optional[dict[str, Any]] = None
+    ) -> dict[str, Any]:
+        _config = config or {}
+        clazz = MetricScorerFactory.get_class(metric_type)
+        return clazz.custom_metric_definition(_config)

datarobot_dome/metrics/metric_scorer.py

@@ -0,0 +1,78 @@
+#  ---------------------------------------------------------------------------------
+#  Copyright (c) 2025 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 abc import abstractmethod
+from enum import Enum
+from typing import Any
+
+import pandas as pd
+
+
+class ScorerType(str, Enum):
+    CITATION_TOKEN_AVERAGE = "CITATION_TOKEN_AVERAGE"
+    CITATION_TOKEN_COUNT = "CITATION_TOKEN_COUNT"
+    DOCUMENT_AVERAGE = "DOCUMENT_AVERAGE"
+    DOCUMENT_COUNT = "DOCUMENT_COUNT"
+
+
+class MetricScorer(ABC):
+    IS_MODEL_SPECIFIC = True
+    TIME_STEP = "hour"
+
+    def __init__(
+        self,
+        config: dict[str, Any],
+    ):
+        self.config = config
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
+
+    @classmethod
+    def custom_metric_definition(cls, config: dict[str, Any]) -> dict[str, Any]:
+        """
+        Generates a custom-metric configuration/definition that is used to create the
+        custom-metric in the DR application.
+
+        This is done as a class method because we create custom-metrics before creating
+        the scorer.
+        """
+        return {
+            "name": config.get("name", cls.NAME),
+            "directionality": cls.DIRECTIONALITY,
+            "units": cls.UNITS,
+            "type": cls.AGGREGATION_TYPE,
+            "baselineValue": config.get("baseline_value", cls.BASELINE_VALUE),
+            "isModelSpecific": config.get("is_model_specific", cls.IS_MODEL_SPECIFIC),
+            "timeStep": cls.TIME_STEP,
+            "description": config.get("description", cls.DESCRIPTION),
+        }
+
+    @property
+    def name(self) -> str:
+        return self.config.get("name", self.NAME)
+
+    @property
+    def per_prediction(self) -> bool:
+        return self.config.get("per_prediction", False)
+
+    @property
+    def input_column(self) -> str:
+        return self.config.get("input_column", self.INPUT_COLUMN)
+
+    @property
+    def encoding(self) -> str:
+        return self.config.get("encoding", "cl100k_base")
+
+    @abstractmethod
+    def score(self, df: pd.DataFrame) -> float:
+        pass

datarobot_dome/pipeline/__init__.py

@@ -0,0 +1,11 @@
+#  ---------------------------------------------------------------------------------
+#  Copyright (c) 2025 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.
+#  ---------------------------------------------------------------------------------