update based on the PR feedback, including removing code duplication

Author: atulikumwenayo

pyproject.toml

@@ -107,6 +107,7 @@ python = ">=3.10,<3.12"
 pyyaml = "^6.0"
 salesforce-cdp-connector = ">=1.0.19"
 setuptools_scm = "^7.1.0"
+requests = "2.33.1"
 
 [tool.poetry.group.dev.dependencies]
 build = "*"

src/datacustomcode/einstein_platform_client.py

@@ -29,12 +29,15 @@
 
 
 class EinsteinPlatformClient:
-    EINSTEIN_PLATFORM_URL = "https://api.salesforce.com/einstein/platform/v1"
+    EINSTEIN_PLATFORM_MODELS_URL = (
+        "https://api.salesforce.com/einstein/platform/v1/models"
+    )
 
     def __init__(
         self,
         credentials_profile: Optional[str] = None,
         sf_cli_org: Optional[str] = None,
+        **kwargs: Any,
     ):
         if sf_cli_org:
             self._token_provider: TokenProvider = SFCLITokenProvider(sf_cli_org)
@@ -44,6 +47,7 @@ def __init__(
             self._token_provider = CredentialsTokenProvider(profile)
             logger.debug(f"Using credentials token provider with profile: {profile}")
         self.token_response = None
+        super().__init__(**kwargs)
 
     def get_headers(self):
         if self.token_response is None:

src/datacustomcode/einstein_platform_config.py

@@ -0,0 +1,41 @@
+# 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,
+    Optional,
+    Type,
+    cast,
+)
+
+from datacustomcode.common_config import BaseObjectConfig
+
+
+class CredentialsObjectConfig(BaseObjectConfig):
+    type_to_create: ClassVar[Type]
+    credentials_profile: Optional[str] = None
+    sf_cli_org: Optional[str] = None
+
+    def to_object(self):
+        """Create an object instance, automatically including credentials in options"""
+
+        options = self.options.copy()
+        if self.credentials_profile is not None:
+            options["credentials_profile"] = self.credentials_profile
+        if self.sf_cli_org is not None:
+            options["sf_cli_org"] = self.sf_cli_org
+
+        type_ = self.type_to_create.subclass_from_config_name(self.type_config_name)
+        return cast(type_, type_(**options))

src/datacustomcode/einstein_predictions/impl/default.py

@@ -18,7 +18,6 @@
     ClassVar,
     Dict,
     List,
-    Optional,
 )
 
 from loguru import logger
@@ -43,17 +42,6 @@ class DefaultEinsteinPredictions(EinsteinPlatformClient, EinsteinPredictions):
         PredictionType.MULTI_OUTCOME: "multi-outcome",
     }
 
-    def __init__(
-        self,
-        credentials_profile: Optional[str] = None,
-        sf_cli_org: Optional[str] = None,
-        **kwargs,
-    ):
-        EinsteinPlatformClient.__init__(
-            self, credentials_profile=credentials_profile, sf_cli_org=sf_cli_org
-        )
-        EinsteinPredictions.__init__(self, **kwargs)
-
     def predict(self, request: PredictionRequest) -> PredictionResponse:
         endpoint = self.ENDPOINT_MAP.get(request.prediction_type)
         if not endpoint:
@@ -63,8 +51,7 @@ def predict(self, request: PredictionRequest) -> PredictionResponse:
             )
 
         api_url = (
-            f"{self.EINSTEIN_PLATFORM_URL}/models/"
-            f"{request.model_api_name}/{endpoint}"
+            f"{self.EINSTEIN_PLATFORM_MODELS_URL}/{request.model_api_name}/{endpoint}"
         )
 
         prediction_columns: List[Dict[str, Any]] = []

src/datacustomcode/einstein_predictions_config.py

@@ -19,25 +19,17 @@
     Type,
     TypeVar,
     Union,
-    cast,
 )
 
-from datacustomcode.common_config import (
-    BaseConfig,
-    BaseObjectConfig,
-    default_config_file,
-)
+from datacustomcode.common_config import BaseConfig, default_config_file
+from datacustomcode.einstein_platform_config import CredentialsObjectConfig
 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 EinsteinPredictionsObjectConfig(CredentialsObjectConfig, Generic[_E]):
+    type_to_create: ClassVar[Type[EinsteinPredictions]] = EinsteinPredictions  # type: ignore[type-abstract]
 
 
 class EinsteinPredictionsConfig(BaseConfig):

src/datacustomcode/llm_gateway/default.py

@@ -13,11 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import (
-    Any,
-    Dict,
-    Optional,
-)
+from typing import Any, Dict
 
 from loguru import logger
 import requests
@@ -26,25 +22,17 @@
 from datacustomcode.llm_gateway.base import LLMGateway
 from datacustomcode.llm_gateway.types.generate_text_request import GenerateTextRequest
 from datacustomcode.llm_gateway.types.generate_text_response import GenerateTextResponse
+from datacustomcode.llm_gateway.types.generate_text_response_builder import (
+    GenerateTextResponseBuilder,
+)
 
 
 class DefaultLLMGateway(EinsteinPlatformClient, LLMGateway):
     CONFIG_NAME = "DefaultLLMGateway"
 
-    def __init__(
-        self,
-        credentials_profile: Optional[str] = None,
-        sf_cli_org: Optional[str] = None,
-        **kwargs,
-    ):
-        EinsteinPlatformClient.__init__(
-            self, credentials_profile=credentials_profile, sf_cli_org=sf_cli_org
-        )
-        LLMGateway.__init__(self, **kwargs)
-
     def generate_text(self, request: GenerateTextRequest) -> GenerateTextResponse:
         api_url = (
-            f"{self.EINSTEIN_PLATFORM_URL}/models/{request.model_name}/generations"
+            f"{self.EINSTEIN_PLATFORM_MODELS_URL}/{request.model_name}/generations"
         )
 
         payload: Dict[str, Any] = {"prompt": request.prompt}
@@ -69,6 +57,8 @@ def generate_text(self, request: GenerateTextRequest) -> GenerateTextResponse:
             logger.error(f"Generate text request failed: {api_url} {e}")
             raise RuntimeError(f"Generate text request failed: {e}") from e
 
-        return GenerateTextResponse(
-            status_code=response.status_code, data=self.parse_response(response)
-        )
+        response_dict = {
+            "status_code": response.status_code,
+            "data": self.parse_response(response),
+        }
+        return GenerateTextResponseBuilder.build(response_dict)

src/datacustomcode/llm_gateway_config.py

@@ -19,25 +19,17 @@
     Type,
     TypeVar,
     Union,
-    cast,
 )
 
-from datacustomcode.common_config import (
-    BaseConfig,
-    BaseObjectConfig,
-    default_config_file,
-)
+from datacustomcode.common_config import BaseConfig, default_config_file
+from datacustomcode.einstein_platform_config import CredentialsObjectConfig
 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 LLMGatewayObjectConfig(CredentialsObjectConfig, Generic[_E]):
+    type_to_create: ClassVar[Type[LLMGateway]] = LLMGateway  # type: ignore[type-abstract]
 
 
 class LLMGatewayConfig(BaseConfig):

src/datacustomcode/templates/function/payload/entrypoint.py

@@ -59,7 +59,7 @@ def make_einstein_prediction(runtime: Runtime) -> None:
     )
 
     prediction_response = runtime.einstein_predictions.predict(prediction_request)
-    print(
+    logger.info(
         f"Einstein prediction results - success: [{prediction_response.is_success}] "
         f"response data: {prediction_response.data}"
     )
@@ -73,11 +73,10 @@ def generate_text(runtime: Runtime):
         .build()
     )
     llm_response = runtime.llm_gateway.generate_text(llm_request)
-
-    if llm_response.is_success:
-        print(llm_response.text)
-    else:
-        print(llm_response.error_code)
+    logger.info(
+        f"LLM Gateway generate text results - success: [{llm_response.is_success}] "
+        f"response data: {llm_response.data}"
+    )
 
 
 def function(request: dict, runtime: Runtime) -> dict: