salesforce-data-customcode 6.0.2.dev1__py3-none-any.whl → 6.0.3.dev1__py3-none-any.whl

datacustomcode/client.py

@@ -59,7 +59,6 @@ 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.
 
@@ -71,7 +70,6 @@ def 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,
         ...     ),
         ... )
 
@@ -81,15 +79,12 @@ 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.
     """
     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 +145,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 +285,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 +302,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:

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",
 ]

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:

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

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,5 @@ 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."""

datacustomcode/llm_gateway/spark_default.py

@@ -31,7 +31,6 @@ if TYPE_CHECKING:
 
 
 _DEFAULT_LLM_MODEL_ID = "sfdc_ai__DefaultGPT4Omni"
-_DEFAULT_LLM_MAX_TOKENS = 200
 
 
 class DefaultSparkLLMGateway(SparkLLMGateway):
@@ -52,16 +51,14 @@ 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":
 
         from pyspark.sql.functions import struct, udf
@@ -83,7 +80,7 @@ class DefaultSparkLLMGateway(SparkLLMGateway):
                 else dict(values_row)
             )
             prompt = template.format(**subs)
-            return _invoke_llm_gateway(gateway, prompt, model_id, max_tokens)
+            return _invoke_llm_gateway(gateway, prompt, model_id)
 
         return udf(_generate, StringType())(values_col)
 
@@ -104,8 +101,8 @@ def _invoke_llm_gateway(
     gateway: "LLMGateway",
     prompt: str,
     model_id: Optional[str],
-    max_tokens: Optional[int],
 ) -> str:
+    from datacustomcode.llm_gateway.errors import LLMGatewayCallError
     from datacustomcode.llm_gateway.types.generate_text_request_builder import (
         GenerateTextRequestBuilder,
     )
@@ -114,6 +111,15 @@ def _invoke_llm_gateway(
         GenerateTextRequestBuilder()
         .set_prompt(prompt)
         .set_model(model_id or _DEFAULT_LLM_MODEL_ID)
-        .set_max_tokens(max_tokens or _DEFAULT_LLM_MAX_TOKENS)
     )
-    return gateway.generate_text(builder.build()).text
+    response = gateway.generate_text(builder.build())
+    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

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"
     )

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,
         )

datacustomcode/templates/script/payload/entrypoint.py

@@ -25,7 +25,6 @@ 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,
             ...     ),
             ... )
 
@@ -35,7 +34,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.dist-info → salesforce_data_customcode-6.0.3.dev1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: salesforce-data-customcode
-Version: 6.0.2.dev1
+Version: 6.0.3.dev1
 Summary: Data Cloud Custom Code SDK
 License-Expression: Apache-2.0
 License-File: LICENSE.txt
@@ -347,7 +347,6 @@ def main():
         "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,
     ),
   )
 

{salesforce_data_customcode-6.0.2.dev1.dist-info → salesforce_data_customcode-6.0.3.dev1.dist-info}/RECORD

@@ -1,7 +1,7 @@
 datacustomcode/__init__.py,sha256=qx-X3U9lU_791ZbAxhGqQeyDw6eqTi9-E9AbYnwJfxU,2071
 datacustomcode/auth.py,sha256=fpSjhIBdv9trC8yq2vuljAix_Euu-4Ah7HDCGhYjOxI,8309
 datacustomcode/cli.py,sha256=ObqhOasZSZ8ZNFP6pywV3FpsUOz45_KaZdMNT4BHOKs,13069
-datacustomcode/client.py,sha256=prAufc2_KAKdy0XX-nYhuD22a93GqQCsQuql9ru-8dY,13294
+datacustomcode/client.py,sha256=NzPtJkHszRh0GqKUkobPyDUCDauUIaeI-zrac9VR7l8,12813
 datacustomcode/cmd.py,sha256=ZMs46aydJw2EaU26JgCtZmnqESQFHvvaJz10hnjZTBk,3537
 datacustomcode/common_config.py,sha256=SAUnxj3kqmOeWwPmFoYq4tuxokMgURVm4QwcGi-avL4,1928
 datacustomcode/config.py,sha256=2Pk61ieQsEWSxKDxlS66_rliKE8iX0j-9LweEm0aaGo,4072
@@ -37,14 +37,15 @@ datacustomcode/io/writer/__init__.py,sha256=gamfOD1VnAtEslRBpqh-yKiVjkG_wYWYdSat
 datacustomcode/io/writer/base.py,sha256=RHafJ-Xmh_TnYcN4QfSh0VdIZeIevKmz9sf05adMgxA,1540
 datacustomcode/io/writer/csv.py,sha256=asty5teBpNQ1fMGHZ7wA3suLhq0sk0lVQPN5x7TOSRk,1555
 datacustomcode/io/writer/print.py,sha256=0g2sP_1Wb95UIyEELWJt3dqM18Iv_CWhDW3mDxcIhns,5105
-datacustomcode/llm_gateway/__init__.py,sha256=ZYEMvQlnWzdSdBwYrhBQGmlfBBqsfmhAy4KViyCBFVA,995
+datacustomcode/llm_gateway/__init__.py,sha256=rcoGpSkv36tZuAOcTxKkhNpM0qV03qgmer_ej07Vmfc,1088
 datacustomcode/llm_gateway/base.py,sha256=CTUhZiZ0JFlmWj6kQpTR4_-mUvqkBKl_Fkdvv9VwxY8,1262
-datacustomcode/llm_gateway/default.py,sha256=HVUc5NihNAQMVgRAJ8w2CBDX4mRLl8FRkpCJWv9FAOU,1952
-datacustomcode/llm_gateway/spark_base.py,sha256=h9fBrhBSSthCMBi6DjryxMKp_WSLdoMDg6a3FgO4oKI,1637
-datacustomcode/llm_gateway/spark_default.py,sha256=2-u9Oy936kj35wv8_qBSc-xtqnC7WNkgprhcuroP7U8,3524
+datacustomcode/llm_gateway/default.py,sha256=QiwrDaUvrh00I-fYYhmogTtqC5r0fmm-JrXxNQLvTc0,1854
+datacustomcode/llm_gateway/errors.py,sha256=ABjztpJwAYMBlmFX808CHHtu4nD6zA4ECelzNJQkJg0,1197
+datacustomcode/llm_gateway/spark_base.py,sha256=apM3FYDPNz3JTgEdlf-7jeryzTYmSzjJsCIt95ZEpvg,1553
+datacustomcode/llm_gateway/spark_default.py,sha256=ctBzkP1vDxgrmoIC6PvfmwxmlPVotkcADK3-M5qyFKI,3794
 datacustomcode/llm_gateway/types/__init__.py,sha256=gamfOD1VnAtEslRBpqh-yKiVjkG_wYWYdSatGIfsN-w,621
-datacustomcode/llm_gateway/types/generate_text_request.py,sha256=pHFp9r12xT_GlFm873kQEfrHCdoJI5sJGrCRVqqN3vg,1647
-datacustomcode/llm_gateway/types/generate_text_request_builder.py,sha256=P2lBwyjGQvqf8G9PR4pNeN02AeYri1XCOvlBeRqAJNw,2791
+datacustomcode/llm_gateway/types/generate_text_request.py,sha256=ChxOCzp7St1EYgwl4tWiwfrPyVci4zPyBhnEjWc1Vfs,1448
+datacustomcode/llm_gateway/types/generate_text_request_builder.py,sha256=R4hCgatMPSfWYqaOmeC1IMQj753iDgu7wU992c4tkmk,2565
 datacustomcode/llm_gateway/types/generate_text_response.py,sha256=NSvj6nvzR4cZ0cVs2-PE9Ct3ZY-xL4MQYknJMd_xdDI,1912
 datacustomcode/llm_gateway/types/generate_text_response_builder.py,sha256=6MV6JZuTyHUId3Slw9k7m0KBgjj-5S9HI4qtd94Jl2A,1272
 datacustomcode/llm_gateway_config.py,sha256=ECm9CFqOFOaLHwenvVFAjiRSzG9Eg9fY4pb8nHbQ34c,3312
@@ -85,13 +86,13 @@ datacustomcode/templates/script/examples/employee_hierarchy/employee_data.csv,sh
 datacustomcode/templates/script/examples/employee_hierarchy/entrypoint.py,sha256=Mfm3iQtEHTQRW6cmZTwRKdTh3IeGWR-teXrd6x-YLSY,2223
 datacustomcode/templates/script/jupyterlab.sh,sha256=IHR3YQ8d_busuyvesByhJgzCKULIFPI-Hqogs0NPhCs,2432
 datacustomcode/templates/script/payload/config.json,sha256=0d2mEMt4NHIeOUTsgiPMuRKdOLIQxmh-Wq-IfEqU-gc,28
-datacustomcode/templates/script/payload/entrypoint.py,sha256=HFrTe7--w5X78Sh4DVkoGGR6-zeqxemqp80RM20Rr4k,1642
+datacustomcode/templates/script/payload/entrypoint.py,sha256=Ox2guNka3r0GW2jzKsmH7AP70j1-1Qqr_jhtH5dkPOM,1590
 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/version.py,sha256=9LlbVrzwBvut1L308QbwNBgxdQk5CrghH39JARVSxms,989
-salesforce_data_customcode-6.0.2.dev1.dist-info/METADATA,sha256=8lK5D1Zf3TelM0vlib8QBQLtdmixh79g4ce2Owuh8QQ,21748
-salesforce_data_customcode-6.0.2.dev1.dist-info/WHEEL,sha256=eY7nduwzv-ldUxpzbRlxwvC693Hg6PX8bWDjEHjZ_dk,88
-salesforce_data_customcode-6.0.2.dev1.dist-info/entry_points.txt,sha256=WpQ94UB7UuRCYGOLtJV3vAgm1tl7iVf43VYRWIuNu5Y,74
-salesforce_data_customcode-6.0.2.dev1.dist-info/licenses/LICENSE.txt,sha256=iOi8EmQpfkFhMENi7VYtkl2EqS14pEOccoXHiW2dyPU,11443
-salesforce_data_customcode-6.0.2.dev1.dist-info/RECORD,,
+salesforce_data_customcode-6.0.3.dev1.dist-info/METADATA,sha256=wmHyC6LrpV7uYwKuMBQEwfrztUX9gLb7v7YRwsYYGAg,21724
+salesforce_data_customcode-6.0.3.dev1.dist-info/WHEEL,sha256=eY7nduwzv-ldUxpzbRlxwvC693Hg6PX8bWDjEHjZ_dk,88
+salesforce_data_customcode-6.0.3.dev1.dist-info/entry_points.txt,sha256=WpQ94UB7UuRCYGOLtJV3vAgm1tl7iVf43VYRWIuNu5Y,74
+salesforce_data_customcode-6.0.3.dev1.dist-info/licenses/LICENSE.txt,sha256=iOi8EmQpfkFhMENi7VYtkl2EqS14pEOccoXHiW2dyPU,11443
+salesforce_data_customcode-6.0.3.dev1.dist-info/RECORD,,