salesforce-data-customcode 6.0.0.dev1__py3-none-any.whl → 6.0.1__py3-none-any.whl

datacustomcode/cli.py

@@ -198,6 +198,11 @@ def zip(path: str, network: str):
     default=None,
     help="SF CLI org alias or username. Fetches credentials via `sf org display`.",
 )
+@click.option(
+    "--use-in-feature",
+    default="SearchIndexChunking",
+    help="Feature where this function will be used.",
+)
 def deploy(
     path: str,
     name: str,
@@ -207,6 +212,7 @@ def deploy(
     profile: str,
     network: str,
     sf_cli_org: Optional[str],
+    use_in_feature: Optional[str],
 ):
     from datacustomcode.constants import USE_IN_FEATURE_MAPPING_FOR_CONNECT_API
     from datacustomcode.deploy import (
@@ -242,20 +248,18 @@ def deploy(
     )
 
     if package_type == "function":
-        # Infer use_in_feature from function signature
+        # Try to infer use_in_feature from function signature; fall back to
+        # the explicit flag value (defaults to SearchIndexChunking)
         entrypoint_path = os.path.join(path, ENTRYPOINT_FILE)
-        use_in_feature = infer_use_in_feature(entrypoint_path)
-        if use_in_feature:
+        inferred = infer_use_in_feature(entrypoint_path)
+        if inferred:
+            use_in_feature = inferred
             logger.info(f"Inferred use_in_feature: {use_in_feature}")
         else:
-            click.secho(
-                "Error: Could not infer function invoke options. "
-                "Please provide --use-in-feature",
-                fg="red",
-            )
-            raise click.Abort()
+            logger.info(f"Using use_in_feature: {use_in_feature}")
 
-        # Map user-provided feature names to API names
+        assert use_in_feature is not None
+        # Map feature names to Connect API names
         mapped_feature = USE_IN_FEATURE_MAPPING_FOR_CONNECT_API.get(
             use_in_feature, use_in_feature
         )

datacustomcode/einstein_platform_client.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from typing import (
     Any,
     Dict,
@@ -30,10 +31,6 @@ from datacustomcode.token_provider import (
 
 
 class EinsteinPlatformClient:
-    EINSTEIN_PLATFORM_MODELS_URL = (
-        "https://api.salesforce.com/einstein/platform/v1/models"
-    )
-
     def __init__(
         self,
         credentials_profile: Optional[str] = None,
@@ -48,8 +45,34 @@ class EinsteinPlatformClient:
             self._token_provider = CredentialsTokenProvider(profile)
             logger.debug(f"Using credentials token provider with profile: {profile}")
         self.token_response = None
+        self._einstein_url_cache: Optional[str] = None
         super().__init__(**kwargs)
 
+    def _get_einstein_platform_url(self) -> str:
+        if self._einstein_url_cache is not None:
+            return self._einstein_url_cache
+
+        env = os.environ.get("SFDC_EINSTEIN_API_ENV", "prod").lower()
+        if env not in ("dev", "test", "stage", "prod"):
+            logger.warning(
+                f"Unknown SFDC_EINSTEIN_API_ENV value '{env}', defaulting to prod"
+            )
+            env = "prod"
+
+        base_url = self._get_base_url_for_env(env)
+        logger.info(f"Using Einstein Platform API endpoint: {base_url} (env={env})")
+        self._einstein_url_cache = f"{base_url}/einstein/platform/v1/models"
+        return self._einstein_url_cache
+
+    def _get_base_url_for_env(self, env: str) -> str:
+        env_map = {
+            "dev": "https://dev.api.salesforce.com",
+            "test": "https://test.api.salesforce.com",
+            "stage": "https://stage.api.salesforce.com",
+            "prod": "https://api.salesforce.com",
+        }
+        return env_map.get(env, "https://api.salesforce.com")
+
     def _get_headers(self):
         if self.token_response is None:
             self.token_response = self._token_provider.get_token()

datacustomcode/einstein_predictions/impl/default.py

@@ -48,7 +48,7 @@ class DefaultEinsteinPredictions(EinsteinPlatformClient, EinsteinPredictions):
             )
 
         api_url = (
-            f"{self.EINSTEIN_PLATFORM_MODELS_URL}/{request.model_api_name}/{endpoint}"
+            f"{self._get_einstein_platform_url()}/{request.model_api_name}/{endpoint}"
         )
 
         prediction_columns: List[Dict[str, Any]] = []

datacustomcode/llm_gateway/default.py

@@ -29,7 +29,7 @@ class DefaultLLMGateway(EinsteinPlatformClient, LLMGateway):
 
     def generate_text(self, request: GenerateTextRequest) -> GenerateTextResponse:
         api_url = (
-            f"{self.EINSTEIN_PLATFORM_MODELS_URL}/{request.model_name}/generations"
+            f"{self._get_einstein_platform_url()}/{request.model_name}/generations"
         )
 
         payload: Dict[str, Any] = {"prompt": request.prompt}

datacustomcode/token_provider.py

@@ -101,48 +101,75 @@ class SFCLITokenProvider(TokenProvider):
 
         from datacustomcode.deploy import AccessTokenResponse
 
-        try:
-            result = subprocess.run(
-                ["sf", "org", "display", "--target-org", self.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. "
-                "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 '{self.sf_cli_org}'"
-            ) from exc
-        except subprocess.CalledProcessError as exc:
-            raise RuntimeError(
-                f"'sf org display' failed for org '{self.sf_cli_org}': {exc.stderr}"
-            ) from exc
+        def _run_sf_command(args: list[str], description: str) -> dict:
+            try:
+                result = subprocess.run(
+                    args,
+                    capture_output=True,
+                    text=True,
+                    check=True,
+                    timeout=30,
+                )
+            except FileNotFoundError as exc:
+                raise RuntimeError(
+                    "The 'sf' command was not found. "
+                    "Install Salesforce CLI: "
+                    "https://developer.salesforce.com/tools/salesforcecli"
+                ) from exc
+            except subprocess.TimeoutExpired as exc:
+                raise RuntimeError(
+                    f"'{description}' timed out for org '{self.sf_cli_org}'"
+                ) from exc
+            except subprocess.CalledProcessError as exc:
+                raise RuntimeError(
+                    f"'{description}' failed for org '{self.sf_cli_org}': "
+                    f"{exc.stderr}"
+                ) from exc
 
-        try:
-            data = json.loads(result.stdout)
-        except json.JSONDecodeError as exc:
-            raise RuntimeError(
-                f"Failed to parse JSON from 'sf org display': {result.stdout}"
-            ) from exc
+            try:
+                data = json.loads(result.stdout)
+            except json.JSONDecodeError as exc:
+                raise RuntimeError(
+                    f"Failed to parse JSON from '{description}': {result.stdout}"
+                ) from exc
 
-        if data.get("status") != 0:
+            if data.get("status") != 0:
+                raise RuntimeError(
+                    f"SF CLI error for org '{self.sf_cli_org}': "
+                    f"{data.get('message', 'unknown error')}"
+                )
+            return dict(data)
+
+        # Get instanceUrl from sf org display
+        display_data = _run_sf_command(
+            ["sf", "org", "display", "--target-org", self.sf_cli_org, "--json"],
+            "sf org display",
+        )
+        instance_url = display_data.get("result", {}).get("instanceUrl")
+        if not instance_url:
             raise RuntimeError(
-                f"SF CLI error for org '{self.sf_cli_org}': "
-                f"{data.get('message', 'unknown error')}"
+                f"'sf org display' did not return an instance URL "
+                f"for org '{self.sf_cli_org}'"
             )
 
-        result_data = data.get("result", {})
-        access_token = result_data.get("accessToken")
-        instance_url = result_data.get("instanceUrl")
-
-        if not access_token or not instance_url:
+        # Get access token via show-access-token (newer SF CLI versions
+        # redact the token in sf org display)
+        token_data = _run_sf_command(
+            [
+                "sf",
+                "org",
+                "auth",
+                "show-access-token",
+                "--target-org",
+                self.sf_cli_org,
+                "--json",
+            ],
+            "sf org auth show-access-token",
+        )
+        access_token = token_data.get("result", {}).get("accessToken")
+        if not access_token:
             raise RuntimeError(
-                f"'sf org display' did not return an access token or instance URL "
+                f"'sf org auth show-access-token' did not return an access token "
                 f"for org '{self.sf_cli_org}'"
             )
 

{salesforce_data_customcode-6.0.0.dev1.dist-info → salesforce_data_customcode-6.0.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: salesforce-data-customcode
-Version: 6.0.0.dev1
+Version: 6.0.1
 Summary: Data Cloud Custom Code SDK
 License-Expression: Apache-2.0
 License-File: LICENSE.txt
@@ -254,12 +254,18 @@ Instead of using `datacustomcode configure`, you can also set credentials via en
 | `SFDC_REFRESH_TOKEN` | OAuth refresh token |
 | `SFDC_ACCESS_TOKEN` | (Optional) OAuth core/access token |
 
+**Einstein Platform API Environment (Optional):**
+| Variable | Description |
+|----------|-------------|
+| `SFDC_EINSTEIN_API_ENV` | Einstein Platform API environment: `dev`, `test`, `stage`, or `prod`. If not set, automatically inferred from login URL. Set this explicitly if auto-detection fails. |
+
 Example usage:
 ```bash
 export SFDC_LOGIN_URL="https://login.salesforce.com"
 export SFDC_CLIENT_ID="your_client_id"
 export SFDC_CLIENT_SECRET="your_client_secret"
 export SFDC_REFRESH_TOKEN="your_refresh_token"
+export SFDC_EINSTEIN_API_ENV="test"  # optional
 
 datacustomcode run ./payload/entrypoint.py
 ```

{salesforce_data_customcode-6.0.0.dev1.dist-info → salesforce_data_customcode-6.0.1.dist-info}/RECORD

@@ -1,6 +1,6 @@
 datacustomcode/__init__.py,sha256=Fj2pIYQqLV_kQrMywrYPj0NsMtHjx4tkvwyqSt-m0E0,1530
 datacustomcode/auth.py,sha256=fpSjhIBdv9trC8yq2vuljAix_Euu-4Ah7HDCGhYjOxI,8309
-datacustomcode/cli.py,sha256=ObqhOasZSZ8ZNFP6pywV3FpsUOz45_KaZdMNT4BHOKs,13069
+datacustomcode/cli.py,sha256=tAtSSUMoVQtME_y1fxpNq4Lw1vfQsP1a6UxlEPiV9Sw,13237
 datacustomcode/client.py,sha256=TO9TkIsJjCEMlQMnD-NBGF26fk2K9jpOZjOT_XWKJ0g,9340
 datacustomcode/cmd.py,sha256=ZMs46aydJw2EaU26JgCtZmnqESQFHvvaJz10hnjZTBk,3537
 datacustomcode/common_config.py,sha256=SAUnxj3kqmOeWwPmFoYq4tuxokMgURVm4QwcGi-avL4,1928
@@ -9,11 +9,11 @@ datacustomcode/config.yaml,sha256=iMs3HPMLHwnesJ9KRD0-CpJqtHRehQvq2w2od9uE0aU,77
 datacustomcode/constants.py,sha256=SskQpUPnQE7LjKh4Z3AdZ2p191_0Kf73jV8gdBXDPLU,1436
 datacustomcode/credentials.py,sha256=D-7Zd3Nh_wStdj8_wUy5cnC8M_mdbfBijhIAi-2EbXs,9467
 datacustomcode/deploy.py,sha256=9yaekNew5ha6bd6rIFERknFROomOeIfikrKXN6jCoq4,18811
-datacustomcode/einstein_platform_client.py,sha256=84GCXNvCtC3lj-Lmqe95smIPgl7tSBomyMgxR7_IboU,3091
+datacustomcode/einstein_platform_client.py,sha256=ON2B_m--vdbCVVMVcLc81T-BRZ5-UuxVCer8E6OEQPA,4083
 datacustomcode/einstein_platform_config.py,sha256=EYnAKz0z1hNN9A8GHtaYm-Ek4GqUi68iYrJ1TSR-J9w,1416
 datacustomcode/einstein_predictions/__init__.py,sha256=o-iask36phwGA7hGi_sxJpwbITO6h_STDe0J92qD2_M,859
 datacustomcode/einstein_predictions/base.py,sha256=OBV445dpZRtCF5cpdHH49w_6RU_jZY1A-e7VtHTQCFU,1061
-datacustomcode/einstein_predictions/impl/default.py,sha256=eeD11YIrh1hqutDnqIoNPvdnNFhA0FQE35kwEsAkErs,3011
+datacustomcode/einstein_predictions/impl/default.py,sha256=fEFhrosozhG0yT87w4Gu8wbe1NgHxBqdAduwWHdWGN4,3011
 datacustomcode/einstein_predictions/types.py,sha256=W7H3sFzm_NdBtnia7O576w0OvM7jwzfGVQmFiZTO8Og,6098
 datacustomcode/einstein_predictions_config.py,sha256=dwY3Q5MFBHm60lL27m-Srt5BpcURc_PuKpaW-ALOqw4,2131
 datacustomcode/file/__init__.py,sha256=gamfOD1VnAtEslRBpqh-yKiVjkG_wYWYdSatGIfsN-w,621
@@ -39,7 +39,7 @@ datacustomcode/io/writer/csv.py,sha256=asty5teBpNQ1fMGHZ7wA3suLhq0sk0lVQPN5x7TOS
 datacustomcode/io/writer/print.py,sha256=0g2sP_1Wb95UIyEELWJt3dqM18Iv_CWhDW3mDxcIhns,5105
 datacustomcode/llm_gateway/__init__.py,sha256=xT7J9qZByD80WWk-835JcJSvQx8_NiX6qsBPhq7mU6k,800
 datacustomcode/llm_gateway/base.py,sha256=CTUhZiZ0JFlmWj6kQpTR4_-mUvqkBKl_Fkdvv9VwxY8,1262
-datacustomcode/llm_gateway/default.py,sha256=M9L-LUZweLcr_NOEG1a5ccLSFGrJ_KCXkymOgsfcHmo,1854
+datacustomcode/llm_gateway/default.py,sha256=QiwrDaUvrh00I-fYYhmogTtqC5r0fmm-JrXxNQLvTc0,1854
 datacustomcode/llm_gateway/types/__init__.py,sha256=gamfOD1VnAtEslRBpqh-yKiVjkG_wYWYdSatGIfsN-w,621
 datacustomcode/llm_gateway/types/generate_text_request.py,sha256=ChxOCzp7St1EYgwl4tWiwfrPyVci4zPyBhnEjWc1Vfs,1448
 datacustomcode/llm_gateway/types/generate_text_request_builder.py,sha256=R4hCgatMPSfWYqaOmeC1IMQj753iDgu7wU992c4tkmk,2565
@@ -90,10 +90,10 @@ datacustomcode/templates/script/payload/config.json,sha256=0d2mEMt4NHIeOUTsgiPMu
 datacustomcode/templates/script/payload/entrypoint.py,sha256=qpFA-jkhaYYF2CFNQ7_7MZVefU7SxnoMQIQYi7ox2dY,691
 datacustomcode/templates/script/requirements-dev.txt,sha256=OWwuy1awesqOZOS3Zizmdd6BDnSePpGFN5bq5IrALvc,176
 datacustomcode/templates/script/requirements.txt,sha256=qJbqzs5z0Hrx590U3T7dHq_m8UQEx2TdhgrJSz-QHIQ,40
-datacustomcode/token_provider.py,sha256=KPj_ca6mcBmhXHaPmZjAI4it8QvOWdwKJL8yfqxfYwU,5446
+datacustomcode/token_provider.py,sha256=a16r5xYeCOPR3M3JKvEV7JVIG7DOzDswObRBI33MQOk,6461
 datacustomcode/version.py,sha256=9LlbVrzwBvut1L308QbwNBgxdQk5CrghH39JARVSxms,989
-salesforce_data_customcode-6.0.0.dev1.dist-info/METADATA,sha256=8Q_ftmjeQ00-wkX2_SrRZW8ar5qEBtHcbDjc7rHXer8,19818
-salesforce_data_customcode-6.0.0.dev1.dist-info/WHEEL,sha256=eY7nduwzv-ldUxpzbRlxwvC693Hg6PX8bWDjEHjZ_dk,88
-salesforce_data_customcode-6.0.0.dev1.dist-info/entry_points.txt,sha256=WpQ94UB7UuRCYGOLtJV3vAgm1tl7iVf43VYRWIuNu5Y,74
-salesforce_data_customcode-6.0.0.dev1.dist-info/licenses/LICENSE.txt,sha256=iOi8EmQpfkFhMENi7VYtkl2EqS14pEOccoXHiW2dyPU,11443
-salesforce_data_customcode-6.0.0.dev1.dist-info/RECORD,,
+salesforce_data_customcode-6.0.1.dist-info/METADATA,sha256=PVoEUlTf9pjQlJq8Yi3Mt7migPd0K16ORl80T2gsMbc,20162
+salesforce_data_customcode-6.0.1.dist-info/WHEEL,sha256=eY7nduwzv-ldUxpzbRlxwvC693Hg6PX8bWDjEHjZ_dk,88
+salesforce_data_customcode-6.0.1.dist-info/entry_points.txt,sha256=WpQ94UB7UuRCYGOLtJV3vAgm1tl7iVf43VYRWIuNu5Y,74
+salesforce_data_customcode-6.0.1.dist-info/licenses/LICENSE.txt,sha256=iOi8EmQpfkFhMENi7VYtkl2EqS14pEOccoXHiW2dyPU,11443
+salesforce_data_customcode-6.0.1.dist-info/RECORD,,