salesforce-data-customcode 6.0.2.dev1__tar.gz → 6.0.3__tar.gz

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: salesforce-data-customcode
-Version: 6.0.2.dev1
+Version: 6.0.3
 Summary: Data Cloud Custom Code SDK
 License-Expression: Apache-2.0
 License-File: LICENSE.txt
@@ -15,6 +15,7 @@ Requires-Dist: click (>=8.1.8,<9.0.0)
 Requires-Dist: loguru (>=0.7.3,<0.8.0)
 Requires-Dist: numpy
 Requires-Dist: pandas
+Requires-Dist: pipreqs
 Requires-Dist: pydantic (==2.13.1)
 Requires-Dist: pyspark (==3.5.1)
 Requires-Dist: pyyaml (>=6.0,<7.0)
@@ -341,14 +342,16 @@ from datacustomcode.client import Client, llm_gateway_generate_text_col
 def main():
   client = Client()
   df = client.read_dlo("Input__dll")
+  # llm_gateway_generate_text_col returns a struct
+  # {status, response, error_code, error_message} per row, so per-row
+  # failures don't abort the Spark job. Pick the field you want with [].
   df_generated = df.withColumn(
     "greeting__c",
     llm_gateway_generate_text_col(
         "In one sentence, greet {name} from {city}.",
         {"name": col("name__c"), "city": col("homecity__c")},
         model_id="sfdc_ai__DefaultGPT4Omni", # An AI model in your org
-        max_tokens=100,
-    ),
+    )["response"],
   )
 
   dlo_name = "Output_dll"

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/README.md

@@ -316,14 +316,16 @@ from datacustomcode.client import Client, llm_gateway_generate_text_col
 def main():
   client = Client()
   df = client.read_dlo("Input__dll")
+  # llm_gateway_generate_text_col returns a struct
+  # {status, response, error_code, error_message} per row, so per-row
+  # failures don't abort the Spark job. Pick the field you want with [].
   df_generated = df.withColumn(
     "greeting__c",
     llm_gateway_generate_text_col(
         "In one sentence, greet {name} from {city}.",
         {"name": col("name__c"), "city": col("homecity__c")},
         model_id="sfdc_ai__DefaultGPT4Omni", # An AI model in your org
-        max_tokens=100,
-    ),
+    )["response"],
   )
 
   dlo_name = "Output_dll"

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/pyproject.toml

@@ -18,7 +18,7 @@ license = "Apache-2.0"
 name = "salesforce-data-customcode"
 readme = "README.md"
 requires-python = ">=3.10,<3.12"
-version = "6.0.2.dev1"
+version = "6.0.3"
 
 [tool.black]
 exclude = '''
@@ -101,6 +101,7 @@ click = "^8.1.8"
 loguru = "^0.7.3"
 numpy = "*"
 pandas = "*"
+pipreqs = "*"
 pydantic = "2.13.1"
 pyspark = "3.5.1"
 python = ">=3.10,<3.12"

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/cli.py

@@ -198,6 +198,7 @@ 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=None, hidden=True, deprecated=True)
 def deploy(
     path: str,
     name: str,
@@ -207,6 +208,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 (
@@ -249,13 +251,14 @@ def deploy(
             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",
+                "Error: Function signature does not match a supported type. "
+                "Use SearchIndexChunkingV1Request and "
+                "SearchIndexChunkingV1Response in function signature.",
                 fg="red",
             )
             raise click.Abort()
 
-        # Map user-provided feature names to API names
+        # Map feature names to Connect API names
         mapped_feature = USE_IN_FEATURE_MAPPING_FOR_CONNECT_API.get(
             use_in_feature, use_in_feature
         )

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/client.py

@@ -59,20 +59,26 @@ def llm_gateway_generate_text_col(
     template: str,
     values: Union[Dict[str, "Column"], "Column"],
     model_id: Optional[str] = None,
-    max_tokens: Optional[int] = None,
 ) -> "Column":
     """Build a Spark Column that runs the LLM Gateway per row.
 
+    The returned Column yields a struct ``{status, response, error_code,
+    error_message}`` for each row. Use ``[...]`` (or ``getField``) to pick the
+    field you want, e.g. ``llm_gateway_generate_text_col(...)["response"]``.
+    Per-row failures populate ``status`` / ``error_code`` / ``error_message``
+    so a single bad row does not abort the whole Spark job.
+
     Example:
 
-        >>> df.withColumn(
-        ...     "greeting__c",
-        ...     llm_gateway_generate_text_col(
-        ...         "In one sentence, greet {name} from {city}.",
-        ...         {"name": col("name__c"), "city": col("homecity__c")},
-        ...         model_id="sfdc_ai__DefaultGPT4Omni",
-        ...         max_tokens=100,
-        ...     ),
+        >>> result = llm_gateway_generate_text_col(
+        ...     "In one sentence, greet {name} from {city}.",
+        ...     {"name": col("name__c"), "city": col("homecity__c")},
+        ...     model_id="sfdc_ai__DefaultGPT4Omni",
+        ... )
+        >>> df.withColumn("greeting__c", result["response"])
+        >>> # …or keep the struct around and inspect failures:
+        >>> df.withColumn("llm", result).select(
+        ...     "llm.status", "llm.response", "llm.error_message"
         ... )
 
     Args:
@@ -81,15 +87,16 @@ def llm_gateway_generate_text_col(
         values: Either a mapping from placeholder name to Spark ``Column``, or
             a single ``Column`` whose value is already a struct.
         model_id: LLM model id. Defaults to ``sfdc_ai__DefaultGPT4Omni``.
-        max_tokens: Maximum tokens to generate. Defaults to 200.
 
     Returns:
-        A Spark ``Column`` that, when evaluated, produces the generated text.
+        A Spark ``Column`` of ``StructType`` with fields ``status``,
+        ``response``, ``error_code``, and ``error_message`` (all nullable
+        strings). On success, ``status == "SUCCESS"`` and ``response`` holds
+        the generated text; on failure, ``status == "ERROR"`` and the
+        ``error_*`` fields carry diagnostic detail.
     """
     gateway = Client()._get_spark_llm_gateway()
-    return gateway.llm_gateway_generate_text_col(
-        template, values, model_id=model_id, max_tokens=max_tokens
-    )
+    return gateway.llm_gateway_generate_text_col(template, values, model_id=model_id)
 
 
 class DataCloudObjectType(Enum):
@@ -150,9 +157,7 @@ class Client:
         finder: Find a file path
         reader: A custom reader to use for reading Data Cloud objects.
         writer: A custom writer to use for writing Data Cloud objects.
-        spark_llm_gateway: Optional custom :class:`SparkLLMGateway`. When
-            omitted, the gateway is lazily resolved from
-            ``spark_llm_gateway_config``.
+        spark_llm_gateway: Optional custom :class:`SparkLLMGateway`.
 
     Example:
     >>> client = Client()
@@ -292,7 +297,6 @@ class Client:
         self,
         prompt: str,
         model_id: Optional[str] = None,
-        max_tokens: Optional[int] = None,
     ) -> str:
         """Issue a one-shot LLM Gateway call. This is the scalar counterpart to
         :func:`llm_gateway_generate_text_col`: it runs **once**  — not per row.
@@ -310,15 +314,13 @@ class Client:
                 ``{field}`` substitution is performed on this string.
             model_id: LLM model id to target. Defaults to
                 ``sfdc_ai__DefaultGPT4Omni`` when ``None``.
-            max_tokens: Hard upper bound on the number of tokens the model
-                may generate. Defaults to 200 when ``None``.
 
         Returns:
             The generated text as a plain Python ``str``; empty when the
             gateway response carries no generated text.
         """
         return self._get_spark_llm_gateway().llm_gateway_generate_text(
-            prompt, model_id=model_id, max_tokens=max_tokens
+            prompt, model_id=model_id
         )
 
     def _get_spark_llm_gateway(self) -> SparkLLMGateway:

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/deploy.py

@@ -26,12 +26,13 @@ from typing import (
     Callable,
     Dict,
     List,
+    Optional,
     Union,
 )
 
 from loguru import logger
 import pydantic
-from pydantic import BaseModel
+from pydantic import BaseModel, model_validator
 import requests
 
 from datacustomcode.cmd import cmd_output
@@ -366,23 +367,86 @@ class BaseConfig(BaseModel):
     entryPoint: str
 
 
+class DataObjectField(BaseModel):
+    name: str
+    label: str
+    dataType: str
+    isPrimaryKey: bool = False
+    keyQualifierFieldName: Optional[str] = None
+
+
+class DataObject(BaseModel):
+    name: str
+    label: str
+    type: str
+    category: str
+    fields: list[DataObjectField]
+
+
 class DataTransformConfig(BaseConfig):
     sdkVersion: str
     dataspace: str
     permissions: Permissions
+    dataObjects: Optional[list[DataObject]] = None
 
 
 class FunctionConfig(BaseConfig):
     pass
 
 
+class DloPermission(BaseModel):
+    dlo: list[str]
+
+
+class DmoPermission(BaseModel):
+    dmo: list[str]
+
+
 class Permissions(BaseModel):
-    read: Union[DloPermission]
-    write: Union[DloPermission]
+    read: Union[DloPermission, DmoPermission]
+    write: Union[DloPermission, DmoPermission]
+
+    @model_validator(mode="after")
+    def _no_mixed_layers(self) -> "Permissions":
+        read_is_dlo = isinstance(self.read, DloPermission)
+        write_is_dlo = isinstance(self.write, DloPermission)
+        if read_is_dlo != write_is_dlo:
+            raise ValueError(
+                "permissions.read and permissions.write must both reference "
+                "DLOs or both reference DMOs (got "
+                f"read={type(self.read).__name__}, "
+                f"write={type(self.write).__name__})"
+            )
+        return self
 
 
-class DloPermission(BaseModel):
-    dlo: list[str]
+def _permission_entries(perm: Union[DloPermission, DmoPermission]) -> list[str]:
+    """Return the list of object names regardless of layer (DLO or DMO)."""
+    if isinstance(perm, DloPermission):
+        return perm.dlo
+    return perm.dmo
+
+
+def _data_object_to_output(obj: DataObject) -> dict[str, Any]:
+    """Convert a config.json DataObject into an outputDataObjects entry."""
+    fields: list[dict[str, Any]] = []
+    for field in obj.fields:
+        entry: dict[str, Any] = {
+            "isPrimaryKey": field.isPrimaryKey,
+            "label": field.label,
+            "name": field.name,
+            "type": field.dataType,
+        }
+        if field.keyQualifierFieldName is not None:
+            entry["keyQualifierField"] = field.keyQualifierFieldName
+        fields.append(entry)
+    return {
+        "category": obj.category,
+        "fields": fields,
+        "label": obj.label,
+        "name": obj.name,
+        "type": obj.type,
+    }
 
 
 def get_config(directory: str) -> BaseConfig:
@@ -404,10 +468,17 @@ def get_config(directory: str) -> BaseConfig:
     except json.JSONDecodeError as err:
         raise ValueError(f"config.json at {config_path} is not valid JSON") from err
     except pydantic.ValidationError as err:
-        missing_fields = [str(err["loc"][0]) for err in err.errors()]
+        errors = err.errors()
+        missing = [e for e in errors if e.get("type") == "missing"]
+        if missing and len(missing) == len(errors):
+            missing_fields = [str(e["loc"][0]) for e in missing]
+            raise ValueError(
+                f"config.json at {config_path} is missing required "
+                f"fields: {', '.join(missing_fields)}"
+            ) from err
+        messages = [str(e.get("msg", "")) for e in errors]
         raise ValueError(
-            f"config.json at {config_path} is missing required "
-            f"fields: {', '.join(missing_fields)}"
+            f"config.json at {config_path} is invalid: {'; '.join(messages)}"
         ) from err
 
 
@@ -421,26 +492,45 @@ def create_data_transform(
     script_name = metadata.name
     request_hydrated = DATA_TRANSFORM_REQUEST_TEMPLATE.copy()
 
-    # Add nodes for each write DLO
-    for i, dlo in enumerate(data_transform_config.permissions.write.dlo, 1):
+    # Add nodes for each write entry (DLO or DMO)
+    for i, name in enumerate(
+        _permission_entries(data_transform_config.permissions.write), 1
+    ):
         request_hydrated["nodes"][f"node{i}"] = {
-            "relation_name": dlo,
+            "relation_name": name,
             "config": {"materialized": "table"},
             "compiled_code": "",
         }
 
-    # Add sources for each read DLO
-    for i, dlo in enumerate(data_transform_config.permissions.read.dlo, 1):
-        request_hydrated["sources"][f"source{i}"] = {"relation_name": dlo}
+    # Add sources for each read entry (DLO or DMO)
+    for i, name in enumerate(
+        _permission_entries(data_transform_config.permissions.read), 1
+    ):
+        request_hydrated["sources"][f"source{i}"] = {"relation_name": name}
 
     request_hydrated["macros"]["macro.byoc"]["arguments"][0]["name"] = script_name
 
+    definition: dict[str, Any] = {
+        "type": "DCSQL",
+        "manifest": request_hydrated,
+        "version": "56.0",
+    }
+
+    # outputDataObjects is only set for DMO-backed transforms. The server requires
+    # the schema of any DMO created/updated by the transform; DLO transforms use
+    # an existing materialized table and must not include this field.
+    if isinstance(data_transform_config.permissions.write, DmoPermission):
+        if not data_transform_config.dataObjects:
+            raise ValueError(
+                "DMO transforms require 'dataObjects' in config.json describing "
+                "the schema of each output DMO."
+            )
+        definition["outputDataObjects"] = [
+            _data_object_to_output(obj) for obj in data_transform_config.dataObjects
+        ]
+
     body = {
-        "definition": {
-            "type": "DCSQL",
-            "manifest": request_hydrated,
-            "version": "56.0",
-        },
+        "definition": definition,
         "label": f"{metadata.name}",
         "name": f"{metadata.name}",
         "type": "BATCH",

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/io/writer/base.py

@@ -27,11 +27,21 @@ if TYPE_CHECKING:
 class WriteMode(str, Enum):
     APPEND = "append"
     OVERWRITE = "overwrite"
-    OVERWRITE_PARTITIONS = "overwrite_partitions"
+    OVERWRITE_PARTITIONS = "overwrite_partitions"  # Deprecated: raises error if used
     MERGE = "merge"
     MERGE_UPSERT_DELETE = "merge_upsert_delete"
 
 
+class MergeRecordType(str, Enum):
+    """Values for the _merge_record_type column used by MERGE_UPSERT_DELETE."""
+
+    UPSERT = "UPSERT"
+    DELETE = "DELETE"
+
+
+MERGE_RECORD_TYPE_COLUMN = "_merge_record_type"
+
+
 class BaseDataCloudWriter(BaseDataAccessLayer):
     """Base class for Data Cloud writers."""
 

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/llm_gateway/__init__.py

@@ -15,6 +15,7 @@
 
 from datacustomcode.llm_gateway.base import LLMGateway
 from datacustomcode.llm_gateway.default import DefaultLLMGateway
+from datacustomcode.llm_gateway.errors import LLMGatewayCallError
 from datacustomcode.llm_gateway.spark_base import SparkLLMGateway
 from datacustomcode.llm_gateway.spark_default import DefaultSparkLLMGateway
 
@@ -22,5 +23,6 @@ __all__ = [
     "DefaultLLMGateway",
     "DefaultSparkLLMGateway",
     "LLMGateway",
+    "LLMGatewayCallError",
     "SparkLLMGateway",
 ]

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/llm_gateway/default.py

@@ -34,8 +34,6 @@ class DefaultLLMGateway(EinsteinPlatformClient, LLMGateway):
 
         payload: Dict[str, Any] = {"prompt": request.prompt}
 
-        if request.max_tokens is not None:
-            payload["max_tokens"] = request.max_tokens
         if request.localization:
             payload["localization"] = request.localization
         if request.tags:

salesforce_data_customcode-6.0.3/src/datacustomcode/llm_gateway/errors.py

@@ -0,0 +1,36 @@
+# 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.
+"""Exceptions raised by LLM Gateway implementations."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+
+class LLMGatewayCallError(RuntimeError):
+    """Raised when an LLM Gateway call returns an error."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        status: Optional[object] = None,
+        error_code: Optional[str] = None,
+        error_message: Optional[str] = None,
+    ) -> None:
+        super().__init__(message)
+        self.status = status
+        self.error_code = error_code
+        self.error_message = error_message

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/llm_gateway/spark_base.py

@@ -40,7 +40,6 @@ class SparkLLMGateway(ABC, UserExtendableNamedConfigMixin):
         self,
         prompt: str,
         model_id: Optional[str] = None,
-        max_tokens: Optional[int] = None,
     ) -> str:
         """Issue a one-shot LLM Gateway call and return the generated text."""
 
@@ -50,6 +49,14 @@ class SparkLLMGateway(ABC, UserExtendableNamedConfigMixin):
         template: str,
         values: Union[Dict[str, "Column"], "Column"],
         model_id: Optional[str] = None,
-        max_tokens: Optional[int] = None,
     ) -> "Column":
-        """Build a Spark ``Column`` that invokes the LLM Gateway per row."""
+        """Build a Spark ``Column`` that invokes the LLM Gateway per row and
+        yields a struct ``{status, response, error_code, error_message}``.
+
+        Select an individual field, e.g.
+        ``llm_gateway_generate_text_col(...)["response"]``. Returning a struct
+        means a single failing row doesn't abort the Spark job.
+        Failing row leaves the rest of the DataFrame intact — callers can
+        inspect ``status`` / ``error_code`` per row instead of having the
+        Spark job abort.
+        """

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/llm_gateway/spark_default.py

@@ -28,10 +28,15 @@ if TYPE_CHECKING:
     from pyspark.sql import Column
 
     from datacustomcode.llm_gateway.base import LLMGateway
+    from datacustomcode.llm_gateway.types.generate_text_response import (
+        GenerateTextResponse,
+    )
 
 
 _DEFAULT_LLM_MODEL_ID = "sfdc_ai__DefaultGPT4Omni"
-_DEFAULT_LLM_MAX_TOKENS = 200
+
+_STATUS_SUCCESS = "SUCCESS"
+_STATUS_ERROR = "ERROR"
 
 
 class DefaultSparkLLMGateway(SparkLLMGateway):
@@ -52,20 +57,26 @@ class DefaultSparkLLMGateway(SparkLLMGateway):
         self,
         prompt: str,
         model_id: Optional[str] = None,
-        max_tokens: Optional[int] = None,
     ) -> str:
-        return _invoke_llm_gateway(self._llm_gateway, prompt, model_id, max_tokens)
+        return _invoke_llm_gateway(self._llm_gateway, prompt, model_id)
 
     def llm_gateway_generate_text_col(
         self,
         template: str,
         values: Union[Dict[str, "Column"], "Column"],
         model_id: Optional[str] = None,
-        max_tokens: Optional[int] = None,
     ) -> "Column":
-
+        """Build a per-row UDF that returns a struct ``{status, response,
+        error_code, error_message}`` so per-row failures do not abort the
+        Spark job. Callers select the field they want, e.g.
+        ``llm_gateway_generate_text_col(...)["response"]``.
+        """
         from pyspark.sql.functions import struct, udf
-        from pyspark.sql.types import StringType
+        from pyspark.sql.types import (
+            StringType,
+            StructField,
+            StructType,
+        )
 
         if isinstance(values, dict):
             values_col = struct(*[v.alias(k) for k, v in values.items()])
@@ -73,19 +84,32 @@ class DefaultSparkLLMGateway(SparkLLMGateway):
             values_col = values
 
         gateway = self._llm_gateway
+        result_schema = StructType(
+            [
+                StructField("status", StringType(), True),
+                StructField("response", StringType(), True),
+                StructField("error_code", StringType(), True),
+                StructField("error_message", StringType(), True),
+            ]
+        )
 
-        def _generate(values_row: Any) -> str:
+        def _generate(values_row: Any) -> Dict[str, Optional[str]]:
             if values_row is None:
-                return ""
+                return {
+                    "status": _STATUS_ERROR,
+                    "response": None,
+                    "error_code": None,
+                    "error_message": "values column was null for this row",
+                }
             subs = (
                 values_row.asDict()
                 if hasattr(values_row, "asDict")
                 else dict(values_row)
             )
             prompt = template.format(**subs)
-            return _invoke_llm_gateway(gateway, prompt, model_id, max_tokens)
+            return _invoke_llm_gateway_as_struct(gateway, prompt, model_id)
 
-        return udf(_generate, StringType())(values_col)
+        return udf(_generate, result_schema)(values_col)
 
 
 def _build_underlying_gateway() -> "LLMGateway":
@@ -100,20 +124,61 @@ def _build_underlying_gateway() -> "LLMGateway":
     return cfg.to_object()
 
 
-def _invoke_llm_gateway(
+def _call_llm_gateway(
     gateway: "LLMGateway",
     prompt: str,
     model_id: Optional[str],
-    max_tokens: Optional[int],
-) -> str:
+) -> "GenerateTextResponse":
+    """Build the request and dispatch it to the underlying gateway."""
     from datacustomcode.llm_gateway.types.generate_text_request_builder import (
         GenerateTextRequestBuilder,
     )
 
-    builder = (
+    request = (
         GenerateTextRequestBuilder()
         .set_prompt(prompt)
         .set_model(model_id or _DEFAULT_LLM_MODEL_ID)
-        .set_max_tokens(max_tokens or _DEFAULT_LLM_MAX_TOKENS)
+        .build()
     )
-    return gateway.generate_text(builder.build()).text
+    return gateway.generate_text(request)
+
+
+def _invoke_llm_gateway(
+    gateway: "LLMGateway",
+    prompt: str,
+    model_id: Optional[str],
+) -> str:
+    from datacustomcode.llm_gateway.errors import LLMGatewayCallError
+
+    response = _call_llm_gateway(gateway, prompt, model_id)
+    if response.is_error:
+        raise LLMGatewayCallError(
+            f"LLM Gateway call failed: status_code={response.status_code}, "
+            f"error_code={response.error_code!r}, "
+            f"message={response.data!r}",
+            status=response.status_code,
+            error_code=response.error_code or None,
+            error_message=str(response.data) if response.data else None,
+        )
+    return response.text
+
+
+def _invoke_llm_gateway_as_struct(
+    gateway: "LLMGateway",
+    prompt: str,
+    model_id: Optional[str],
+) -> Dict[str, Optional[str]]:
+    response = _call_llm_gateway(gateway, prompt, model_id)
+    if response.is_error:
+        return {
+            "status": _STATUS_ERROR,
+            "response": None,
+            "error_code": response.error_code or None,
+            "error_message": str(response.data) if response.data else None,
+        }
+    return {
+        "status": _STATUS_SUCCESS,
+        "response": response.text,
+        "error_code": None,
+        "error_message": None,
+    }

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/llm_gateway/types/generate_text_request.py

@@ -40,13 +40,6 @@ class GenerateTextRequest(BaseModel):
     )
     model_name: str = Field(..., min_length=1, description="Name of the model to use")
     prompt: str = Field(..., description="Input prompt")
-    max_tokens: Optional[int] = Field(
-        default=None,
-        ge=1,
-        description=(
-            "Maximum number of tokens to generate. If None, server default applies."
-        ),
-    )
     localization: Optional[Dict[str, Any]] = Field(
         default=None, description="Localization settings"
     )

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/llm_gateway/types/generate_text_request_builder.py

@@ -26,7 +26,6 @@ class GenerateTextRequestBuilder:
     def __init__(self) -> None:
         self._prompt = ""
         self._model_name = ""
-        self._max_tokens: Optional[int] = None
         self._localization: Optional[Dict[str, Any]] = None
         self._tags: Optional[Dict[str, Any]] = None
 
@@ -38,10 +37,6 @@ class GenerateTextRequestBuilder:
         self._model_name = model_name
         return self
 
-    def set_max_tokens(self, max_tokens: int) -> "GenerateTextRequestBuilder":
-        self._max_tokens = max_tokens
-        return self
-
     def set_localization(
         self,
         localization: Optional[Dict[str, Any]] = None,
@@ -80,7 +75,6 @@ class GenerateTextRequestBuilder:
         request = GenerateTextRequest(
             prompt=self._prompt,
             model_name=self._model_name,
-            max_tokens=self._max_tokens,
             localization=self._localization,
             tags=self._tags,
         )

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/templates/function/.devcontainer/devcontainer.json

@@ -5,6 +5,6 @@
 		"dockerfile": "../Dockerfile"
 	},
 	"features": {
-		"ghcr.io/devcontainers/features/git:1": {},
+		"ghcr.io/devcontainers/features/git:1": {}
 	}
 }

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/templates/script/.devcontainer/devcontainer.json

@@ -5,6 +5,6 @@
 		"dockerfile": "../Dockerfile"
 	},
 	"features": {
-		"ghcr.io/devcontainers/features/git:1": {},
+		"ghcr.io/devcontainers/features/git:1": {}
 	}
 }

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/templates/script/payload/entrypoint.py

@@ -16,7 +16,9 @@ def main():
     You can use your AI models configured in Salesforce to generate column
     values. See README.md for how to test locally before deploying to Data Cloud.
 
-    Example:
+    Example (the per-row helper returns a struct
+    ``{status, response, error_code, error_message}`` — pick the field you
+    want with ``[...]``):
 
         >>> from datacustomcode.client import llm_gateway_generate_text_col
             df_generated = df.withColumn(
@@ -25,8 +27,7 @@ def main():
             ...         "In one sentence, greet {name} from {city}.",
             ...         {"name": col("name__c"), "city": col("homecity__c")},
             ...         model_id="sfdc_ai__DefaultGPT4Omni",
-            ...         max_tokens=100,
-            ...     ),
+            ...     )["response"],
             ... )
 
     You can also invoke the LLM with a literal plain text prompt — no
@@ -35,7 +36,7 @@ def main():
     Example:
 
         >>> generated_text = client.llm_gateway_generate_text(
-        ...     prompt, model_id, max_tokens
+        ...     prompt, model_id
         ... )
     """
 

{salesforce_data_customcode-6.0.2.dev1 → salesforce_data_customcode-6.0.3}/src/datacustomcode/token_provider.py

@@ -101,49 +101,83 @@ 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 org info from sf org display
+        display_data = _run_sf_command(
+            ["sf", "org", "display", "--target-org", self.sf_cli_org, "--json"],
+            "sf org display",
+        )
+        result_data = display_data.get("result", {})
+        instance_url = result_data.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")
+        # Try show-access-token first (SF CLI >= 2.136.6); fall back to the
+        # token from sf org display (older CLIs don't redact it).
+        access_token = None
+        try:
+            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")
+        except RuntimeError:
+            # Command not available on older SF CLI versions
+            access_token = result_data.get("accessToken")
 
-        if not access_token or not instance_url:
+        if not access_token:
             raise RuntimeError(
-                f"'sf org display' did not return an access token or instance URL "
-                f"for org '{self.sf_cli_org}'"
+                f"Could not obtain an access token for org '{self.sf_cli_org}'. "
+                f"Upgrade SF CLI to 2.136.6+ or ensure the org is authenticated."
             )
 
         return AccessTokenResponse(access_token=access_token, instance_url=instance_url)