datarobot-moderations 11.1.12__py3-none-any.whl

datarobot_dome/__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/async_http_client.py

@@ -0,0 +1,248 @@
+#  ---------------------------------------------------------------------------------
+#  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 asyncio
+import atexit
+import datetime
+import logging
+import os
+import traceback
+from http import HTTPStatus
+from io import StringIO
+
+import aiohttp
+import backoff
+import nest_asyncio
+import pandas as pd
+
+from datarobot_dome.constants import DATAROBOT_ACTUAL_ON_PREM_ST_SAAS_URL
+from datarobot_dome.constants import DATAROBOT_CONFIGURED_ON_PREM_ST_SAAS_URL
+from datarobot_dome.constants import DATAROBOT_SERVERLESS_PLATFORM
+from datarobot_dome.constants import DEFAULT_GUARD_PREDICTION_TIMEOUT_IN_SEC
+from datarobot_dome.constants import LOGGER_NAME_PREFIX
+from datarobot_dome.constants import RETRY_COUNT
+from datarobot_dome.constants import ModerationEventTypes
+
+RETRY_STATUS_CODES = [
+    HTTPStatus.REQUEST_ENTITY_TOO_LARGE,
+    HTTPStatus.BAD_GATEWAY,
+    HTTPStatus.GATEWAY_TIMEOUT,
+]
+RETRY_AFTER_STATUS_CODES = [HTTPStatus.TOO_MANY_REQUESTS, HTTPStatus.SERVICE_UNAVAILABLE]
+
+
+# We want this logger to be available for backoff too, hence defining outside the class
+logger = logging.getLogger(LOGGER_NAME_PREFIX + ".AsyncHTTPClient")
+
+
+# Event handlers for backoff
+def _timeout_backoff_handler(details):
+    logger.warning(
+        f"HTTP Timeout: Backing off {details['wait']} seconds after {details['tries']} tries"
+    )
+
+
+def _timeout_giveup_handler(details):
+    url = details["args"][1]
+    logger.error(f"Giving up predicting on {url}, Retried {details['tries']} after HTTP Timeout")
+
+
+def _retry_backoff_handler(details):
+    status_code = details["value"].status
+    message = details["value"].reason
+    retry_after_value = details["value"].headers.get("Retry-After")
+    logger.warning(
+        f"Received status code {status_code}, message {message},"
+        f" Retry-After val: {retry_after_value} "
+        f"Backing off {details['wait']} seconds after {details['tries']} tries"
+    )
+
+
+def _retry_giveup_handler(details):
+    message = (
+        f"Giving up predicting on {details['args'][1]}, Retried {details['tries']} retries, "
+        f"elapsed time {details['elapsed']} sec, but couldn't get predictions"
+    )
+    raise Exception(message)
+
+
+class AsyncHTTPClient:
+    def __init__(self, timeout=DEFAULT_GUARD_PREDICTION_TIMEOUT_IN_SEC):
+        self._logger = logging.getLogger(LOGGER_NAME_PREFIX + "." + self.__class__.__name__)
+        self.csv_headers = {
+            "Content-Type": "text/csv",
+            "Accept": "text/csv",
+            "Authorization": f"Bearer {os.environ['DATAROBOT_API_TOKEN']}",
+        }
+        self.json_headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+            "Authorization": f"Bearer {os.environ['DATAROBOT_API_TOKEN']}",
+        }
+        self.session = None
+        self.events_url = f"{os.environ['DATAROBOT_ENDPOINT']}/remoteEvents/"
+
+        try:
+            self.loop = asyncio.get_event_loop()
+        except RuntimeError as e:
+            if str(e).startswith("There is no current event loop in thread") or str(e).startswith(
+                "Event loop is closed"
+            ):
+                self.loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(self.loop)
+            else:
+                raise
+        self.loop.run_until_complete(self.__create_client_session(timeout))
+        self.loop.set_debug(True)
+        nest_asyncio.apply(loop=self.loop)
+
+        atexit.register(self.shutdown)
+
+    async def __create_client_session(self, timeout):
+        client_timeout = aiohttp.ClientTimeout(
+            connect=timeout, sock_connect=timeout, sock_read=timeout
+        )
+        # Creation of client session needs to happen within in async function
+        self.session = aiohttp.ClientSession(timeout=client_timeout)
+
+    def shutdown(self):
+        asyncio.run(self.session.close())
+
+    @staticmethod
+    def is_serverless_deployment(deployment):
+        if not deployment.prediction_environment:
+            return False
+
+        if deployment.prediction_environment.get("platform") == DATAROBOT_SERVERLESS_PLATFORM:
+            return True
+
+        return False
+
+    @staticmethod
+    def is_on_prem_st_saas_endpoint():
+        return os.environ.get("DATAROBOT_ENDPOINT").startswith(
+            DATAROBOT_CONFIGURED_ON_PREM_ST_SAAS_URL
+        )
+
+    @backoff.on_predicate(
+        backoff.runtime,
+        predicate=lambda r: r.status in RETRY_AFTER_STATUS_CODES,
+        value=lambda r: int(r.headers.get("Retry-After", DEFAULT_GUARD_PREDICTION_TIMEOUT_IN_SEC)),
+        max_time=RETRY_COUNT * DEFAULT_GUARD_PREDICTION_TIMEOUT_IN_SEC,
+        max_tries=RETRY_COUNT,
+        jitter=None,
+        logger=logger,
+        on_backoff=_retry_backoff_handler,
+        on_giveup=_retry_giveup_handler,
+    )
+    @backoff.on_predicate(
+        backoff.fibo,
+        predicate=lambda r: r.status in RETRY_STATUS_CODES,
+        jitter=None,
+        max_tries=RETRY_COUNT,
+        max_time=RETRY_COUNT * DEFAULT_GUARD_PREDICTION_TIMEOUT_IN_SEC,
+        logger=logger,
+        on_backoff=_retry_backoff_handler,
+        on_giveup=_retry_giveup_handler,
+    )
+    @backoff.on_exception(
+        backoff.fibo,
+        asyncio.TimeoutError,
+        max_tries=RETRY_COUNT,
+        max_time=RETRY_COUNT * DEFAULT_GUARD_PREDICTION_TIMEOUT_IN_SEC,
+        logger=logger,
+        on_backoff=_timeout_backoff_handler,
+        on_giveup=_timeout_giveup_handler,
+        raise_on_giveup=True,
+    )
+    async def post_predict_request(self, url_path, input_df):
+        return await self.session.post(
+            url_path, data=input_df.to_csv(index=False), headers=self.csv_headers
+        )
+
+    async def predict(self, deployment, input_df):
+        deployment_id = str(deployment.id)
+        if self.is_on_prem_st_saas_endpoint():
+            url_path = DATAROBOT_ACTUAL_ON_PREM_ST_SAAS_URL
+        elif self.is_serverless_deployment(deployment):
+            url_path = f"{os.environ['DATAROBOT_ENDPOINT']}"
+        else:
+            prediction_server = deployment.default_prediction_server
+            if not prediction_server:
+                raise ValueError(
+                    "Can't make prediction request because Deployment object doesn't contain "
+                    "default prediction server"
+                )
+            datarobot_key = prediction_server.get("datarobot-key")
+            if datarobot_key:
+                self.csv_headers["datarobot-key"] = datarobot_key
+
+            url_path = f"{prediction_server['url']}/predApi/v1.0"
+
+        url_path += f"/deployments/{deployment_id}/predictions"
+        response = await self.post_predict_request(url_path, input_df)
+        if not response.ok:
+            raise Exception(
+                f"Failed to get guard predictions: {response.reason} Status: {response.status}"
+            )
+        csv_data = await response.text()
+        return pd.read_csv(StringIO(csv_data))
+
+    async def async_report_event(
+        self, title, message, event_type, deployment_id, guard_name=None, metric_name=None
+    ):
+        payload = {
+            "title": title,
+            "message": message,
+            "timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat(),
+            "deploymentId": str(deployment_id),
+            "eventType": event_type,
+            "moderationData": {"guardName": "", "metricName": ""},
+        }
+        error_text = ""
+        if metric_name:
+            payload["moderationData"]["metricName"] = metric_name
+            error_text = f"for metric {metric_name}"
+        if guard_name:
+            payload["moderationData"]["guardName"] = guard_name
+            error_text = f"for guard {guard_name}"
+
+        response = await self.session.post(self.events_url, json=payload, headers=self.json_headers)
+        if response.status != HTTPStatus.CREATED:
+            # Lets not raise - we just failed to report an event, let the moderation
+            # continue
+            logger.error(
+                f"Failed to post event {event_type} {error_text} "
+                f" Status: {response.status} Message: {response.reason}"
+            )
+
+    async def bulk_upload_custom_metrics(self, url, payload, deployment_id):
+        self._logger.debug("Uploading custom metrics")
+        try:
+            response = await self.session.post(url, json=payload, headers=self.json_headers)
+            if response.status != HTTPStatus.ACCEPTED:
+                raise Exception(
+                    f"Error uploading custom metrics: Status Code: {response.status_code}"
+                    f"Message: {response.text}"
+                )
+            self._logger.info("Successfully uploaded custom metrics")
+        except Exception as e:
+            title = "Failed to upload custom metrics"
+            message = f"Exception: {e} Payload: {payload}"
+            self._logger.error(title + " " + message)
+            self._logger.error(traceback.format_exc())
+            await self.async_report_event(
+                title,
+                message,
+                ModerationEventTypes.MODERATION_METRIC_REPORTING_ERROR,
+                deployment_id,
+            )
+            # Lets not raise the exception, just walk off

datarobot_dome/chat_helper.py

@@ -0,0 +1,227 @@
+#  ---------------------------------------------------------------------------------
+#  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 logging
+import time
+import traceback
+from re import match
+
+import tiktoken
+
+from datarobot_dome.constants import AGENTIC_PIPELINE_INTERACTIONS_ATTR
+from datarobot_dome.constants import NONE_CUSTOM_PY_RESPONSE
+from datarobot_dome.constants import PROMPT_TOKEN_COUNT_COLUMN_NAME_FROM_USAGE
+from datarobot_dome.constants import RESPONSE_TOKEN_COUNT_COLUMN_NAME_FROM_USAGE
+from datarobot_dome.constants import GuardStage
+from datarobot_dome.guard_executor import AsyncGuardExecutor
+from datarobot_dome.guard_helpers import calculate_token_counts_for_cost_calculations
+from datarobot_dome.guard_helpers import get_citation_columns
+from datarobot_dome.guard_helpers import get_rouge_1_score
+
+_logger = logging.getLogger("chat_helper")
+
+
+def get_all_citation_columns(df):
+    citation_columns = []
+    for pattern in [
+        "CITATION_CONTENT_",
+        "CITATION_SOURCE_",
+        "CITATION_PAGE_",
+        "CITATION_CHUNK_ID_",
+        "CITATION_START_INDEX_",
+        "CITATION_SIMILARITY_SCORE_",
+    ]:
+        citation_columns.extend(list(filter(lambda column: match(pattern, column), df.columns)))
+    return citation_columns
+
+
+def build_moderations_attribute_for_completion(pipeline, df):
+    """
+    Given the dataframe build a moderation attribute to be returned with
+    chat completion or chat completion chunk
+    """
+    if df is None or df.empty:
+        return None
+
+    prompt_column_name = pipeline.get_input_column(GuardStage.PROMPT)
+    blocked_message_prompt_column_name = f"blocked_message_{prompt_column_name}"
+    response_column_name = pipeline.get_input_column(GuardStage.RESPONSE)
+    replaced_message_prompt_column_name = f"replaced_message_{prompt_column_name}"
+    blocked_message_completion_column_name = f"blocked_message_{response_column_name}"
+    replaced_message_response_column_name = f"replaced_message_{response_column_name}"
+
+    moderations = df.to_dict(orient="records")[0]
+    columns_to_drop = [
+        pipeline.get_input_column(GuardStage.PROMPT),
+        # Its already copied as part of the completion.choices[0].message.content
+        pipeline.get_input_column(GuardStage.RESPONSE),
+        blocked_message_prompt_column_name,
+        blocked_message_completion_column_name,
+        replaced_message_prompt_column_name,
+        replaced_message_response_column_name,
+        f"Noneed_{prompt_column_name}",
+        f"Noneed_{response_column_name}",
+    ]
+    citation_columns = get_all_citation_columns(df)
+    if len(citation_columns) > 0:
+        columns_to_drop += citation_columns
+    for column in columns_to_drop:
+        if column in moderations:
+            moderations.pop(column)
+
+    return moderations
+
+
+def run_postscore_guards(pipeline, predictions_df, postscore_guards=None):
+    """Run postscore guards on the input data."""
+    if not postscore_guards:
+        postscore_guards = pipeline.get_postscore_guards()
+    response_column_name = pipeline.get_input_column(GuardStage.RESPONSE)
+    blocked_completion_column_name = f"blocked_{response_column_name}"
+    input_df = predictions_df.copy(deep=True)
+    if len(postscore_guards) == 0:
+        input_df[blocked_completion_column_name] = False
+        return input_df, 0
+
+    start_time = time.time()
+    try:
+        postscore_df, postscore_latency = AsyncGuardExecutor(pipeline).run_guards(
+            input_df, postscore_guards, GuardStage.RESPONSE
+        )
+    except Exception as ex:
+        end_time = time.time()
+        _logger.error(f"Failed to run postscore guards: {ex}")
+        _logger.error(traceback.format_exc())
+        postscore_df = input_df
+        postscore_df[blocked_completion_column_name] = False
+        postscore_latency = end_time - start_time
+
+    # Again ensure the indexing matches the input dataframe indexing
+    postscore_df.index = predictions_df.index
+    _logger.debug("After passing completions through post score guards")
+    _logger.debug(postscore_df)
+    _logger.debug(f"Post Score Guard Latency: {postscore_latency} sec")
+
+    return postscore_df, postscore_latency
+
+
+def get_response_message_and_finish_reason(pipeline, postscore_df, streaming=False):
+    response_column_name = pipeline.get_input_column(GuardStage.RESPONSE)
+    blocked_completion_column_name = f"blocked_{response_column_name}"
+    replaced_response_column_name = f"replaced_{response_column_name}"
+    if postscore_df.empty:
+        response_message = NONE_CUSTOM_PY_RESPONSE
+        finish_reason = "stop"
+    elif postscore_df.loc[0, blocked_completion_column_name]:
+        blocked_message_completion_column_name = f"blocked_message_{response_column_name}"
+        response_message = postscore_df.loc[0, blocked_message_completion_column_name]
+        finish_reason = "content_filter"
+    elif (
+        replaced_response_column_name in postscore_df.columns
+        and postscore_df.loc[0, replaced_response_column_name]
+    ):
+        replaced_message_response_column_name = f"replaced_message_{response_column_name}"
+        response_message = postscore_df.loc[0, replaced_message_response_column_name]
+        # In case of streaming - if the guard replaces the text, we don't want to
+        # stop streaming - so don't put finish_reason in case of streaming
+        finish_reason = None if streaming else "content_filter"
+    else:
+        response_message = postscore_df.loc[0, response_column_name]
+        finish_reason = None if streaming else "stop"
+
+    return response_message, finish_reason
+
+
+def calculate_token_counts_and_confidence_score(pipeline, result_df):
+    prompt_column_name = pipeline.get_input_column(GuardStage.PROMPT)
+    blocked_prompt_column_name = f"blocked_{prompt_column_name}"
+    response_column_name = pipeline.get_input_column(GuardStage.RESPONSE)
+    blocked_completion_column_name = f"blocked_{response_column_name}"
+
+    encoding = tiktoken.get_encoding("cl100k_base")
+
+    citation_columns = get_citation_columns(result_df.columns)
+
+    def _get_llm_contexts(index):
+        contexts = []
+        if len(citation_columns) >= 0:
+            for column in citation_columns:
+                contexts.append(result_df.loc[index][column])
+        return contexts
+
+    for index, row in result_df.iterrows():
+        if not (
+            row.get(blocked_prompt_column_name, False)
+            or row.get(blocked_completion_column_name, False)
+        ):
+            completion = result_df.loc[index][response_column_name]
+            if completion != NONE_CUSTOM_PY_RESPONSE:
+                result_df.loc[index, "datarobot_token_count"] = len(
+                    encoding.encode(str(completion), disallowed_special=())
+                )
+                result_df.loc[index, "datarobot_confidence_score"] = get_rouge_1_score(
+                    pipeline.rouge_scorer, _get_llm_contexts(index), [completion]
+                )
+            else:
+                result_df.loc[index, "datarobot_confidence_score"] = 0.0
+        else:
+            # If the row is blocked, set default value
+            result_df.loc[index, "datarobot_confidence_score"] = 0.0
+
+
+def add_citations_to_df(citations, df):
+    if not citations:
+        return df
+
+    for index, citation in enumerate(citations):
+        df[f"CITATION_CONTENT_{index}"] = citation["content"]
+    return df
+
+
+def add_token_count_columns_to_df(pipeline, df, usage=None):
+    if not usage:
+        prompt_column_name = pipeline.get_input_column(GuardStage.PROMPT)
+        response_column_name = pipeline.get_input_column(GuardStage.RESPONSE)
+        df = calculate_token_counts_for_cost_calculations(
+            prompt_column_name, response_column_name, df
+        )
+    else:
+        df[PROMPT_TOKEN_COUNT_COLUMN_NAME_FROM_USAGE] = [usage.prompt_tokens]
+        df[RESPONSE_TOKEN_COUNT_COLUMN_NAME_FROM_USAGE] = [usage.completion_tokens]
+    return df
+
+
+def remove_unnecessary_columns(pipeline, result_df):
+    prompt_column_name = pipeline.get_input_column(GuardStage.PROMPT)
+    blocked_message_prompt_column_name = f"blocked_message_{prompt_column_name}"
+    response_column_name = pipeline.get_input_column(GuardStage.RESPONSE)
+    replaced_message_prompt_column_name = f"replaced_message_{prompt_column_name}"
+    blocked_message_completion_column_name = f"blocked_message_{response_column_name}"
+    replaced_message_response_column_name = f"replaced_message_{response_column_name}"
+    # We don't need these columns, because they have already been copied into
+    # 'completion' column
+    columns_to_remove = [
+        blocked_message_prompt_column_name,
+        blocked_message_completion_column_name,
+        replaced_message_prompt_column_name,
+        replaced_message_response_column_name,
+        f"Noneed_{prompt_column_name}",
+        f"Noneed_{response_column_name}",
+        PROMPT_TOKEN_COUNT_COLUMN_NAME_FROM_USAGE,
+        RESPONSE_TOKEN_COUNT_COLUMN_NAME_FROM_USAGE,
+        AGENTIC_PIPELINE_INTERACTIONS_ATTR,
+    ]
+    columns_to_remove.extend(get_all_citation_columns(result_df))
+    for column in columns_to_remove:
+        if column in result_df.columns:
+            result_df = result_df.drop(column, axis=1)
+
+    return result_df