test Einstein Prediction with actual production API

Author: atulikumwenayo

src/datacustomcode/cli.py

@@ -179,14 +179,15 @@ def deploy(
     function_invoke_opt: str,
     sf_cli_org: Optional[str],
 ):
-    from datacustomcode.credentials import Credentials
     from datacustomcode.deploy import (
         COMPUTE_TYPES,
-        AccessTokenResponse,
         CodeExtensionMetadata,
-        _retrieve_access_token_from_sf_cli,
         deploy_full,
     )
+    from datacustomcode.token_provider import (
+        CredentialsTokenProvider,
+        SFCLITokenProvider,
+    )
 
     logger.debug("Deploying project")
 
@@ -220,22 +221,15 @@ def deploy(
             function_invoke_options = function_invoke_opt.split(",")
             metadata.functionInvokeOptions = function_invoke_options
 
-    auth: Union[Credentials, AccessTokenResponse]
-    if sf_cli_org:
-        try:
-            auth = _retrieve_access_token_from_sf_cli(sf_cli_org)
-        except RuntimeError as e:
-            click.secho(f"Error: {e}", fg="red")
-            raise click.Abort() from None
-    else:
-        try:
-            auth = Credentials.from_available(profile=profile)
-        except ValueError as e:
-            click.secho(
-                f"Error: {e}",
-                fg="red",
-            )
-            raise click.Abort() from None
+    try:
+        if sf_cli_org:
+            auth = SFCLITokenProvider(sf_cli_org).get_token()
+        else:
+            auth = CredentialsTokenProvider(profile).get_token()
+    except RuntimeError as e:
+        click.secho(f"Error: {e}", fg="red")
+        raise click.Abort() from None
+
     deploy_full(path, metadata, auth, network)
 
 

src/datacustomcode/config.yaml

@@ -26,4 +26,10 @@ proxy_config:
 
 einstein_predictions_config:
   type_config_name: DefaultEinsteinPredictions
-  options: {}
+  options:
+    credentials_profile: default
+
+llm_gateway_config:
+  type_config_name: DefaultLLMGateway
+  options:
+    credentials_profile: default

src/datacustomcode/deploy.py

@@ -19,11 +19,9 @@
 import os
 import re
 import shutil
-import subprocess
 import tempfile
 import time
 from typing import (
-    TYPE_CHECKING,
     Any,
     Callable,
     Dict,
@@ -37,15 +35,10 @@
 import requests
 
 from datacustomcode.cmd import cmd_output
-from datacustomcode.credentials import AuthType
 from datacustomcode.scan import find_base_directory, get_package_type
 
-if TYPE_CHECKING:
-    from datacustomcode.credentials import Credentials
-
 DATA_CUSTOM_CODE_PATH = "services/data/v63.0/ssot/data-custom-code"
 DATA_TRANSFORMS_PATH = "services/data/v63.0/ssot/data-transforms"
-AUTH_PATH = "services/oauth2/token"
 WAIT_FOR_DEPLOYMENT_TIMEOUT = 3000
 
 # Available compute types for Data Cloud deployments.
@@ -163,80 +156,6 @@ class AccessTokenResponse(BaseModel):
     instance_url: str
 
 
-def _retrieve_access_token(credentials: Credentials) -> AccessTokenResponse:
-    """Get an access token for the Salesforce API."""
-    logger.debug("Getting oauth token...")
-
-    url = f"{credentials.login_url.rstrip('/')}/{AUTH_PATH.lstrip('/')}"
-
-    if credentials.auth_type == AuthType.OAUTH_TOKENS:
-        data = {
-            "grant_type": "refresh_token",
-            "refresh_token": credentials.refresh_token,
-            "client_id": credentials.client_id,
-            "client_secret": credentials.client_secret,
-        }
-    elif credentials.auth_type == AuthType.CLIENT_CREDENTIALS:
-        data = {
-            "grant_type": "client_credentials",
-            "client_id": credentials.client_id,
-            "client_secret": credentials.client_secret,
-        }
-    else:
-        raise ValueError(f"Unsupported auth_type: {credentials.auth_type}")
-
-    response = _make_api_call(url, "POST", data=data)
-    return AccessTokenResponse(**response)
-
-
-def _retrieve_access_token_from_sf_cli(sf_cli_org: str) -> AccessTokenResponse:
-    """Get an access token from the Salesforce CLI."""
-    try:
-        result = subprocess.run(
-            ["sf", "org", "display", "--target-org", sf_cli_org, "--json"],
-            capture_output=True,
-            text=True,
-            check=True,
-            timeout=30,
-        )
-    except FileNotFoundError as exc:
-        raise RuntimeError(
-            "The 'sf' command was not found. "
-            "Please install Salesforce CLI: https://developer.salesforce.com/tools/salesforcecli"
-        ) from exc
-    except subprocess.TimeoutExpired as exc:
-        raise RuntimeError(
-            f"'sf org display' timed out for org '{sf_cli_org}'"
-        ) from exc
-    except subprocess.CalledProcessError as exc:
-        raise RuntimeError(
-            f"'sf org display' failed for org '{sf_cli_org}'.\n"
-            f"Ensure the org is authenticated via 'sf org login web'.\n"
-            f"stderr: {exc.stderr.strip()}"
-        ) from exc
-
-    try:
-        data = json.loads(result.stdout)
-    except json.JSONDecodeError as exc:
-        raise RuntimeError(f"Failed to parse 'sf org display' output: {exc}") from exc
-
-    if data.get("status") != 0:
-        raise RuntimeError(
-            f"SF CLI error for org '{sf_cli_org}': "
-            f"{data.get('message', 'unknown error')}"
-        )
-
-    org_result = data.get("result", {})
-    access_token = org_result.get("accessToken")
-    instance_url = org_result.get("instanceUrl")
-    if not access_token or not instance_url:
-        raise RuntimeError(
-            f"'sf org display' did not return an access token or instance URL "
-            f"for org '{sf_cli_org}'"
-        )
-    return AccessTokenResponse(access_token=access_token, instance_url=instance_url)
-
-
 class CreateDeploymentResponse(BaseModel):
     fileUploadUrl: str
 
@@ -567,16 +486,11 @@ def zip(
 def deploy_full(
     directory: str,
     metadata: CodeExtensionMetadata,
-    credentials: Union["Credentials", AccessTokenResponse],
+    access_token: AccessTokenResponse,
     docker_network: str,
     callback=None,
 ) -> AccessTokenResponse:
     """Deploy a data transform in the DataCloud."""
-    if isinstance(credentials, AccessTokenResponse):
-        access_token = credentials
-    else:
-        access_token = _retrieve_access_token(credentials)
-
     # prepare payload
     config = get_config(directory)
 
@@ -587,7 +501,6 @@ def deploy_full(
     wait_for_deployment(access_token, metadata, callback)
 
     # create data transform
-
     if isinstance(config, DataTransformConfig):
         create_data_transform(directory, access_token, metadata, config)
     return access_token

src/datacustomcode/einstein_predictions/impl/default.py

@@ -13,23 +13,131 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import (
+    Any,
+    ClassVar,
+    Dict,
+    List,
+    Optional,
+)
+
+from loguru import logger
+import requests
+
 from datacustomcode.einstein_predictions.base import EinsteinPredictions
 from datacustomcode.einstein_predictions.types import (
     PredictionRequest,
     PredictionResponse,
+    PredictionType,
+)
+from datacustomcode.token_provider import (
+    CredentialsTokenProvider,
+    SFCLITokenProvider,
+    TokenProvider,
 )
 
 
 class DefaultEinsteinPredictions(EinsteinPredictions):
     CONFIG_NAME = "DefaultEinsteinPredictions"
+    EINSTEIN_PLATFORM_URL = "https://api.salesforce.com/einstein/platform/v1"
 
-    def __init__(self, **kwargs):
+    ENDPOINT_MAP: ClassVar[dict[PredictionType, str]] = {
+        PredictionType.REGRESSION: "regression",
+        PredictionType.CLUSTERING: "clustering",
+        PredictionType.CLASSIFICATION: "classification",
+        PredictionType.BINARY_CLASSIFICATION: "binary-classification",
+        PredictionType.MULTI_OUTCOME: "multi-outcome",
+    }
+
+    def __init__(
+        self,
+        credentials_profile: Optional[str] = None,
+        sf_cli_org: Optional[str] = None,
+        **kwargs,
+    ):
         super().__init__(**kwargs)
 
+        if sf_cli_org:
+            self._token_provider: TokenProvider = SFCLITokenProvider(sf_cli_org)
+            logger.debug(f"Using SF CLI token provider for org: {sf_cli_org}")
+        else:
+            profile = credentials_profile or "default"
+            self._token_provider = CredentialsTokenProvider(profile)
+            logger.debug(f"Using credentials token provider with profile: {profile}")
+
     def predict(self, request: PredictionRequest) -> PredictionResponse:
+        """Make a prediction request to the Einstein Predictions API"""
+        token_response = self._token_provider.get_token()
+        access_token = token_response.access_token
+
+        endpoint = self.ENDPOINT_MAP.get(request.prediction_type)
+        if not endpoint:
+            raise RuntimeError(
+                f"Unknown prediction type: {request.prediction_type}. "
+                f"Valid types: {list(self.ENDPOINT_MAP.keys())}"
+            )
+
+        api_url = (
+            f"{self.EINSTEIN_PLATFORM_URL}/models/"
+            f"{request.model_api_name}/{endpoint}"
+        )
+
+        prediction_columns: List[Dict[str, Any]] = []
+        for col in request.prediction_columns:
+            col_data: Dict[str, Any] = {"columnName": col.column_name}
+            if col.string_values:
+                col_data["stringValues"] = col.string_values
+            if col.double_values:
+                col_data["doubleValues"] = col.double_values
+            if col.boolean_values:
+                col_data["booleanValues"] = col.boolean_values
+            if col.date_values:
+                col_data["dateValues"] = col.date_values
+            if col.datetime_values:
+                col_data["datetimeValues"] = col.datetime_values
+            prediction_columns.append(col_data)
+
+        payload: Dict[str, Any] = {"predictionColumns": prediction_columns}
+
+        if request.settings:
+            payload["settings"] = request.settings
+
+        headers = {
+            "Authorization": f"Bearer {access_token}",
+            "Content-Type": "application/json",
+            "x-sfdc-app-context": "EinsteinGPT",
+            "x-client-feature-id": "ai-platform-models-connected-app",
+        }
+
+        logger.debug(f"Making Einstein prediction request to: {api_url}")
+        try:
+            response = requests.post(api_url, json=payload, headers=headers, timeout=60)
+            if not response.ok and not response.text:
+                error_msg = (
+                    f"Einstein Prediction request failed: {api_url} - "
+                    f"{response.status_code} {response.reason}. "
+                    "If your code uses Einstein APIs, make sure you have "
+                    'configured the SDK to use "client_credentials" auth type. '
+                    "Refer to https://developer.salesforce.com/docs/ai/agentforce/"
+                    "guide/agent-api-get-started.html#create-a-salesforce-app "
+                    "to create your external client app."
+                )
+                logger.error(error_msg)
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Prediction API request failed: {api_url} {e}")
+            raise RuntimeError(f"Prediction API request failed: {e}") from e
+
+        response_data: Dict[str, Any] = {}
+        if response.content:
+            try:
+                response_data = response.json()
+            except ValueError:
+                logger.warning("Failed to parse response as JSON")
+                response_data = {"raw_response": response.text}
+
         return PredictionResponse(
             version="v1",
             prediction_type=request.prediction_type,
-            status_code=200,
-            data={"results": [{"prediction": {"predictedValue": 1.0}}]},
+            status_code=response.status_code,
+            data=response_data,
         )