Merge pull request #119 from forcedotcom/einstein-predict-script

Support einstein predict for script in public SDK

Author: jcatt-sf

README.md

@@ -373,6 +373,56 @@ datacustomcode run ./payload/entrypoint.py --sf-cli-org myorg
 ```
 
 
+## Testing Einstein Predictions
+
+You can use AI models configured in Einstein Studio to score your data while
+transforming it. As with the LLM Gateway, there are two flavors: a one-shot
+scalar call (`client.einstein_predict`) and a per-row column helper
+(`einstein_predict_col`). Below is a sample code example:
+
+```
+from datacustomcode.client import Client, einstein_predict_col
+from datacustomcode.einstein_predictions.types import PredictionType
+
+
+def main():
+  client = Client()
+  df = client.read_dlo("Input__dll")
+  # einstein_predict_col returns a struct
+  # {status, response, error_code, error_message} per row, so per-row
+  # failures don't abort the Spark job. `response` is the prediction
+  # payload as a JSON string. Pick the field you want with [].
+  df_scored = df.withColumn(
+    "prediction__c",
+    einstein_predict_col(
+        "my_regression_model",  # An AI model in your org
+        PredictionType.REGRESSION,
+        {"square_feet": col("square_feet__c"), "beds": col("beds__c")},
+    )["response"],
+  )
+
+  dlo_name = "Output_dll"
+  client.write_to_dlo(dlo_name, df_scored, write_mode=WriteMode.APPEND)
+
+  # One-shot scalar prediction returns the response payload as a dict
+  prediction = client.einstein_predict(
+    "my_regression_model",
+    PredictionType.REGRESSION,
+    {"square_feet": 1800, "beds": 3},
+  )
+
+if __name__ == "__main__":
+  main()
+```
+
+Testing this code locally uses the same External Client App setup described in
+[Testing LLM Gateway](#testing-llm-gateway). Once your `myorg` alias is set up,
+run:
+```
+datacustomcode run ./payload/entrypoint.py --sf-cli-org myorg
+```
+
+
 ## Docker usage
 
 The SDK provides Docker-based development options that allow you to test your code in an environment that closely resembles Data Cloud's execution environment.

src/datacustomcode/__init__.py

@@ -17,10 +17,13 @@
     "AuthType",
     "Client",
     "Credentials",
+    "DefaultSparkEinsteinPredictions",
     "DefaultSparkLLMGateway",
     "PrintDataCloudWriter",
     "QueryAPIDataCloudReader",
+    "SparkEinsteinPredictions",
     "SparkLLMGateway",
+    "einstein_predict_col",
     "llm_gateway_generate_text_col",
 ]
 
@@ -59,4 +62,16 @@ def __getattr__(name: str):
         from datacustomcode.client import llm_gateway_generate_text_col
 
         return llm_gateway_generate_text_col
+    elif name == "SparkEinsteinPredictions":
+        from datacustomcode.einstein_predictions import SparkEinsteinPredictions
+
+        return SparkEinsteinPredictions
+    elif name == "DefaultSparkEinsteinPredictions":
+        from datacustomcode.einstein_predictions import DefaultSparkEinsteinPredictions
+
+        return DefaultSparkEinsteinPredictions
+    elif name == "einstein_predict_col":
+        from datacustomcode.client import einstein_predict_col
+
+        return einstein_predict_col
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/datacustomcode/client.py

@@ -17,13 +17,15 @@
 from enum import Enum
 from typing import (
     TYPE_CHECKING,
+    Any,
     ClassVar,
     Dict,
     Optional,
     Union,
 )
 
 from datacustomcode.config import config
+from datacustomcode.einstein_predictions_config import spark_einstein_predictions_config
 from datacustomcode.file.path.default import DefaultFindFilePath
 from datacustomcode.io.reader.base import BaseDataCloudReader
 from datacustomcode.llm_gateway_config import spark_llm_gateway_config
@@ -34,6 +36,8 @@
 
     from pyspark.sql import Column, DataFrame as PySparkDataFrame
 
+    from datacustomcode.einstein_predictions.spark_base import SparkEinsteinPredictions
+    from datacustomcode.einstein_predictions.types import PredictionType
     from datacustomcode.io.reader.base import BaseDataCloudReader
     from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
     from datacustomcode.llm_gateway.spark_base import SparkLLMGateway
@@ -99,6 +103,70 @@ def llm_gateway_generate_text_col(
     return gateway.llm_gateway_generate_text_col(template, values, model_id=model_id)
 
 
+def _build_spark_einstein_predictions() -> "SparkEinsteinPredictions":
+    """Instantiate the SDK-configured :class:`SparkEinsteinPredictions`.
+
+    Raises:
+        RuntimeError: If no ``spark_einstein_predictions_config`` has been loaded.
+    """
+    cfg = spark_einstein_predictions_config.spark_einstein_predictions_config
+    if cfg is None:
+        raise RuntimeError(
+            "spark_einstein_predictions_config is not configured. Add a "
+            "'spark_einstein_predictions_config' section to config.yaml."
+        )
+    return cfg.to_object()
+
+
+def einstein_predict_col(
+    model_api_name: str,
+    prediction_type: "PredictionType",
+    features: Dict[str, "Column"],
+    settings: Optional[Dict[str, Any]] = None,
+) -> "Column":
+    """Build a Spark Column that runs an Einstein prediction per row.
+
+    The returned Column yields a struct ``{status, response, error_code,
+    error_message}`` for each row. Use ``[...]`` (or ``getField``) to pick the
+    field you want, e.g. ``einstein_predict_col(...)["response"]``. ``response``
+    holds the prediction response payload as a JSON string. Per-row failures
+    populate ``status`` / ``error_code`` / ``error_message`` so a single bad row
+    does not abort the whole Spark job.
+
+    Example:
+
+        >>> from datacustomcode.einstein_predictions.types import PredictionType
+        >>> result = einstein_predict_col(
+        ...     "my_regression_model",
+        ...     PredictionType.REGRESSION,
+        ...     {"square_feet": col("square_feet__c"), "beds": col("beds__c")},
+        ... )
+        >>> df.withColumn("prediction__c", result["response"])
+        >>> # …or keep the struct around and inspect failures:
+        >>> df.withColumn("pred", result).select(
+        ...     "pred.status", "pred.response", "pred.error_message"
+        ... )
+
+    Args:
+        model_api_name: API name of the Einstein model to invoke.
+        prediction_type: The :class:`PredictionType` of the model.
+        features: A mapping from model feature column name to a Spark ``Column``
+            supplying that feature's per-row value.
+        settings: Optional prediction settings forwarded to the model.
+
+    Returns:
+        A Spark ``Column`` of ``StructType`` with fields ``status``,
+        ``response``, ``error_code``, and ``error_message`` (all nullable
+        strings). On success, ``status == "SUCCESS"`` and ``response`` holds
+        the JSON-serialized prediction payload; on failure, ``status ==
+        "ERROR"`` and the ``error_*`` fields carry diagnostic detail.
+    """
+    predictions = Client()._get_spark_einstein_predictions()
+    return predictions.einstein_predict_col(
+        model_api_name, prediction_type, features, settings=settings
+    )
+
+
 class DataCloudObjectType(Enum):
     DLO = "dlo"
     DMO = "dmo"
@@ -158,6 +226,8 @@ class Client:
         reader: A custom reader to use for reading Data Cloud objects.
         writer: A custom writer to use for writing Data Cloud objects.
         spark_llm_gateway: Optional custom :class:`SparkLLMGateway`.
+        spark_einstein_predictions: Optional custom
+            :class:`SparkEinsteinPredictions`.
 
     Example:
     >>> client = Client()
@@ -172,6 +242,7 @@ class Client:
     _writer: BaseDataCloudWriter
     _file: DefaultFindFilePath
     _spark_llm_gateway: Optional[SparkLLMGateway]
+    _spark_einstein_predictions: Optional[SparkEinsteinPredictions]
     _data_layer_history: dict[DataCloudObjectType, set[str]]
     _code_type: str
 
@@ -181,12 +252,14 @@ def __new__(
         writer: Optional[BaseDataCloudWriter] = None,
         spark_provider: Optional[BaseSparkSessionProvider] = None,
         spark_llm_gateway: Optional[SparkLLMGateway] = None,
+        spark_einstein_predictions: Optional[SparkEinsteinPredictions] = None,
         code_type: str = "script",
     ) -> Client:
 
         if cls._instance is None:
             cls._instance = super().__new__(cls)
             cls._instance._spark_llm_gateway = spark_llm_gateway
+            cls._instance._spark_einstein_predictions = spark_einstein_predictions
             # Initialize Readers and Writers from config
             # and/or provided reader and writer
             if reader is None or writer is None:
@@ -358,6 +431,49 @@ def _get_spark_llm_gateway(self) -> SparkLLMGateway:
             self._spark_llm_gateway = _build_spark_llm_gateway()
         return self._spark_llm_gateway
 
+    def einstein_predict(
+        self,
+        model_api_name: str,
+        prediction_type: "PredictionType",
+        features: Dict[str, Any],
+        settings: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """Issue a one-shot Einstein prediction. This is the scalar counterpart
+        to :func:`einstein_predict_col`: it runs **once** — not per row. Use the
+        column helper method instead when you want to fan a prediction out
+        across every row of a DataFrame.
+
+        Example:
+
+            >>> from datacustomcode.einstein_predictions.types import PredictionType
+            >>> response = Client().einstein_predict(
+            ...     "my_regression_model",
+            ...     PredictionType.REGRESSION,
+            ...     {"square_feet": 1800, "beds": 3},
+            ... )
+
+        Args:
+            model_api_name: API name of the Einstein model to invoke.
+            prediction_type: The :class:`PredictionType` of the model.
+            features: A mapping from model feature column name to a single
+                scalar value (``str`` / ``float`` / ``bool``).
+            settings: Optional prediction settings forwarded to the model.
+
+        Returns:
+            The prediction response payload as a plain Python ``dict``.
+
+        Raises:
+            EinsteinPredictionsCallError: If the prediction call fails.
+        """
+        return self._get_spark_einstein_predictions().einstein_predict(
+            model_api_name, prediction_type, features, settings=settings
+        )
+
+    def _get_spark_einstein_predictions(self) -> SparkEinsteinPredictions:
+        if self._spark_einstein_predictions is None:
+            self._spark_einstein_predictions = _build_spark_einstein_predictions()
+        return self._spark_einstein_predictions
+
     def _validate_data_layer_history_does_not_contain(
         self, data_cloud_object_type: DataCloudObjectType
     ) -> None:

src/datacustomcode/config.yaml

@@ -24,6 +24,9 @@ einstein_predictions_config:
   options:
     credentials_profile: default
 
+spark_einstein_predictions_config:
+  type_config_name: DefaultSparkEinsteinPredictions
+
 llm_gateway_config:
   type_config_name: DefaultLLMGateway
   options:

src/datacustomcode/einstein_predictions/__init__.py

@@ -14,9 +14,17 @@
 # limitations under the License.
 
 from datacustomcode.einstein_predictions.base import EinsteinPredictions
+from datacustomcode.einstein_predictions.errors import EinsteinPredictionsCallError
 from datacustomcode.einstein_predictions.impl.default import DefaultEinsteinPredictions
+from datacustomcode.einstein_predictions.spark_base import SparkEinsteinPredictions
+from datacustomcode.einstein_predictions.spark_default import (
+    DefaultSparkEinsteinPredictions,
+)
 
 __all__ = [
     "DefaultEinsteinPredictions",
+    "DefaultSparkEinsteinPredictions",
     "EinsteinPredictions",
+    "EinsteinPredictionsCallError",
+    "SparkEinsteinPredictions",
 ]

src/datacustomcode/einstein_predictions/errors.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Exceptions raised by Einstein Predictions implementations."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+
+class EinsteinPredictionsCallError(RuntimeError):
+    """Raised when an Einstein Predictions call returns an error."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        status: Optional[object] = None,
+        error_code: Optional[str] = None,
+        error_message: Optional[str] = None,
+    ) -> None:
+        super().__init__(message)
+        self.status = status
+        self.error_code = error_code
+        self.error_message = error_message