salesforce-data-customcode 6.0.5.dev1__py3-none-any.whl → 6.0.6.dev2__py3-none-any.whl

datacustomcode/__init__.py

@@ -17,10 +17,14 @@ __all__ = [
     "AuthType",
     "Client",
     "Credentials",
+    "DefaultSparkEinsteinPredictions",
     "DefaultSparkLLMGateway",
     "PrintDataCloudWriter",
     "QueryAPIDataCloudReader",
+    "SparkEinsteinPredictions",
     "SparkLLMGateway",
+    "StreamingClient",
+    "einstein_predict_col",
     "llm_gateway_generate_text_col",
 ]
 
@@ -31,6 +35,10 @@ def __getattr__(name: str):
         from datacustomcode.client import Client
 
         return Client
+    elif name == "StreamingClient":
+        from datacustomcode.client import StreamingClient
+
+        return StreamingClient
     elif name == "AuthType":
         from datacustomcode.credentials import AuthType
 
@@ -59,4 +67,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,13 +17,17 @@ from __future__ import annotations
 from enum import Enum
 from typing import (
     TYPE_CHECKING,
+    Any,
     ClassVar,
     Dict,
     Optional,
+    TypeVar,
     Union,
+    cast,
 )
 
 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
@@ -32,14 +36,55 @@ from datacustomcode.spark.default import DefaultSparkSessionProvider
 if TYPE_CHECKING:
     from pathlib import Path
 
-    from pyspark.sql import Column, DataFrame as PySparkDataFrame
+    from pyspark.sql import (
+        Column,
+        DataFrame as PySparkDataFrame,
+        SparkSession,
+    )
+    from pyspark.sql.streaming import StreamingQuery
 
+    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
     from datacustomcode.spark.base import BaseSparkSessionProvider
 
 
+def _streaming_source_name() -> str:
+    """Return the streaming transform's read-source name.
+
+    Resolved from ``config.streaming_source``, which ``run_entrypoint``
+    populates from config.json's ``streamingSource`` field.
+
+    Raises:
+        RuntimeError: If no ``streaming_source`` has been configured (e.g. the
+            transform's config.json has no ``streamingSource`` field).
+    """
+    source = config.streaming_source
+    if not source:
+        raise RuntimeError(
+            "No streaming source configured. A streaming transform must declare "
+            "its read source in config.json under 'streamingSource'."
+        )
+    return source
+
+
+def _active_client() -> "_BaseClient":
+    """Return the client backing the module-level Spark column helpers.
+
+    Prefers an already-initialized singleton so a streaming job reuses its
+    :class:`StreamingClient` (and a batch job its :class:`Client`) rather than
+    forcing an unrelated client into existence. Falls back to building the
+    batch :class:`Client` when neither has been created yet.
+    """
+    if Client._instance is not None:
+        return Client._instance
+    if StreamingClient._instance is not None:
+        return StreamingClient._instance
+    return Client()
+
+
 def _build_spark_llm_gateway() -> "SparkLLMGateway":
     """Instantiate the SDK-configured :class:`SparkLLMGateway`.
 
@@ -95,10 +140,74 @@ def llm_gateway_generate_text_col(
         the generated text; on failure, ``status == "ERROR"`` and the
         ``error_*`` fields carry diagnostic detail.
     """
-    gateway = Client()._get_spark_llm_gateway()
+    gateway = _active_client()._get_spark_llm_gateway()
     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 = _active_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"
@@ -137,74 +246,86 @@ class DataCloudAccessLayerException(Exception):
         return msg
 
 
-class Client:
-    """Entrypoint for accessing DataCloud objects.
+_ClientT = TypeVar("_ClientT", bound="_BaseClient")
+
 
-    This is the object used to access Data Cloud DLOs and DMOs. Accessing DLOs/DMOs
-    are tracked and will throw an exception if they are mixed. In other words, you
-    can read from DLOs and write to DLOs, read from DMOs and write to DMOs, but you
-    cannot read from DLOs and write to DMOs or read from DMOs and write to DLOs.
-    Furthermore you cannot mix during merging tables. This class is a singleton to
-    prevent accidental mixing of DLOs and DMOs.
+class _BaseClient:
+    """Shared machinery for the Data Cloud client singletons.
 
-    You can provide custom readers and writers to the client for advanced use
-    cases, but this is not recommended for testing as they may result in unexpected
-    behavior once deployed to Data Cloud. By default, the client intercepts all
-    read/write operations and mocks access to Data Cloud. For example, during
-    writing, we print to the console instead of writing to Data Cloud.
+    Holds the wiring common to :class:`Client` (batch) and
+    :class:`StreamingClient`
+
+    This base class is not meant to be instantiated directly; use
+    :class:`Client` or :class:`StreamingClient`.
 
     Args:
-        finder: Find a file path
         reader: A custom reader to use for reading Data Cloud objects.
         writer: A custom writer to use for writing Data Cloud objects.
+        spark_provider: Optional custom :class:`BaseSparkSessionProvider`.
         spark_llm_gateway: Optional custom :class:`SparkLLMGateway`.
-
-    Example:
-    >>> client = Client()
-    >>> file_path = client.find_file_path("data.csv")
-    >>> dlo = client.read_dlo("my_dlo")
-    >>> client.write_to_dmo("my_dmo", dlo)
-    >>> answer = client.llm_gateway_generate_text("Generate a greeting message")
+        spark_einstein_predictions: Optional custom
+            :class:`SparkEinsteinPredictions`.
     """
 
-    _instance: ClassVar[Optional[Client]] = None
+    # Each concrete subclass gets its own ``_instance`` slot: reads fall through
+    # to this base default of ``None``, but ``cls._instance = ...`` in __new__
+    # always writes to the subclass, so ``Client`` and ``StreamingClient`` never
+    # share an instance.
+    _instance: ClassVar[Optional[_BaseClient]] = None
+    # Process-wide Spark session shared across BOTH client types. Unlike
+    # ``_instance``, this is written via ``_BaseClient._shared_spark`` (never
+    # ``cls._shared_spark``), so the slot lives on the base class and a
+    # ``Client`` and a ``StreamingClient`` in the same process reuse one session
+    # — and therefore one underlying connection — instead of opening two
+    # containing differing state
+    _shared_spark: ClassVar[Optional[SparkSession]] = None
     _reader: BaseDataCloudReader
     _writer: BaseDataCloudWriter
     _file: DefaultFindFilePath
     _spark_llm_gateway: Optional[SparkLLMGateway]
+    _spark_einstein_predictions: Optional[SparkEinsteinPredictions]
     _data_layer_history: dict[DataCloudObjectType, set[str]]
     _code_type: str
 
     def __new__(
-        cls,
+        cls: type[_ClientT],
         reader: Optional[BaseDataCloudReader] = None,
         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:
+    ) -> _ClientT:
 
         if cls._instance is None:
-            cls._instance = super().__new__(cls)
-            cls._instance._spark_llm_gateway = spark_llm_gateway
+            instance = super().__new__(cls)
+            instance._spark_llm_gateway = spark_llm_gateway
+            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:
-                # We need a spark because we will initialize readers and writers
-                if config.spark_config is None:
-                    raise ValueError(
-                        "Spark config is required when reader/writer is not provided"
-                    )
-
-                provider: BaseSparkSessionProvider
-                if spark_provider is not None:
-                    provider = spark_provider
-                elif config.spark_provider_config is not None:
-                    provider = config.spark_provider_config.to_object()
+                # We need a spark because we will initialize readers and writers.
+                # Reuse the process-wide session if one client already built it,
+                # so a Client and a StreamingClient share a single connection.
+                if _BaseClient._shared_spark is not None:
+                    spark = _BaseClient._shared_spark
                 else:
-                    provider = DefaultSparkSessionProvider()
-
-                spark = provider.get_session(config.spark_config)
+                    if config.spark_config is None:
+                        raise ValueError(
+                            "Spark config is required when reader/writer is not "
+                            "provided"
+                        )
+
+                    provider: BaseSparkSessionProvider
+                    if spark_provider is not None:
+                        provider = spark_provider
+                    elif config.spark_provider_config is not None:
+                        provider = config.spark_provider_config.to_object()
+                    else:
+                        provider = DefaultSparkSessionProvider()
+
+                    spark = provider.get_session(config.spark_config)
+                    _BaseClient._shared_spark = spark
 
             if config.reader_config is None and reader is None:
                 raise ValueError(
@@ -227,66 +348,17 @@ class Client:
             else:
                 writer_init = writer
 
-            cls._instance._reader = reader_init
-            cls._instance._writer = writer_init
-            cls._instance._file = DefaultFindFilePath()
-            cls._instance._data_layer_history = {
+            instance._reader = reader_init
+            instance._writer = writer_init
+            instance._file = DefaultFindFilePath()
+            instance._data_layer_history = {
                 DataCloudObjectType.DLO: set(),
                 DataCloudObjectType.DMO: set(),
             }
-        elif (reader is not None or writer is not None) and cls._instance is not None:
+            cls._instance = instance
+        elif reader is not None or writer is not None:
             raise ValueError("Cannot set reader or writer after client is initialized")
-        return cls._instance
-
-    def read_dlo(self, name: str) -> PySparkDataFrame:
-        """Read a DLO from Data Cloud.
-
-        Args:
-            name: The name of the DLO to read.
-
-        Returns:
-            A PySpark DataFrame containing the DLO data.
-        """
-        self._record_dlo_access(name)
-        return self._reader.read_dlo(name)  # type: ignore[no-any-return]
-
-    def read_dmo(self, name: str) -> PySparkDataFrame:
-        """Read a DMO from Data Cloud.
-
-        Args:
-            name: The name of the DMO to read.
-
-        Returns:
-            A PySpark DataFrame containing the DMO data.
-        """
-        self._record_dmo_access(name)
-        return self._reader.read_dmo(name)  # type: ignore[no-any-return]
-
-    def write_to_dlo(
-        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
-    ) -> None:
-        """Write a PySpark DataFrame to a DLO in Data Cloud.
-
-        Args:
-            name: The name of the DLO to write to.
-            dataframe: The PySpark DataFrame to write.
-            write_mode: The write mode to use for writing to the DLO.
-        """
-        self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO)
-        return self._writer.write_to_dlo(name, dataframe, write_mode, **kwargs)  # type: ignore[no-any-return]
-
-    def write_to_dmo(
-        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
-    ) -> None:
-        """Write a PySpark DataFrame to a DMO in Data Cloud.
-
-        Args:
-            name: The name of the DMO to write to.
-            dataframe: The PySpark DataFrame to write.
-            write_mode: The write mode to use for writing to the DMO.
-        """
-        self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO)
-        return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs)  # type: ignore[no-any-return]
+        return cast(_ClientT, cls._instance)
 
     def find_file_path(self, file_name: str) -> Path:
         """Resolve a bundled file shipped in the package to an absolute path.
@@ -358,6 +430,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:
@@ -371,3 +486,115 @@ class Client:
 
     def _record_dmo_access(self, name: str) -> None:
         self._data_layer_history[DataCloudObjectType.DMO].add(name)
+
+
+class Client(_BaseClient):
+    """Entrypoint for batch access to Data Cloud objects.
+
+    This is the object used to read and write bounded snapshots of Data Cloud
+    DLOs and DMOs.
+    """
+
+    _instance: ClassVar[Optional[Client]] = None
+
+    def read_dlo(self, name: str) -> PySparkDataFrame:
+        """Read a DLO from Data Cloud.
+
+        Args:
+            name: The name of the DLO to read.
+
+        Returns:
+            A PySpark DataFrame containing the DLO data.
+        """
+        self._record_dlo_access(name)
+        return self._reader.read_dlo(name)  # type: ignore[no-any-return]
+
+    def read_dmo(self, name: str) -> PySparkDataFrame:
+        """Read a DMO from Data Cloud.
+
+        Args:
+            name: The name of the DMO to read.
+
+        Returns:
+            A PySpark DataFrame containing the DMO data.
+        """
+        self._record_dmo_access(name)
+        return self._reader.read_dmo(name)  # type: ignore[no-any-return]
+
+    def write_to_dlo(
+        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
+    ) -> None:
+        """Write a PySpark DataFrame to a DLO in Data Cloud.
+
+        Args:
+            name: The name of the DLO to write to.
+            dataframe: The PySpark DataFrame to write.
+            write_mode: The write mode to use for writing to the DLO.
+        """
+        self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO)
+        return self._writer.write_to_dlo(name, dataframe, write_mode, **kwargs)  # type: ignore[no-any-return]
+
+    def write_to_dmo(
+        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
+    ) -> None:
+        """Write a PySpark DataFrame to a DMO in Data Cloud.
+
+        Args:
+            name: The name of the DMO to write to.
+            dataframe: The PySpark DataFrame to write.
+            write_mode: The write mode to use for writing to the DMO.
+        """
+        self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO)
+        return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs)  # type: ignore[no-any-return]
+
+
+class StreamingClient(_BaseClient):
+    """Entrypoint for streaming (``DELTA_SYNC``) access to Data Cloud objects.
+
+    This is the streaming counterpart to :class:`Client`. Instead of reading and
+    writing bounded snapshots, it reads a DLO/DMO change feed as a streaming
+    DataFrame and writes the transformed stream back via a ``StreamingQuery``.
+    """
+
+    _instance: ClassVar[Optional[StreamingClient]] = None
+
+    def read_dlo_deltas(self) -> PySparkDataFrame:
+        """Read the streaming change feed (deltas) for a DLO from Data Cloud.
+
+        For use in a streaming (``DELTA_SYNC``) BYOC transform. Returns a
+        streaming DataFrame whose rows carry the change-feed metadata columns
+        (``_record_type``, ``_commit_*``) alongside the source columns.
+
+        Returns:
+            A streaming PySpark DataFrame over the DLO change feed.
+        """
+        self._record_dlo_access(_streaming_source_name())
+        return self._reader.read_dlo_deltas()  # type: ignore[no-any-return]
+
+    def read_dmo_deltas(self) -> PySparkDataFrame:
+        """Read the streaming change feed (deltas) for a DMO from Data Cloud.
+
+        Returns:
+            A streaming PySpark DataFrame over the DMO change feed.
+        """
+        self._record_dmo_access(_streaming_source_name())
+        return self._reader.read_dmo_deltas()  # type: ignore[no-any-return]
+
+    def write_dlo_deltas(
+        self, name: str, dataframe: PySparkDataFrame, **kwargs
+    ) -> StreamingQuery:
+        """Write a streaming DataFrame of deltas to a DLO in Data Cloud.
+
+        Starts a streaming query that writes each micro-batch to the
+        target DLO and returns the  ``StreamingQuery`` handle; the caller
+        typically calls ``query.awaitTermination()``.
+
+        Args:
+            name: The name of the DLO to write to.
+            dataframe: The streaming PySpark DataFrame to write.
+
+        Returns:
+            The started ``StreamingQuery``.
+        """
+        self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO)
+        return self._writer.write_dlo_deltas(name, dataframe, **kwargs)  # type: ignore[no-any-return]

datacustomcode/config.py

@@ -89,6 +89,9 @@ class ClientConfig(BaseConfig):
     spark_provider_config: Union[
         SparkProviderConfig[BaseSparkSessionProvider], None
     ] = None
+    # Source object name for a streaming (DELTA_SYNC) transform, populated by
+    # ``run_entrypoint`` from config.json's ``streamingSource`` field
+    streaming_source: Union[str, None] = None
 
     def update(self, other: ClientConfig) -> ClientConfig:
         """Merge this ClientConfig with another, respecting force flags.
@@ -116,6 +119,8 @@ class ClientConfig(BaseConfig):
         self.spark_provider_config = merge(
             self.spark_provider_config, other.spark_provider_config
         )
+        if other.streaming_source is not None:
+            self.streaming_source = other.streaming_source
         return self
 
 

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