salesforce-data-customcode 3.0.1__tar.gz → 3.0.2__tar.gz

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-3.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: salesforce-data-customcode
-Version: 3.0.1
+Version: 3.0.2
 Summary: Data Cloud Custom Code SDK
 License-Expression: Apache-2.0
 License-File: LICENSE.txt

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-3.0.2}/pyproject.toml

@@ -18,7 +18,7 @@ license = "Apache-2.0"
 name = "salesforce-data-customcode"
 readme = "README.md"
 requires-python = ">=3.10,<3.12"
-version = "3.0.1"
+version = "3.0.2"
 
 [tool.black]
 exclude = '''

salesforce_data_customcode-3.0.2/src/datacustomcode/common_config.py

@@ -0,0 +1,63 @@
+# 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 abc import ABC, abstractmethod
+import os
+from typing import Any
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+)
+import yaml
+
+DEFAULT_CONFIG_NAME = "config.yaml"
+
+
+def default_config_file() -> str:
+    return os.path.join(os.path.dirname(__file__), DEFAULT_CONFIG_NAME)
+
+
+class ForceableConfig(BaseModel):
+    force: bool = Field(
+        default=False,
+        description="If True, this takes precedence over parameters passed to the "
+        "initializer of the client",
+    )
+
+
+class BaseObjectConfig(ForceableConfig):
+    model_config = ConfigDict(validate_default=True, extra="forbid")
+    type_config_name: str = Field(
+        description="The config name of the object to create",
+    )
+    options: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Options passed to the constructor.",
+    )
+
+
+class BaseConfig(ABC, BaseModel):
+    @abstractmethod
+    def update(self, other: Any) -> "BaseConfig": ...
+
+    def load(self, config_path: str) -> "BaseConfig":
+        """Load configuration from a YAML file and merge with existing config"""
+        with open(config_path, "r") as f:
+            config_data = yaml.safe_load(f)
+
+        loaded_config = self.__class__.model_validate(config_data)
+        self.update(loaded_config)
+        return self

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-3.0.2}/src/datacustomcode/config.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 from __future__ import annotations
 
-import os
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -26,12 +25,14 @@ from typing import (
     cast,
 )
 
-from pydantic import (
-    BaseModel,
-    ConfigDict,
-    Field,
+from pydantic import Field
+
+from datacustomcode.common_config import (
+    BaseConfig,
+    BaseObjectConfig,
+    ForceableConfig,
+    default_config_file,
 )
-import yaml
 
 # This lets all readers and writers to be findable via config
 from datacustomcode.io import *  # noqa: F403
@@ -42,36 +43,15 @@ from datacustomcode.proxy.base import BaseProxyAccessLayer
 from datacustomcode.proxy.client.base import BaseProxyClient  # noqa: TCH002
 from datacustomcode.spark.base import BaseSparkSessionProvider
 
-DEFAULT_CONFIG_NAME = "config.yaml"
-
-
 if TYPE_CHECKING:
     from pyspark.sql import SparkSession
 
 
-class ForceableConfig(BaseModel):
-    force: bool = Field(
-        default=False,
-        description="If True, this takes precedence over parameters passed to the "
-        "initializer of the client.",
-    )
-
-
 _T = TypeVar("_T", bound="BaseDataAccessLayer")
 
 
-class AccessLayerObjectConfig(ForceableConfig, Generic[_T]):
-    model_config = ConfigDict(validate_default=True, extra="forbid")
+class AccessLayerObjectConfig(BaseObjectConfig, Generic[_T]):
     type_base: ClassVar[Type[BaseDataAccessLayer]] = BaseDataAccessLayer
-    type_config_name: str = Field(
-        description="The config name of the object to create. "
-        "For metrics, this would might be 'ipmnormal'. For custom classes, you can "
-        "assign a name to a class variable `CONFIG_NAME` and reference it here.",
-    )
-    options: dict[str, Any] = Field(
-        default_factory=dict,
-        description="Options passed to the constructor.",
-    )
 
     def to_object(self, spark: SparkSession) -> _T:
         type_ = self.type_base.subclass_from_config_name(self.type_config_name)
@@ -97,35 +77,25 @@ _P = TypeVar("_P", bound=BaseSparkSessionProvider)
 _PX = TypeVar("_PX", bound=BaseProxyAccessLayer)
 
 
-class ProxyAccessLayerObjectConfig(ForceableConfig, Generic[_PX]):
+class ProxyAccessLayerObjectConfig(BaseObjectConfig, Generic[_PX]):
     """Config for proxy clients that take no constructor args (e.g. no spark)."""
 
-    model_config = ConfigDict(validate_default=True, extra="forbid")
     type_base: ClassVar[Type[BaseProxyAccessLayer]] = BaseProxyAccessLayer
-    type_config_name: str = Field(
-        description="CONFIG_NAME of the proxy client (e.g. 'LocalProxyClient').",
-    )
-    options: dict[str, Any] = Field(default_factory=dict)
 
     def to_object(self) -> _PX:
         type_ = self.type_base.subclass_from_config_name(self.type_config_name)
         return cast(_PX, type_(**self.options))
 
 
-class SparkProviderConfig(ForceableConfig, Generic[_P]):
-    model_config = ConfigDict(validate_default=True, extra="forbid")
+class SparkProviderConfig(BaseObjectConfig, Generic[_P]):
     type_base: ClassVar[Type[BaseSparkSessionProvider]] = BaseSparkSessionProvider
-    type_config_name: str = Field(
-        description="CONFIG_NAME of the Spark session provider."
-    )
-    options: dict[str, Any] = Field(default_factory=dict)
 
     def to_object(self) -> _P:
         type_ = self.type_base.subclass_from_config_name(self.type_config_name)
         return cast(_P, type_(**self.options))
 
 
-class ClientConfig(BaseModel):
+class ClientConfig(BaseConfig):
     reader_config: Union[AccessLayerObjectConfig[BaseDataCloudReader], None] = None
     writer_config: Union[AccessLayerObjectConfig[BaseDataCloudWriter], None] = None
     proxy_config: Union[ProxyAccessLayerObjectConfig[BaseProxyClient], None] = None
@@ -163,31 +133,10 @@ class ClientConfig(BaseModel):
         )
         return self
 
-    def load(self, config_path: str) -> ClientConfig:
-        """Load a config from a file and update this config with it.
 
-        Args:
-            config_path: The path to the config file
-
-        Returns:
-            Self, with updated values from the loaded config.
-        """
-        with open(config_path, "r") as f:
-            config_data = yaml.safe_load(f)
-        loaded_config = ClientConfig.model_validate(config_data)
-
-        return self.update(loaded_config)
-
-
-config = ClientConfig()
 """Global config object.
 
 This is the object that makes config accessible globally and globally mutable.
 """
-
-
-def _defaults() -> str:
-    return os.path.join(os.path.dirname(__file__), DEFAULT_CONFIG_NAME)
-
-
-config.load(_defaults())
+config = ClientConfig()
+config.load(default_config_file())

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-3.0.2}/src/datacustomcode/config.yaml

@@ -23,3 +23,7 @@ proxy_config:
   type_config_name: LocalProxyClientProvider
   options:
     credentials_profile: default
+
+einstein_predictions_config:
+  type_config_name: DefaultEinsteinPredictions
+  options: {}

salesforce_data_customcode-3.0.2/src/datacustomcode/einstein_predictions/__init__.py

@@ -0,0 +1,22 @@
+# 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 datacustomcode.einstein_predictions.base import EinsteinPredictions
+from datacustomcode.einstein_predictions.impl.default import DefaultEinsteinPredictions
+
+__all__ = [
+    "EinsteinPredictions",
+    "DefaultEinsteinPredictions",
+]

salesforce_data_customcode-3.0.2/src/datacustomcode/einstein_predictions/base.py

@@ -0,0 +1,32 @@
+# 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 abc import ABC, abstractmethod
+
+from datacustomcode.einstein_predictions.types import (
+    PredictionRequest,
+    PredictionResponse,
+)
+from datacustomcode.mixin import UserExtendableNamedConfigMixin
+
+
+class EinsteinPredictions(ABC, UserExtendableNamedConfigMixin):
+    CONFIG_NAME: str
+
+    def __init__(self, **kwargs):
+        pass
+
+    @abstractmethod
+    def predict(self, request: PredictionRequest) -> PredictionResponse: ...

salesforce_data_customcode-3.0.2/src/datacustomcode/einstein_predictions/impl/default.py

@@ -0,0 +1,35 @@
+# 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 datacustomcode.einstein_predictions.base import EinsteinPredictions
+from datacustomcode.einstein_predictions.types import (
+    PredictionRequest,
+    PredictionResponse,
+)
+
+
+class DefaultEinsteinPredictions(EinsteinPredictions):
+    CONFIG_NAME = "DefaultEinsteinPredictions"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def predict(self, request: PredictionRequest) -> PredictionResponse:
+        return PredictionResponse(
+            version="v1",
+            prediction_type=request.prediction_type,
+            status_code=200,
+            data={"results": [{"prediction": {"predictedValue": 1.0}}]},
+        )

salesforce_data_customcode-3.0.2/src/datacustomcode/einstein_predictions/types.py

@@ -0,0 +1,184 @@
+# 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 enum import Enum, unique
+from typing import (
+    Any,
+    Dict,
+    Literal,
+    Optional,
+)
+
+from pydantic import (
+    BaseModel,
+    Field,
+    model_validator,
+)
+
+
+@unique
+class PredictionType(Enum):
+    REGRESSION = 1
+    CLUSTERING = 2
+    CLASSIFICATION = 3
+    MULTI_OUTCOME = 4
+    BINARY_CLASSIFICATION = 5
+
+
+class PredictionColumn(BaseModel):
+    column_name: str = Field(min_length=1, description="Column name")
+    string_values: Optional[list[str]] = Field(
+        default=None, min_length=1, description="Column string values"
+    )
+    double_values: Optional[list[float]] = Field(
+        default=None, min_length=1, description="Column double values"
+    )
+    boolean_values: Optional[list[bool]] = Field(
+        default=None, min_length=1, description="Column boolean values"
+    )
+    date_values: Optional[list[str]] = Field(
+        default=None, min_length=1, description="Column date values"
+    )
+    datetime_values: Optional[list[str]] = Field(
+        default=None, min_length=1, description="Column datetime values"
+    )
+
+    @model_validator(mode="after")
+    def validate_exactly_one_value_type(self):
+        set_count = sum(
+            [
+                self.string_values is not None,
+                self.double_values is not None,
+                self.boolean_values is not None,
+                self.date_values is not None,
+                self.datetime_values is not None,
+            ]
+        )
+
+        if set_count != 1:
+            raise ValueError("Exactly one value type must be set")
+
+        return self
+
+
+class PredictionColumBuilder:
+    def __init__(self) -> None:
+        self._column_name: Optional[str] = None
+        self._string_values: Optional[list[str]] = None
+        self._double_values: Optional[list[float]] = None
+        self._boolean_values: Optional[list[bool]] = None
+        self._date_values: Optional[list[str]] = None
+        self._datetime_values: Optional[list[str]] = None
+
+    def set_column_name(self, column_name: str) -> "PredictionColumBuilder":
+        self._column_name = column_name
+        return self
+
+    def set_string_values(self, string_values: list[str]) -> "PredictionColumBuilder":
+        self._string_values = string_values
+        return self
+
+    def set_double_values(self, double_values: list[float]) -> "PredictionColumBuilder":
+        self._double_values = double_values
+        return self
+
+    def set_boolean_values(
+        self, boolean_values: list[bool]
+    ) -> "PredictionColumBuilder":
+        self._boolean_values = boolean_values
+        return self
+
+    def set_date_values(self, date_values: list[str]) -> "PredictionColumBuilder":
+        self._date_values = date_values
+        return self
+
+    def set_datetime_values(
+        self, datetime_values: list[str]
+    ) -> "PredictionColumBuilder":
+        self._datetime_values = datetime_values
+        return self
+
+    def build(self) -> PredictionColumn:
+        return PredictionColumn(
+            column_name=self._column_name,
+            string_values=self._string_values,
+            double_values=self._double_values,
+            boolean_values=self._boolean_values,
+            date_values=self._date_values,
+            datetime_values=self._datetime_values,
+        )
+
+
+class PredictionRequest(BaseModel):
+    version: Literal["v1"] = Field(
+        default="v1", description="API version, must be 'v1'"
+    )
+    prediction_type: PredictionType = Field(description="Prediction type")
+    model_api_name: str = Field(
+        min_length=1, description="API name of the model to use"
+    )
+    prediction_columns: list[PredictionColumn] = Field(
+        min_length=1, description="List of prediction columns"
+    )
+    settings: Optional[Dict[str, Any]] = Field(
+        default=None, description="Settings for the prediction request"
+    )
+
+
+class PredictionRequestBuilder:
+    def __init__(self) -> None:
+        self._prediction_type: Optional[PredictionType] = None
+        self._model_api_name: Optional[str] = None
+        self._prediction_columns: list[PredictionColumn] = []
+        self._settings: Optional[Dict[str, Any]] = None
+
+    def set_prediction_type(
+        self, prediction_type: PredictionType
+    ) -> "PredictionRequestBuilder":
+        self._prediction_type = prediction_type
+        return self
+
+    def set_model_api_name(self, model_api_name: str) -> "PredictionRequestBuilder":
+        self._model_api_name = model_api_name
+        return self
+
+    def set_prediction_columns(
+        self, prediction_columns: list[PredictionColumn]
+    ) -> "PredictionRequestBuilder":
+        self._prediction_columns = prediction_columns
+        return self
+
+    def set_settings(self, settings: Dict[str, Any]):
+        self._settings = settings
+        return self
+
+    def build(self) -> PredictionRequest:
+        return PredictionRequest(
+            prediction_type=self._prediction_type,
+            model_api_name=self._model_api_name,
+            prediction_columns=self._prediction_columns,
+            settings=self._settings,
+        )
+
+
+class PredictionResponse(BaseModel):
+    version: Literal["v1"] = Field(default="v1", description="API version")
+    prediction_type: PredictionType = Field(description="Prediction type")
+    status_code: int = Field(description="HTTP status code")
+    data: Optional[Dict[str, Any]] = Field(default=None, description="Response data")
+
+    @property
+    def is_success(self) -> bool:
+        return self.status_code == 200

salesforce_data_customcode-3.0.2/src/datacustomcode/einstein_predictions_config.py

@@ -0,0 +1,67 @@
+# 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 typing import (
+    ClassVar,
+    Generic,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from datacustomcode.common_config import (
+    BaseConfig,
+    BaseObjectConfig,
+    default_config_file,
+)
+from datacustomcode.einstein_predictions.base import EinsteinPredictions
+
+_E = TypeVar("_E", bound=EinsteinPredictions)
+
+
+class EinsteinPredictionsObjectConfig(BaseObjectConfig, Generic[_E]):
+    type_base: ClassVar[Type[EinsteinPredictions]] = EinsteinPredictions  # type: ignore[type-abstract]
+
+    def to_object(self) -> _E:
+        type_ = self.type_base.subclass_from_config_name(self.type_config_name)
+        return cast(_E, type_(**self.options))
+
+
+class EinsteinPredictionsConfig(BaseConfig):
+    einstein_predictions_config: Union[
+        EinsteinPredictionsObjectConfig[EinsteinPredictions], None
+    ] = None
+
+    def update(self, other: "EinsteinPredictionsConfig") -> "EinsteinPredictionsConfig":
+        def merge(
+            config_a: Union[EinsteinPredictionsObjectConfig, None],
+            config_b: Union[EinsteinPredictionsObjectConfig, None],
+        ) -> Union[EinsteinPredictionsObjectConfig, None]:
+            if config_a is not None and config_a.force:
+                return config_a
+            if config_b:
+                return config_b
+            return config_a
+
+        self.einstein_predictions_config = merge(
+            self.einstein_predictions_config, other.einstein_predictions_config
+        )
+        return self
+
+
+# Global Einstein Predictions config instance
+einstein_predictions_config = EinsteinPredictionsConfig()
+einstein_predictions_config.load(default_config_file())

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-3.0.2}/src/datacustomcode/function/runtime.py

@@ -17,6 +17,8 @@
 import threading
 from typing import Optional
 
+from datacustomcode.einstein_predictions.base import EinsteinPredictions
+from datacustomcode.einstein_predictions_config import einstein_predictions_config
 from datacustomcode.file.path.default import DefaultFindFilePath
 from datacustomcode.function.base import BaseRuntime
 from datacustomcode.llm_gateway.default import DefaultLLMGateway
@@ -65,6 +67,7 @@ class Runtime(BaseRuntime):
         # Initialize resources
         self._llm_gateway = DefaultLLMGateway()
         self._file = DefaultFindFilePath()
+        self._einstein_predictions: Optional[EinsteinPredictions] = None
 
     @property
     def llm_gateway(self) -> DefaultLLMGateway:
@@ -75,3 +78,16 @@ class Runtime(BaseRuntime):
     def file(self) -> DefaultFindFilePath:
         """Access file operations."""
         return self._file
+
+    @property
+    def einstein_predictions(self) -> EinsteinPredictions:
+        if self._einstein_predictions is None:
+            if einstein_predictions_config.einstein_predictions_config is None:
+                raise RuntimeError(
+                    "Einstein Predictions is not configured. Add "
+                    "'einstein_predictions_config' section to config.yaml"
+                )
+            self._einstein_predictions = (
+                einstein_predictions_config.einstein_predictions_config.to_object()
+            )
+        return self._einstein_predictions

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-3.0.2}/src/datacustomcode/llm_gateway/base.py

@@ -14,9 +14,11 @@
 # limitations under the License.
 from __future__ import annotations
 
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 
+from datacustomcode.mixin import UserExtendableNamedConfigMixin
+
 if TYPE_CHECKING:
     from datacustomcode.llm_gateway.types.generate_text_request import (
         GenerateTextRequest,
@@ -26,8 +28,10 @@ if TYPE_CHECKING:
     )
 
 
-class LLMGateway:
-    def __init__(self) -> None:
+class LLMGateway(ABC, UserExtendableNamedConfigMixin):
+    CONFIG_NAME: str
+
+    def __init__(self, **kwargs):
         pass
 
     @abstractmethod

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-3.0.2}/src/datacustomcode/llm_gateway/default.py

@@ -22,6 +22,8 @@ from datacustomcode.llm_gateway.types.generate_text_response_builder import (
 
 
 class DefaultLLMGateway(LLMGateway):
+    CONFIG_NAME = "DefaultLLMGateway"
+
     def generate_text(self, request: GenerateTextRequest) -> GenerateTextResponse:
 
         response_data = {

salesforce_data_customcode-3.0.2/src/datacustomcode/llm_gateway_config.py

@@ -0,0 +1,65 @@
+# 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 typing import (
+    ClassVar,
+    Generic,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from datacustomcode.common_config import (
+    BaseConfig,
+    BaseObjectConfig,
+    default_config_file,
+)
+from datacustomcode.llm_gateway.base import LLMGateway
+
+_E = TypeVar("_E", bound=LLMGateway)
+
+
+class LLMGatewayObjectConfig(BaseObjectConfig, Generic[_E]):
+    type_base: ClassVar[Type[LLMGateway]] = LLMGateway  # type: ignore[type-abstract]
+
+    def to_object(self) -> _E:
+        type_ = self.type_base.subclass_from_config_name(self.type_config_name)
+        return cast(_E, type_(**self.options))
+
+
+class LLMGatewayConfig(BaseConfig):
+    llm_gateway_config: Union[LLMGatewayObjectConfig[LLMGateway], None] = None
+
+    def update(self, other: "LLMGatewayConfig") -> "LLMGatewayConfig":
+        def merge(
+            config_a: Union[LLMGatewayObjectConfig, None],
+            config_b: Union[LLMGatewayObjectConfig, None],
+        ) -> Union[LLMGatewayObjectConfig, None]:
+            if config_a is not None and config_a.force:
+                return config_a
+            if config_b:
+                return config_b
+            return config_a
+
+        self.llm_gateway_config = merge(
+            self.llm_gateway_config, other.llm_gateway_config
+        )
+        return self
+
+
+# Global LLM Gateway config instance
+llm_gateway_config = LLMGatewayConfig()
+llm_gateway_config.load(default_config_file())

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-3.0.2}/src/datacustomcode/templates/function/payload/entrypoint.py

@@ -2,6 +2,11 @@ import logging
 from typing import List
 from uuid import uuid4
 
+from datacustomcode.einstein_predictions.types import (
+    PredictionColumBuilder,
+    PredictionRequestBuilder,
+    PredictionType,
+)
 from datacustomcode.function import Runtime
 from datacustomcode.llm_gateway.types.generate_text_request_builder import (
     GenerateTextRequestBuilder,
@@ -38,6 +43,28 @@ def chunk_text(text: str, chunk_size: int = 1000) -> List[str]:
     return chunks
 
 
+def make_einstein_prediction(runtime: Runtime) -> None:
+    column = (
+        PredictionColumBuilder()
+        .set_column_name("col1")
+        .set_string_values(["str1", "str2"])
+        .build()
+    )
+    prediction_request = (
+        PredictionRequestBuilder()
+        .set_prediction_type(PredictionType.REGRESSION)
+        .set_model_api_name("regressionModel")
+        .set_prediction_columns([column])
+        .build()
+    )
+
+    prediction_response = runtime.einstein_predictions.predict(prediction_request)
+    print(
+        f"Einstein prediction results - success: {prediction_response.is_success} \
+           response data: {prediction_response.data}"
+    )
+
+
 def function(request: dict, runtime: Runtime) -> dict:
     logger.info("Inside Function")
     logger.info(request)
@@ -46,6 +73,8 @@ def function(request: dict, runtime: Runtime) -> dict:
     output_chunks = []
     current_seq_no = 1  # Start sequence number from 1
 
+    make_einstein_prediction(runtime)
+
     builder = GenerateTextRequestBuilder()
     llm_request = builder.set_prompt("Hello").set_model("modelName").build()
     llm_response = runtime.llm_gateway.generate_text(llm_request)