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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: salesforce-data-customcode
-Version: 6.0.2
+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)
@@ -330,6 +331,55 @@ Options:
 - `--function-invoke-opt TEXT`: Currently we support only `UnstructuredChunking` for functions.
 
 
+## Testing LLM Gateway
+
+You can use AI models configured in Salesforce to generate responses while transforming your data. Below is a sample code example:
+
+```
+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
+    )["response"],
+  )
+
+  dlo_name = "Output_dll"
+  client.write_to_dlo(dlo_name, df_upper1, write_mode=WriteMode.APPEND)
+
+  greeting = client.llm_gateway_generate_text("In one sentence, generate a greeting message", "sfdc_ai__DefaultGPT52")
+
+if __name__ == "__main__":
+  main()
+```
+
+In order to test this code on your local machine before deploying it to Data Cloud, you must first set up an External Client App that allows access to the Agent API. Follow this guide to create the ECA https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html#create-a-salesforce-app. You must use `http://localhost:1717/OauthRedirect` as the callback URL.
+
+Once the ECA is set up, log in to your org using this ECA
+```
+sf org login web \
+  --alias myorg \
+  --instance-url https://{MY_DOMAIN_URL} \
+  --client-id {CONSUMER_KEY} \
+  --scopes "sfap_api api"
+```
+
+then you can test your code using `myorg` alias
+```
+datacustomcode run ./payload/entrypoint.py --sf-cli-org myorg
+```
+
+
 ## Docker usage
 
 The SDK provides Docker-based development options that allow you to test your code in an environment that closely resembles Data Cloud's execution environment.

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

@@ -305,6 +305,55 @@ Options:
 - `--function-invoke-opt TEXT`: Currently we support only `UnstructuredChunking` for functions.
 
 
+## Testing LLM Gateway
+
+You can use AI models configured in Salesforce to generate responses while transforming your data. Below is a sample code example:
+
+```
+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
+    )["response"],
+  )
+
+  dlo_name = "Output_dll"
+  client.write_to_dlo(dlo_name, df_upper1, write_mode=WriteMode.APPEND)
+
+  greeting = client.llm_gateway_generate_text("In one sentence, generate a greeting message", "sfdc_ai__DefaultGPT52")
+
+if __name__ == "__main__":
+  main()
+```
+
+In order to test this code on your local machine before deploying it to Data Cloud, you must first set up an External Client App that allows access to the Agent API. Follow this guide to create the ECA https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html#create-a-salesforce-app. You must use `http://localhost:1717/OauthRedirect` as the callback URL.
+
+Once the ECA is set up, log in to your org using this ECA
+```
+sf org login web \
+  --alias myorg \
+  --instance-url https://{MY_DOMAIN_URL} \
+  --client-id {CONSUMER_KEY} \
+  --scopes "sfap_api api"
+```
+
+then you can test your code using `myorg` alias
+```
+datacustomcode run ./payload/entrypoint.py --sf-cli-org myorg
+```
+
+
 ## Docker usage
 
 The SDK provides Docker-based development options that allow you to test your code in an environment that closely resembles Data Cloud's execution environment.

{salesforce_data_customcode-6.0.2 → 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"
+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 → salesforce_data_customcode-6.0.3}/src/datacustomcode/__init__.py

@@ -17,8 +17,11 @@ __all__ = [
     "AuthType",
     "Client",
     "Credentials",
+    "DefaultSparkLLMGateway",
     "PrintDataCloudWriter",
     "QueryAPIDataCloudReader",
+    "SparkLLMGateway",
+    "llm_gateway_generate_text_col",
 ]
 
 
@@ -44,4 +47,16 @@ def __getattr__(name: str):
         from datacustomcode.io.reader.query_api import QueryAPIDataCloudReader
 
         return QueryAPIDataCloudReader
+    elif name == "SparkLLMGateway":
+        from datacustomcode.llm_gateway import SparkLLMGateway
+
+        return SparkLLMGateway
+    elif name == "DefaultSparkLLMGateway":
+        from datacustomcode.llm_gateway import DefaultSparkLLMGateway
+
+        return DefaultSparkLLMGateway
+    elif name == "llm_gateway_generate_text_col":
+        from datacustomcode.client import llm_gateway_generate_text_col
+
+        return llm_gateway_generate_text_col
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

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

@@ -18,24 +18,87 @@ from enum import Enum
 from typing import (
     TYPE_CHECKING,
     ClassVar,
+    Dict,
     Optional,
+    Union,
 )
 
 from datacustomcode.config import config
 from datacustomcode.file.path.default import DefaultFindFilePath
 from datacustomcode.io.reader.base import BaseDataCloudReader
+from datacustomcode.llm_gateway_config import spark_llm_gateway_config
 from datacustomcode.spark.default import DefaultSparkSessionProvider
 
 if TYPE_CHECKING:
     from pathlib import Path
 
-    from pyspark.sql import DataFrame as PySparkDataFrame
+    from pyspark.sql import Column, DataFrame as PySparkDataFrame
 
     from datacustomcode.io.reader.base import BaseDataCloudReader
     from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
+    from datacustomcode.llm_gateway.spark_base import SparkLLMGateway
     from datacustomcode.spark.base import BaseSparkSessionProvider
 
 
+def _build_spark_llm_gateway() -> "SparkLLMGateway":
+    """Instantiate the SDK-configured :class:`SparkLLMGateway`.
+
+    Raises:
+        RuntimeError: If no ``spark_llm_gateway_config`` has been loaded.
+    """
+    cfg = spark_llm_gateway_config.spark_llm_gateway_config
+    if cfg is None:
+        raise RuntimeError(
+            "spark_llm_gateway_config is not configured. Add a "
+            "'spark_llm_gateway_config' section to config.yaml."
+        )
+    return cfg.to_object()
+
+
+def llm_gateway_generate_text_col(
+    template: str,
+    values: Union[Dict[str, "Column"], "Column"],
+    model_id: Optional[str] = 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:
+
+        >>> 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:
+        template: The prompt template, with ``{field}`` placeholders matching
+            keys in ``values``. Substitution uses ``str.format``.
+        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``.
+
+    Returns:
+        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)
+
+
 class DataCloudObjectType(Enum):
     DLO = "dlo"
     DMO = "dmo"
@@ -94,18 +157,21 @@ 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`.
 
     Example:
     >>> client = Client()
     >>> file_path = client.find_file_path("data.csv")
     >>> dlo = client.read_dlo("my_dlo")
     >>> client.write_to_dmo("my_dmo", dlo)
+    >>> answer = client.llm_gateway_generate_text("Generate a greeting message")
     """
 
     _instance: ClassVar[Optional[Client]] = None
     _reader: BaseDataCloudReader
     _writer: BaseDataCloudWriter
     _file: DefaultFindFilePath
+    _spark_llm_gateway: Optional[SparkLLMGateway]
     _data_layer_history: dict[DataCloudObjectType, set[str]]
     _code_type: str
 
@@ -114,11 +180,13 @@ class Client:
         reader: Optional[BaseDataCloudReader] = None,
         writer: Optional[BaseDataCloudWriter] = None,
         spark_provider: Optional[BaseSparkSessionProvider] = None,
+        spark_llm_gateway: Optional[SparkLLMGateway] = None,
         code_type: str = "script",
     ) -> Client:
 
         if cls._instance is None:
             cls._instance = super().__new__(cls)
+            cls._instance._spark_llm_gateway = spark_llm_gateway
             # Initialize Readers and Writers from config
             # and/or provided reader and writer
             if reader is None or writer is None:
@@ -225,6 +293,41 @@ class Client:
 
         return self._file.find_file_path(file_name)  # type: ignore[no-any-return]
 
+    def llm_gateway_generate_text(
+        self,
+        prompt: str,
+        model_id: Optional[str] = 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.
+        Use the column helper method instead when you want to fan a prompt out across
+        every row of a DataFrame.
+
+        Example:
+
+            >>> response = Client().llm_gateway_generate_text(
+            ...     "Generate a greeting message"
+            ... )
+
+        Args:
+            prompt: The literal prompt to send. Plain text — no
+                ``{field}`` substitution is performed on this string.
+            model_id: LLM model id to target. Defaults to
+                ``sfdc_ai__DefaultGPT4Omni`` 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
+        )
+
+    def _get_spark_llm_gateway(self) -> SparkLLMGateway:
+        if self._spark_llm_gateway is None:
+            self._spark_llm_gateway = _build_spark_llm_gateway()
+        return self._spark_llm_gateway
+
     def _validate_data_layer_history_does_not_contain(
         self, data_cloud_object_type: DataCloudObjectType
     ) -> None:

{salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3}/src/datacustomcode/config.yaml

@@ -28,3 +28,6 @@ llm_gateway_config:
   type_config_name: DefaultLLMGateway
   options:
     credentials_profile: default
+
+spark_llm_gateway_config:
+  type_config_name: DefaultSparkLLMGateway

{salesforce_data_customcode-6.0.2 → 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 → salesforce_data_customcode-6.0.3}/src/datacustomcode/einstein_platform_config.py

@@ -15,20 +15,23 @@
 
 from typing import (
     ClassVar,
+    Generic,
     Optional,
     Type,
-    cast,
+    TypeVar,
 )
 
 from datacustomcode.common_config import BaseObjectConfig
 
+_T = TypeVar("_T")
 
-class CredentialsObjectConfig(BaseObjectConfig):
+
+class CredentialsObjectConfig(BaseObjectConfig, Generic[_T]):
     type_to_create: ClassVar[Type]
     credentials_profile: Optional[str] = None
     sf_cli_org: Optional[str] = None
 
-    def to_object(self):
+    def to_object(self) -> _T:
         """Create an object instance, automatically including credentials in options"""
 
         options = self.options.copy()
@@ -38,4 +41,5 @@ class CredentialsObjectConfig(BaseObjectConfig):
             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))
+        instance: _T = type_(**options)
+        return instance

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

@@ -28,7 +28,7 @@ from datacustomcode.einstein_predictions.base import EinsteinPredictions
 _E = TypeVar("_E", bound=EinsteinPredictions)
 
 
-class EinsteinPredictionsObjectConfig(CredentialsObjectConfig, Generic[_E]):
+class EinsteinPredictionsObjectConfig(CredentialsObjectConfig[_E], Generic[_E]):
     type_to_create: ClassVar[Type[EinsteinPredictions]] = EinsteinPredictions  # type: ignore[type-abstract]
 
 

{salesforce_data_customcode-6.0.2 → salesforce_data_customcode-6.0.3}/src/datacustomcode/function/feature_types/chunking.py

@@ -16,6 +16,7 @@
 """
 Pydantic models for Search Index Chunking V1
 """
+
 from enum import Enum
 from typing import (
     Dict,

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

@@ -15,8 +15,14 @@
 
 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
 
 __all__ = [
     "DefaultLLMGateway",
+    "DefaultSparkLLMGateway",
     "LLMGateway",
+    "LLMGatewayCallError",
+    "SparkLLMGateway",
 ]

salesforce_data_customcode-6.0.2/src/datacustomcode/proxy/base.py → salesforce_data_customcode-6.0.3/src/datacustomcode/llm_gateway/errors.py

@@ -12,13 +12,25 @@
 # 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 abc import ABC
+from typing import Optional
 
-from datacustomcode.mixin import UserExtendableNamedConfigMixin
 
+class LLMGatewayCallError(RuntimeError):
+    """Raised when an LLM Gateway call returns an error."""
 
-class BaseProxyAccessLayer(ABC, UserExtendableNamedConfigMixin):
-    def __init__(self):
-        pass
+    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