add EP to configuration

Author: atulikumwenayo

src/datacustomcode/common_config.py

@@ -0,0 +1,44 @@
+import os
+import yaml
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+)
+
+from typing import Any
+
+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(BaseModel):
+    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

src/datacustomcode/config.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 from __future__ import annotations
 
-import os
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -27,11 +26,8 @@
 )
 
 from pydantic import (
-    BaseModel,
-    ConfigDict,
     Field,
 )
-import yaml
 
 # This lets all readers and writers to be findable via config
 from datacustomcode.io import *  # noqa: F403
@@ -41,38 +37,18 @@
 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"
+from datacustomcode.common_config import ForceableConfig, BaseObjectConfig, BaseConfig, default_config_file
 
 
 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)
         return cast(_T, type_(spark=spark, **self.options))
@@ -97,35 +73,22 @@ class SparkConfig(ForceableConfig):
 _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 +126,9 @@ def merge(
         )
         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())

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: {}

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

src/datacustomcode/einstein_predictions/base.py

@@ -16,7 +16,13 @@
 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
 
-class EinsteinPredictions(ABC):
     @abstractmethod
     def predict(self, request: PredictionRequest) -> PredictionResponse: ...

src/datacustomcode/einstein_predictions/impl/default.py

@@ -14,14 +14,22 @@
 # 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:
-    def __init__(self, base_url: str, access_token: str) -> None:
-        pass
+class DefaultEinsteinPredictions(EinsteinPredictions):
+    CONFIG_NAME = "DefaultEinsteinPredictions"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
 
     def predict(self, request: PredictionRequest) -> PredictionResponse:
-        pass
+        return PredictionResponse(
+              version="v1",
+              prediction_type=request.prediction_type,
+              status_code=200,
+              data={"results": [{"prediction": {"predictedValue": "1"}}]}
+        )

src/datacustomcode/einstein_predictions_config.py

@@ -0,0 +1,57 @@
+# 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.einstein_predictions.base import EinsteinPredictions
+from datacustomcode.common_config import BaseObjectConfig, BaseConfig, default_config_file
+
+_E = TypeVar("_E", bound=EinsteinPredictions)
+
+class EinsteinPredictionsObjectConfig(BaseObjectConfig, Generic[_E]):
+    type_base: ClassVar[Type[EinsteinPredictions]] = EinsteinPredictions
+    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())

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 @@ def __init__(self) -> None:
         # Initialize resources
         self._llm_gateway = DefaultLLMGateway()
         self._file = DefaultFindFilePath()
+        self._einstein_predictions: Optional[EinsteinPredictions] = None
 
     @property
     def llm_gateway(self) -> DefaultLLMGateway:
@@ -75,3 +78,13 @@ def llm_gateway(self) -> DefaultLLMGateway:
     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

tests/test_config.py

@@ -11,7 +11,7 @@
     AccessLayerObjectConfig,
     ClientConfig,
     SparkConfig,
-    _defaults,
+    default_config_file,
     config,
 )
 from datacustomcode.io.base import BaseDataAccessLayer
@@ -298,8 +298,8 @@ def test_load_config_from_file(self):
             os.unlink(temp_path)
 
     def test_defaults(self):
-        # Just verify that _defaults function exists and returns a string path
-        result = _defaults()
+        # Just verify that default_config_file function exists and returns a string path
+        result = default_config_file()
         assert isinstance(result, str)
         assert result.endswith("config.yaml")