salesforce-data-customcode 6.0.4.dev3__py3-none-any.whl → 6.0.6.dev1__py3-none-any.whl

datacustomcode/__init__.py

@@ -17,10 +17,13 @@ __all__ = [
     "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}")

datacustomcode/client.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 from enum import Enum
 from typing import (
     TYPE_CHECKING,
+    Any,
     ClassVar,
     Dict,
     Optional,
@@ -24,6 +25,7 @@ from typing import (
 )
 
 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 @@ if TYPE_CHECKING:
 
     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 @@ class Client:
         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 @@ class Client:
             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:

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:

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",
 ]

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

datacustomcode/einstein_predictions/spark_base.py

@@ -0,0 +1,72 @@
+# 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.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Optional,
+)
+
+from datacustomcode.mixin import UserExtendableNamedConfigMixin
+
+if TYPE_CHECKING:
+    from pyspark.sql import Column
+
+    from datacustomcode.einstein_predictions.types import PredictionType
+
+
+class SparkEinsteinPredictions(ABC, UserExtendableNamedConfigMixin):
+    CONFIG_NAME: str
+
+    def __init__(self, **kwargs: Any) -> None:
+        pass
+
+    @abstractmethod
+    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 and return the response data.
+
+        ``features`` maps each model feature column name to a single scalar
+        value (``str``/``float``/``bool``). The value is wrapped into a
+        single-element prediction column of the appropriate type.
+        """
+
+    @abstractmethod
+    def einstein_predict_col(
+        self,
+        model_api_name: str,
+        prediction_type: PredictionType,
+        features: Dict[str, "Column"],
+        settings: Optional[Dict[str, Any]] = None,
+    ) -> "Column":
+        """Build a Spark ``Column`` that invokes Einstein predict per row and
+        yields a struct ``{status, response, error_code, error_message}``.
+
+        ``features`` maps each model feature column name to a Spark ``Column``
+        supplying that feature's per-row value. Select an individual field,
+        e.g. ``einstein_predict_col(...)["response"]``. ``response`` holds the
+        prediction response payload as a JSON string. Returning a struct means
+        a single failing row leaves the rest of the DataFrame intact — callers
+        can inspect ``status`` / ``error_code`` per row instead of having the
+        Spark job abort.
+        """

datacustomcode/einstein_predictions/spark_default.py

@@ -0,0 +1,239 @@
+# 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.
+from __future__ import annotations
+
+import json
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Optional,
+)
+
+from datacustomcode.einstein_predictions.spark_base import SparkEinsteinPredictions
+from datacustomcode.einstein_predictions.types import (
+    PredictionColumn,
+    PredictionRequest,
+    PredictionType,
+)
+
+if TYPE_CHECKING:
+    from pyspark.sql import Column
+
+    from datacustomcode.einstein_predictions.base import EinsteinPredictions
+    from datacustomcode.einstein_predictions.types import PredictionResponse
+
+
+_STATUS_SUCCESS = "SUCCESS"
+_STATUS_ERROR = "ERROR"
+
+
+class DefaultSparkEinsteinPredictions(SparkEinsteinPredictions):
+
+    CONFIG_NAME = "DefaultSparkEinsteinPredictions"
+
+    def __init__(
+        self,
+        einstein_predictions: Optional["EinsteinPredictions"] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        if einstein_predictions is None:
+            einstein_predictions = _build_underlying_predictions()
+        self._einstein_predictions: "EinsteinPredictions" = einstein_predictions
+
+    def einstein_predict(
+        self,
+        model_api_name: str,
+        prediction_type: PredictionType,
+        features: Dict[str, Any],
+        settings: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        return _invoke_predictions(
+            self._einstein_predictions,
+            model_api_name,
+            prediction_type,
+            features,
+            settings,
+        )
+
+    def einstein_predict_col(
+        self,
+        model_api_name: str,
+        prediction_type: PredictionType,
+        features: Dict[str, "Column"],
+        settings: Optional[Dict[str, Any]] = None,
+    ) -> "Column":
+        """Build a per-row UDF that returns a struct ``{status, response,
+        error_code, error_message}`` so per-row failures do not abort the
+        Spark job. Callers select the field they want, e.g.
+        ``einstein_predict_col(...)["response"]``. ``response`` carries the
+        prediction response payload serialized as a JSON string.
+        """
+        from pyspark.sql.functions import struct, udf
+        from pyspark.sql.types import (
+            StringType,
+            StructField,
+            StructType,
+        )
+
+        feature_names = list(features.keys())
+        values_col = struct(*[features[name].alias(name) for name in feature_names])
+
+        predictions = self._einstein_predictions
+        result_schema = StructType(
+            [
+                StructField("status", StringType(), True),
+                StructField("response", StringType(), True),
+                StructField("error_code", StringType(), True),
+                StructField("error_message", StringType(), True),
+            ]
+        )
+
+        def _predict(values_row: Any) -> Dict[str, Optional[str]]:
+            if values_row is None:
+                return {
+                    "status": _STATUS_ERROR,
+                    "response": None,
+                    "error_code": None,
+                    "error_message": "features column was null for this row",
+                }
+            row_features = (
+                values_row.asDict()
+                if hasattr(values_row, "asDict")
+                else dict(values_row)
+            )
+            return _invoke_predictions_as_struct(
+                predictions,
+                model_api_name,
+                prediction_type,
+                row_features,
+                settings,
+            )
+
+        return udf(_predict, result_schema)(values_col)
+
+
+def _build_underlying_predictions() -> "EinsteinPredictions":
+    from datacustomcode.einstein_predictions_config import einstein_predictions_config
+
+    cfg = einstein_predictions_config.einstein_predictions_config
+    if cfg is None:
+        raise RuntimeError(
+            "einstein_predictions_config is not configured. Add an "
+            "'einstein_predictions_config' section to config.yaml."
+        )
+    return cfg.to_object()
+
+
+def _feature_to_prediction_column(name: str, value: Any) -> PredictionColumn:
+    """Wrap a single scalar feature value into a one-element prediction column.
+
+    The value's Python type selects the prediction column value type. ``bool``
+    is checked before ``int``/``float`` because ``bool`` is a subclass of
+    ``int`` in Python.
+    """
+    if isinstance(value, bool):
+        return PredictionColumn(column_name=name, boolean_values=[value])
+    if isinstance(value, (int, float)):
+        return PredictionColumn(column_name=name, double_values=[float(value)])
+    return PredictionColumn(column_name=name, string_values=[str(value)])
+
+
+def _build_request(
+    model_api_name: str,
+    prediction_type: PredictionType,
+    features: Dict[str, Any],
+    settings: Optional[Dict[str, Any]],
+) -> PredictionRequest:
+    prediction_columns = [
+        _feature_to_prediction_column(name, value) for name, value in features.items()
+    ]
+    return PredictionRequest(
+        prediction_type=prediction_type,
+        model_api_name=model_api_name,
+        prediction_columns=prediction_columns,
+        settings=settings,
+    )
+
+
+def _call_predictions(
+    predictions: "EinsteinPredictions",
+    model_api_name: str,
+    prediction_type: PredictionType,
+    features: Dict[str, Any],
+    settings: Optional[Dict[str, Any]],
+) -> "PredictionResponse":
+    """Build the request and dispatch it to the underlying predictions resource."""
+    request = _build_request(model_api_name, prediction_type, features, settings)
+    return predictions.predict(request)
+
+
+def _invoke_predictions(
+    predictions: "EinsteinPredictions",
+    model_api_name: str,
+    prediction_type: PredictionType,
+    features: Dict[str, Any],
+    settings: Optional[Dict[str, Any]],
+) -> Dict[str, Any]:
+    from datacustomcode.einstein_predictions.errors import EinsteinPredictionsCallError
+
+    response = _call_predictions(
+        predictions, model_api_name, prediction_type, features, settings
+    )
+    if not response.is_success:
+        error_code = _extract_error_code(response)
+        raise EinsteinPredictionsCallError(
+            f"Einstein Predictions call failed: "
+            f"status_code={response.status_code}, "
+            f"error_code={error_code!r}, message={response.data!r}",
+            status=response.status_code,
+            error_code=error_code,
+            error_message=str(response.data) if response.data else None,
+        )
+    return response.data or {}
+
+
+def _invoke_predictions_as_struct(
+    predictions: "EinsteinPredictions",
+    model_api_name: str,
+    prediction_type: PredictionType,
+    features: Dict[str, Any],
+    settings: Optional[Dict[str, Any]],
+) -> Dict[str, Optional[str]]:
+    response = _call_predictions(
+        predictions, model_api_name, prediction_type, features, settings
+    )
+    if not response.is_success:
+        return {
+            "status": _STATUS_ERROR,
+            "response": None,
+            "error_code": _extract_error_code(response),
+            "error_message": str(response.data) if response.data else None,
+        }
+    return {
+        "status": _STATUS_SUCCESS,
+        "response": json.dumps(response.data) if response.data is not None else None,
+        "error_code": None,
+        "error_message": None,
+    }
+
+
+def _extract_error_code(response: "PredictionResponse") -> Optional[str]:
+    if response.data:
+        error_code = response.data.get("errorCode")
+        if error_code is not None:
+            return str(error_code)
+    return None

datacustomcode/einstein_predictions_config.py

@@ -21,11 +21,17 @@ from typing import (
     Union,
 )
 
-from datacustomcode.common_config import BaseConfig, default_config_file
+from datacustomcode.common_config import (
+    BaseConfig,
+    BaseObjectConfig,
+    default_config_file,
+)
 from datacustomcode.einstein_platform_config import CredentialsObjectConfig
 from datacustomcode.einstein_predictions.base import EinsteinPredictions
+from datacustomcode.einstein_predictions.spark_base import SparkEinsteinPredictions
 
 _E = TypeVar("_E", bound=EinsteinPredictions)
+_S = TypeVar("_S", bound=SparkEinsteinPredictions)
 
 
 class EinsteinPredictionsObjectConfig(CredentialsObjectConfig[_E], Generic[_E]):
@@ -54,6 +60,44 @@ class EinsteinPredictionsConfig(BaseConfig):
         return self
 
 
+class SparkEinsteinPredictionsObjectConfig(BaseObjectConfig, Generic[_S]):
+    type_to_create: ClassVar[Type[SparkEinsteinPredictions]] = SparkEinsteinPredictions  # type: ignore[type-abstract]
+
+    def to_object(self) -> SparkEinsteinPredictions:
+        type_ = self.type_to_create.subclass_from_config_name(self.type_config_name)
+        return type_(**self.options)
+
+
+class SparkEinsteinPredictionsConfig(BaseConfig):
+    spark_einstein_predictions_config: Union[
+        SparkEinsteinPredictionsObjectConfig[SparkEinsteinPredictions], None
+    ] = None
+
+    def update(
+        self, other: "SparkEinsteinPredictionsConfig"
+    ) -> "SparkEinsteinPredictionsConfig":
+        def merge(
+            config_a: Union[SparkEinsteinPredictionsObjectConfig, None],
+            config_b: Union[SparkEinsteinPredictionsObjectConfig, None],
+        ) -> Union[SparkEinsteinPredictionsObjectConfig, None]:
+            if config_a is not None and config_a.force:
+                return config_a
+            if config_b:
+                return config_b
+            return config_a
+
+        self.spark_einstein_predictions_config = merge(
+            self.spark_einstein_predictions_config,
+            other.spark_einstein_predictions_config,
+        )
+        return self
+
+
 # Global Einstein Predictions config instance
 einstein_predictions_config = EinsteinPredictionsConfig()
 einstein_predictions_config.load(default_config_file())
+
+
+# Global Spark Einstein Predictions config instance
+spark_einstein_predictions_config = SparkEinsteinPredictionsConfig()
+spark_einstein_predictions_config.load(default_config_file())

datacustomcode/templates/script/payload/entrypoint.py

@@ -38,6 +38,23 @@ def main():
         >>> generated_text = client.llm_gateway_generate_text(
         ...     prompt, model_id
         ... )
+
+    You can also score rows with an Einstein Studio model. The per-row helper
+    returns the same ``{status, response, error_code, error_message}`` struct,
+    where ``response`` is the prediction payload as a JSON string.
+
+    Example:
+
+        >>> from datacustomcode.client import einstein_predict_col
+        >>> from datacustomcode.einstein_predictions.types import PredictionType
+            df_scored = df.withColumn(
+            ...     "prediction__c",
+            ...     einstein_predict_col(
+            ...         "my_regression_model",
+            ...         PredictionType.REGRESSION,
+            ...         {"beds": col("beds__c"), "baths": col("baths__c")},
+            ...     )["response"],
+            ... )
     """
 
     # Drop specific columns related to relationships

{salesforce_data_customcode-6.0.4.dev3.dist-info → salesforce_data_customcode-6.0.6.dev1.dist-info}/METADATA

@@ -1,11 +1,11 @@
 Metadata-Version: 2.4
 Name: salesforce-data-customcode
-Version: 6.0.4.dev3
+Version: 6.0.6.dev1
 Summary: Data Cloud Custom Code SDK
 License-Expression: Apache-2.0
 License-File: LICENSE.txt
 Requires-Python: >=3.10,<3.12
-Classifier: Development Status :: 4 - Beta
+Classifier: Development Status :: 5 - Production/Stable
 Classifier: Operating System :: Unix
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
@@ -23,7 +23,7 @@ Requires-Dist: salesforce-cdp-connector (>=1.0.19)
 Requires-Dist: setuptools_scm (>=7.1.0,<8.0.0)
 Description-Content-Type: text/markdown
 
-# Data Cloud Custom Code SDK (BETA)
+# Data Cloud Custom Code SDK
 
 This package provides a development kit for creating custom data transformations in [Data Cloud](https://www.salesforce.com/data/). It allows you to write your own data processing logic in Python while leveraging Data Cloud's infrastructure for data access and running data transformations, mapping execution into Data Cloud data structures like [Data Model Objects](https://help.salesforce.com/s/articleView?id=data.c360_a_data_model_objects.htm&type=5) and [Data Lake Objects](https://help.salesforce.com/s/articleView?id=sf.c360_a_data_lake_objects.htm&language=en_US&type=5).
 
@@ -36,7 +36,7 @@ Use of this project with Salesforce is subject to the [TERMS OF USE](./TERMS_OF_
 - **Python 3.11 only** (currently supported version - if your system version is different, we recommend using [pyenv](https://github.com/pyenv/pyenv) to configure 3.11)
 - JDK 17
 - Docker support like [Docker Desktop](https://docs.docker.com/desktop/)
-- A salesforce org with some DLOs or DMOs with data and this feature enabled (it is not GA)
+- A salesforce org with some DLOs or DMOs with data and this feature enabled
 - **One of the following** for authentication:
   - A Salesforce org already authenticated via the [Salesforce CLI](https://developer.salesforce.com/tools/salesforcecli)
     (simplest — no External Client App needed)
@@ -116,7 +116,7 @@ After modifying the `entrypoint.py` as needed, using any dependencies you add in
 ```zsh
 cd my_package
 datacustomcode scan ./payload/entrypoint.py
-datacustomcode deploy --path ./payload --name my_custom_script --cpu-size CPU_L
+datacustomcode deploy --path ./payload --name my_custom_script --cpu-size CPU_L --sf-cli-org myorg
 ```
 
 > [!TIP]
@@ -345,6 +345,7 @@ Options:
 - `--description TEXT`: Description of the transformation job (default: "")
 - `--network TEXT`: docker network (default: "default")
 - `--cpu-size TEXT`: CPU size for the deployment (default: `CPU_2XL`). Available options: CPU_L(Large), CPU_XL(Extra Large), CPU_2XL(2X Large), CPU_4XL(4X Large)
+- `--sf-cli-org TEXT`: Salesforce CLI org alias or username (e.g. `myorg`). Fetches credentials via `sf org display` — no `datacustomcode configure` step needed. Takes precedence over `--profile` if both are supplied.
 - `--function-invoke-opt TEXT`: Currently we support only `UnstructuredChunking` for functions.
 
 
@@ -397,6 +398,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.
@@ -508,7 +559,7 @@ sf --version
 **Browser-based (recommended for developer orgs and sandboxes):**
 ```zsh
 # Production / Developer Edition
-sf org login web --alias myorg
+sf org login web --alias myorg --instance-url <your-instance-url>
 
 # Sandbox
 sf org login web --alias mysandbox --instance-url https://test.salesforce.com

{salesforce_data_customcode-6.0.4.dev3.dist-info → salesforce_data_customcode-6.0.6.dev1.dist-info}/RECORD

@@ -1,21 +1,24 @@
-datacustomcode/__init__.py,sha256=qx-X3U9lU_791ZbAxhGqQeyDw6eqTi9-E9AbYnwJfxU,2071
+datacustomcode/__init__.py,sha256=SGsDnWDTihWkLNnWdqvH08zklQgxykar73JAcPL2plI,2666
 datacustomcode/auth.py,sha256=fpSjhIBdv9trC8yq2vuljAix_Euu-4Ah7HDCGhYjOxI,8309
 datacustomcode/cli.py,sha256=i7hPoQqJskTmx3Odz2wVgkxBnf_PFaIawI12FXhaRUE,13264
-datacustomcode/client.py,sha256=fc3JmDbAFcTT-yP_k_0jNiGqNp_KwZaIpEsI-NPxNcw,15115
+datacustomcode/client.py,sha256=x0JN_hHXfpEdjHeINQEdUQt-as3j17v6uoDJHjfLe-Y,20241
 datacustomcode/cmd.py,sha256=ZMs46aydJw2EaU26JgCtZmnqESQFHvvaJz10hnjZTBk,3537
 datacustomcode/common_config.py,sha256=SAUnxj3kqmOeWwPmFoYq4tuxokMgURVm4QwcGi-avL4,1928
 datacustomcode/config.py,sha256=2Pk61ieQsEWSxKDxlS66_rliKE8iX0j-9LweEm0aaGo,4072
-datacustomcode/config.yaml,sha256=Wd8Zd2f1__hRv7166H5IBbU5YoOEOdQNHgAb4ONbFnA,845
+datacustomcode/config.yaml,sha256=ldssmQYPghmYSSPaelm3F2mC99njavhrWsINRSzWymM,933
 datacustomcode/constants.py,sha256=SskQpUPnQE7LjKh4Z3AdZ2p191_0Kf73jV8gdBXDPLU,1436
 datacustomcode/credentials.py,sha256=D-7Zd3Nh_wStdj8_wUy5cnC8M_mdbfBijhIAi-2EbXs,9467
 datacustomcode/deploy.py,sha256=TZ6L4VrAnL-SlmAUMbEYguVpCgOm_Ncv8NdaBR3LX4Y,21906
 datacustomcode/einstein_platform_client.py,sha256=ON2B_m--vdbCVVMVcLc81T-BRZ5-UuxVCer8E6OEQPA,4083
 datacustomcode/einstein_platform_config.py,sha256=6lb_FRbEzdt5a8-2cR15f13ZKw674uvRJ4Xq0o_pQXE,1490
-datacustomcode/einstein_predictions/__init__.py,sha256=o-iask36phwGA7hGi_sxJpwbITO6h_STDe0J92qD2_M,859
+datacustomcode/einstein_predictions/__init__.py,sha256=_RIDTqcEHDRu6AsCsj9lRBykNkHI8Ch8JBILL1ioPtA,1237
 datacustomcode/einstein_predictions/base.py,sha256=OBV445dpZRtCF5cpdHH49w_6RU_jZY1A-e7VtHTQCFU,1061
+datacustomcode/einstein_predictions/errors.py,sha256=EoxVU5Mf14AIAyVBedI1JnbsS_o853kY-iMA1u-KPcw,1224
 datacustomcode/einstein_predictions/impl/default.py,sha256=fEFhrosozhG0yT87w4Gu8wbe1NgHxBqdAduwWHdWGN4,3011
+datacustomcode/einstein_predictions/spark_base.py,sha256=MnFH2-pJ7ugFnc9iUvCdMt1vPAPAmiImptjjIGh9bvo,2567
+datacustomcode/einstein_predictions/spark_default.py,sha256=tYL6u1DY3Ogu8xF7o1VqXqTGrcP2VUUMI6BShjXWvns,8119
 datacustomcode/einstein_predictions/types.py,sha256=W7H3sFzm_NdBtnia7O576w0OvM7jwzfGVQmFiZTO8Og,6098
-datacustomcode/einstein_predictions_config.py,sha256=oxKgp2u7t5SdHVsasBBxsvd3KAfXuZ9uQXTrtpcCUlU,2135
+datacustomcode/einstein_predictions_config.py,sha256=7zUot5fUnXiihiu_9NStvk0-e9Otv55-Xk7dKqrMfBY,3758
 datacustomcode/file/__init__.py,sha256=gamfOD1VnAtEslRBpqh-yKiVjkG_wYWYdSatGIfsN-w,621
 datacustomcode/file/base.py,sha256=2IUlQufPvpHccgyoeZlKhcktO4pS4NKdha03Whq33Zk,745
 datacustomcode/file/path/__init__.py,sha256=gamfOD1VnAtEslRBpqh-yKiVjkG_wYWYdSatGIfsN-w,621
@@ -87,13 +90,13 @@ datacustomcode/templates/script/examples/employee_hierarchy/employee_data.csv,sh
 datacustomcode/templates/script/examples/employee_hierarchy/entrypoint.py,sha256=Mfm3iQtEHTQRW6cmZTwRKdTh3IeGWR-teXrd6x-YLSY,2223
 datacustomcode/templates/script/jupyterlab.sh,sha256=IHR3YQ8d_busuyvesByhJgzCKULIFPI-Hqogs0NPhCs,2432
 datacustomcode/templates/script/payload/config.json,sha256=0d2mEMt4NHIeOUTsgiPMuRKdOLIQxmh-Wq-IfEqU-gc,28
-datacustomcode/templates/script/payload/entrypoint.py,sha256=bLUhC4l03H-OQJKdMf11vGXmBnga6QitRXo-g9HibZg,1741
+datacustomcode/templates/script/payload/entrypoint.py,sha256=4Uph0ILa5ukk_6C90paK3uzQn0zZvNLr9ho3o_ZwErQ,2474
 datacustomcode/templates/script/requirements-dev.txt,sha256=OWwuy1awesqOZOS3Zizmdd6BDnSePpGFN5bq5IrALvc,176
 datacustomcode/templates/script/requirements.txt,sha256=qJbqzs5z0Hrx590U3T7dHq_m8UQEx2TdhgrJSz-QHIQ,40
 datacustomcode/token_provider.py,sha256=kGPxdxGZEr6VFknW7jFb9P6wOvRaRo-sfdBN2cUjEZY,6792
 datacustomcode/version.py,sha256=9LlbVrzwBvut1L308QbwNBgxdQk5CrghH39JARVSxms,989
-salesforce_data_customcode-6.0.4.dev3.dist-info/METADATA,sha256=Tg_qVJWkMhhj1pz1hSHwVpLh00hppX1cDMkmts7ycJM,23234
-salesforce_data_customcode-6.0.4.dev3.dist-info/WHEEL,sha256=eY7nduwzv-ldUxpzbRlxwvC693Hg6PX8bWDjEHjZ_dk,88
-salesforce_data_customcode-6.0.4.dev3.dist-info/entry_points.txt,sha256=WpQ94UB7UuRCYGOLtJV3vAgm1tl7iVf43VYRWIuNu5Y,74
-salesforce_data_customcode-6.0.4.dev3.dist-info/licenses/LICENSE.txt,sha256=iOi8EmQpfkFhMENi7VYtkl2EqS14pEOccoXHiW2dyPU,11443
-salesforce_data_customcode-6.0.4.dev3.dist-info/RECORD,,
+salesforce_data_customcode-6.0.6.dev1.dist-info/METADATA,sha256=k11RadLqAMw2inLBjUV6rLcjor-Itqf3wX7u1EJhuO8,25115
+salesforce_data_customcode-6.0.6.dev1.dist-info/WHEEL,sha256=eY7nduwzv-ldUxpzbRlxwvC693Hg6PX8bWDjEHjZ_dk,88
+salesforce_data_customcode-6.0.6.dev1.dist-info/entry_points.txt,sha256=WpQ94UB7UuRCYGOLtJV3vAgm1tl7iVf43VYRWIuNu5Y,74
+salesforce_data_customcode-6.0.6.dev1.dist-info/licenses/LICENSE.txt,sha256=iOi8EmQpfkFhMENi7VYtkl2EqS14pEOccoXHiW2dyPU,11443
+salesforce_data_customcode-6.0.6.dev1.dist-info/RECORD,,