salesforce-data-customcode 2.0.0__tar.gz → 3.0.1__tar.gz

{salesforce_data_customcode-2.0.0 → salesforce_data_customcode-3.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: salesforce-data-customcode
-Version: 2.0.0
+Version: 3.0.1
 Summary: Data Cloud Custom Code SDK
 License-Expression: Apache-2.0
 License-File: LICENSE.txt
@@ -15,7 +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: pydantic (>=1.8.2,<3.0.0)
+Requires-Dist: pydantic (==2.13.1)
 Requires-Dist: pyspark (==3.5.1)
 Requires-Dist: pyyaml (>=6.0,<7.0)
 Requires-Dist: salesforce-cdp-connector (>=1.0.19)

{salesforce_data_customcode-2.0.0 → salesforce_data_customcode-3.0.1}/pyproject.toml

@@ -18,7 +18,7 @@ license = "Apache-2.0"
 name = "salesforce-data-customcode"
 readme = "README.md"
 requires-python = ">=3.10,<3.12"
-version = "2.0.0"
+version = "3.0.1"
 
 [tool.black]
 exclude = '''
@@ -73,7 +73,9 @@ use_parentheses = true
 
 [tool.mypy]
 check_untyped_defs = false
+explicit_package_bases = true
 ignore_missing_imports = true
+mypy_path = "src"
 no_implicit_optional = true
 plugins = [
   'pydantic.mypy'
@@ -99,7 +101,7 @@ click = "^8.1.8"
 loguru = "^0.7.3"
 numpy = "*"
 pandas = "*"
-pydantic = "^1.8.2 || ^2.0.0"
+pydantic = "2.13.1"
 pyspark = "3.5.1"
 python = ">=3.10,<3.12"
 pyyaml = "^6.0"

{salesforce_data_customcode-2.0.0 → salesforce_data_customcode-3.0.1}/src/datacustomcode/client.py

@@ -119,8 +119,6 @@ class Client:
         spark_provider: Optional["BaseSparkSessionProvider"] = None,
         code_type: str = "script",
     ) -> Client:
-        if "function" in code_type:
-            return cls._new_function_client()
 
         if cls._instance is None:
             cls._instance = super().__new__(cls)
@@ -175,16 +173,6 @@ class Client:
             raise ValueError("Cannot set reader or writer after client is initialized")
         return cls._instance
 
-    @classmethod
-    def _new_function_client(cls) -> Client:
-        cls._instance = super().__new__(cls)
-        cls._instance._proxy = (
-            config.proxy_config.to_object()  # type: ignore
-            if config.proxy_config is not None
-            else None
-        )
-        return cls._instance
-
     def read_dlo(self, name: str) -> PySparkDataFrame:
         """Read a DLO from Data Cloud.
 
@@ -195,7 +183,7 @@ class Client:
             A PySpark DataFrame containing the DLO data.
         """
         self._record_dlo_access(name)
-        return self._reader.read_dlo(name)
+        return self._reader.read_dlo(name)  # type: ignore[no-any-return]
 
     def read_dmo(self, name: str) -> PySparkDataFrame:
         """Read a DMO from Data Cloud.
@@ -207,7 +195,7 @@ class Client:
             A PySpark DataFrame containing the DMO data.
         """
         self._record_dmo_access(name)
-        return self._reader.read_dmo(name)
+        return self._reader.read_dmo(name)  # type: ignore[no-any-return]
 
     def write_to_dlo(
         self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
@@ -220,7 +208,7 @@ class Client:
             write_mode: The write mode to use for writing to the DLO.
         """
         self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO)
-        return self._writer.write_to_dlo(name, dataframe, write_mode, **kwargs)
+        return self._writer.write_to_dlo(name, dataframe, write_mode, **kwargs)  # type: ignore[no-any-return]
 
     def write_to_dmo(
         self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
@@ -233,17 +221,17 @@ class Client:
             write_mode: The write mode to use for writing to the DMO.
         """
         self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO)
-        return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs)
+        return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs)  # type: ignore[no-any-return]
 
     def call_llm_gateway(self, LLM_MODEL_ID: str, prompt: str, maxTokens: int) -> str:
         if self._proxy is None:
             raise ValueError("No proxy configured; set proxy or proxy_config")
-        return self._proxy.call_llm_gateway(LLM_MODEL_ID, prompt, maxTokens)
+        return self._proxy.call_llm_gateway(LLM_MODEL_ID, prompt, maxTokens)  # type: ignore[no-any-return]
 
     def find_file_path(self, file_name: str) -> Path:
         """Return a file path"""
 
-        return self._file.find_file_path(file_name)
+        return self._file.find_file_path(file_name)  # type: ignore[no-any-return]
 
     def _validate_data_layer_history_does_not_contain(
         self, data_cloud_object_type: DataCloudObjectType

{salesforce_data_customcode-2.0.0 → salesforce_data_customcode-3.0.1}/src/datacustomcode/config.py

@@ -36,10 +36,10 @@ import yaml
 # This lets all readers and writers to be findable via config
 from datacustomcode.io import *  # noqa: F403
 from datacustomcode.io.base import BaseDataAccessLayer
-from datacustomcode.io.reader.base import BaseDataCloudReader  # noqa: TCH001
-from datacustomcode.io.writer.base import BaseDataCloudWriter  # noqa: TCH001
+from datacustomcode.io.reader.base import BaseDataCloudReader  # noqa: TCH002
+from datacustomcode.io.writer.base import BaseDataCloudWriter  # noqa: TCH002
 from datacustomcode.proxy.base import BaseProxyAccessLayer
-from datacustomcode.proxy.client.base import BaseProxyClient  # noqa: TCH001
+from datacustomcode.proxy.client.base import BaseProxyClient  # noqa: TCH002
 from datacustomcode.spark.base import BaseSparkSessionProvider
 
 DEFAULT_CONFIG_NAME = "config.yaml"

salesforce_data_customcode-3.0.1/src/datacustomcode/function/__init__.py

@@ -0,0 +1,20 @@
+# 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.
+
+"""Function runtime for Data Cloud Custom Code."""
+
+from datacustomcode.function.runtime import Runtime
+
+__all__ = ["Runtime"]

salesforce_data_customcode-3.0.1/src/datacustomcode/function/base.py

@@ -0,0 +1,18 @@
+# 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.
+
+
+class BaseRuntime:
+    """Base class for datacustomcode run time"""

salesforce_data_customcode-3.0.1/src/datacustomcode/function/feature_types/chunking.py

@@ -0,0 +1,89 @@
+# 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.
+
+"""
+Pydantic models for byoc-function-proto (uds_chunking.proto)
+Auto-generated - validation rules from buf.validate
+"""
+
+from typing import (
+    Any,
+    Dict,
+    List,
+    Literal,
+)
+
+from pydantic import BaseModel, Field
+
+
+class DocElement(BaseModel):
+    """Document element to be chunked"""
+
+    text: str = Field(..., description="Text content to be chunked")
+    metadata: Dict[str, Any] = Field(
+        default_factory=dict, description="Source document metadata"
+    )
+
+
+class ChunkOutput(BaseModel):
+    """Output chunk from the chunking process"""
+
+    chunk_id: str = Field(..., description="UUID for this chunk")
+    chunk_type: str = Field(..., description="Type: 'text'")
+    text: str = Field(..., description="Chunk text content")
+    seq_no: int = Field(..., description="Sequential chunk number (1-based)")
+    metadata: Dict[str, str] = Field(
+        default_factory=dict, description="Metadata from source (DMO fields)"
+    )
+    tag_metadata: Dict[str, Any] = Field(
+        default_factory=dict, description="Additional tags"
+    )
+    citations: Dict[str, Any] = Field(
+        default_factory=dict, description="Citation information"
+    )
+
+
+class StatusResponse(BaseModel):
+    """Status response for operation"""
+
+    status_type: str = Field(..., description="'success' or 'error'")
+    status_message: str = Field(..., description="Human-readable status")
+
+
+class UdsChunkingV1BatchRequest(BaseModel):
+    """Batch request for UDS chunking"""
+
+    version: Literal["v1"] = Field(
+        default="v1", description="API version, must be 'v1'"
+    )
+    input: List[DocElement] = Field(
+        ..., min_length=1, description="List of documents (min 1)"
+    )
+    max_characters: int = Field(..., description="Max chars per chunk (default: 100)")
+    additional_params: Dict[str, Any] = Field(
+        default_factory=dict, description="Future extension point"
+    )
+
+
+class UdsChunkingV1BatchResponse(BaseModel):
+    """Batch response for UDS chunking"""
+
+    version: Literal["v1"] = Field(
+        default="v1", description="API version, must be 'v1'"
+    )
+    output: List[ChunkOutput] = Field(
+        default_factory=list, description="Flat list of chunks from all docs"
+    )
+    status: StatusResponse = Field(..., description="Overall operation status")

salesforce_data_customcode-3.0.1/src/datacustomcode/function/runtime.py

@@ -0,0 +1,77 @@
+# 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.
+
+
+import threading
+from typing import Optional
+
+from datacustomcode.file.path.default import DefaultFindFilePath
+from datacustomcode.function.base import BaseRuntime
+from datacustomcode.llm_gateway.default import DefaultLLMGateway
+
+
+class Runtime(BaseRuntime):
+    """Client for Function code type.
+
+    NOTE: Do not instantiate this class directly.
+    It will be provided to your function by the SDK:
+
+        def function(request: dict, runtime: RunTime) -> dict:
+            response = {...}
+            return response
+
+    """
+
+    _instance: Optional["Runtime"] = None
+    _lock = threading.Lock()
+
+    def __new__(cls):
+        """Create singleton instance (thread-safe)."""
+        with cls._lock:
+            if cls._instance is not None:
+                raise RuntimeError(
+                    "Runtime can only be instantiated once by the SDK.\n\n"
+                    "Do not instantiate it yourself. Accept it as a parameter:\n\n"
+                    "  from datacustomcode.runtime.function.RunTime import Function\n"
+                    "  \n"
+                    "  def function(request: dict, runtime: Runtime) -> dict:\n"
+                    "      response = {...}\n"
+                    "      return response"
+                )
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self) -> None:
+        # Prevent re-initialization
+        if hasattr(self, "_initialized"):
+            return
+
+        self._initialized = True
+
+        super().__init__()
+
+        # Initialize resources
+        self._llm_gateway = DefaultLLMGateway()
+        self._file = DefaultFindFilePath()
+
+    @property
+    def llm_gateway(self) -> DefaultLLMGateway:
+        """Access LLM operations."""
+        return self._llm_gateway
+
+    @property
+    def file(self) -> DefaultFindFilePath:
+        """Access file operations."""
+        return self._file

{salesforce_data_customcode-2.0.0 → salesforce_data_customcode-3.0.1}/src/datacustomcode/io/reader/query_api.py

@@ -197,7 +197,7 @@ class QueryAPIDataCloudReader(BaseDataCloudReader):
             self, "_sf_cli_reader", None
         )
         if sf_cli_reader is not None:
-            return sf_cli_reader.read_dlo(name, schema)
+            return sf_cli_reader.read_dlo(name, schema)  # type: ignore[no-any-return]
 
         query = self._build_query(name)
 
@@ -230,7 +230,7 @@ class QueryAPIDataCloudReader(BaseDataCloudReader):
             self, "_sf_cli_reader", None
         )
         if sf_cli_reader is not None:
-            return sf_cli_reader.read_dmo(name, schema)
+            return sf_cli_reader.read_dmo(name, schema)  # type: ignore[no-any-return]
 
         query = self._build_query(name)
 

salesforce_data_customcode-3.0.1/src/datacustomcode/llm_gateway/base.py

@@ -0,0 +1,34 @@
+# 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.
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from datacustomcode.llm_gateway.types.generate_text_request import (
+        GenerateTextRequest,
+    )
+    from datacustomcode.llm_gateway.types.generate_text_response import (
+        GenerateTextResponse,
+    )
+
+
+class LLMGateway:
+    def __init__(self) -> None:
+        pass
+
+    @abstractmethod
+    def generate_text(self, request: GenerateTextRequest) -> GenerateTextResponse: ...

salesforce_data_customcode-3.0.1/src/datacustomcode/llm_gateway/default.py

@@ -0,0 +1,33 @@
+# 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.
+
+from datacustomcode.llm_gateway.base import LLMGateway
+from datacustomcode.llm_gateway.types.generate_text_request import GenerateTextRequest
+from datacustomcode.llm_gateway.types.generate_text_response import GenerateTextResponse
+from datacustomcode.llm_gateway.types.generate_text_response_builder import (
+    GenerateTextResponseBuilder,
+)
+
+
+class DefaultLLMGateway(LLMGateway):
+    def generate_text(self, request: GenerateTextRequest) -> GenerateTextResponse:
+
+        response_data = {
+            "version": "v1",
+            "status_code": 200,
+            "data": {"generation": {"generatedText": "Hello World"}},
+        }
+
+        return GenerateTextResponseBuilder.build(response_data)

salesforce_data_customcode-3.0.1/src/datacustomcode/llm_gateway/types/__init__.py

@@ -0,0 +1,14 @@
+# 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.

salesforce_data_customcode-3.0.1/src/datacustomcode/llm_gateway/types/generate_text_request.py

@@ -0,0 +1,46 @@
+# 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.
+
+from typing import (
+    Any,
+    Dict,
+    Literal,
+    Optional,
+)
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+)
+from pydantic.alias_generators import to_camel
+
+
+class GenerateTextRequest(BaseModel):
+
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,  # Allows both snake_case and camelCase input
+    )
+
+    version: Literal["v1"] = Field(
+        default="v1", description="API version, must be 'v1'"
+    )
+    model_name: str = Field(..., min_length=1, description="Name of the model to use")
+    prompt: str = Field(..., description="Input prompt")
+    localization: Optional[Dict[str, Any]] = Field(
+        default=None, description="Localization settings"
+    )
+    tags: Optional[Dict[str, Any]] = Field(default=None, description="Additional tags")

salesforce_data_customcode-3.0.1/src/datacustomcode/llm_gateway/types/generate_text_request_builder.py

@@ -0,0 +1,82 @@
+# 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.
+
+from typing import (
+    Any,
+    Dict,
+    Optional,
+)
+
+from datacustomcode.llm_gateway.types.generate_text_request import GenerateTextRequest
+
+
+class GenerateTextRequestBuilder:
+    def __init__(self) -> None:
+        self._prompt = ""
+        self._model_name = ""
+        self._localization: Optional[Dict[str, Any]] = None
+        self._tags: Optional[Dict[str, Any]] = None
+
+    def set_prompt(self, prompt: str) -> "GenerateTextRequestBuilder":
+        self._prompt = prompt
+        return self
+
+    def set_model(self, model_name: str) -> "GenerateTextRequestBuilder":
+        self._model_name = model_name
+        return self
+
+    def set_localization(
+        self,
+        localization: Optional[Dict[str, Any]] = None,
+        locale: Optional[str] = None,
+    ) -> "GenerateTextRequestBuilder":
+        """
+        Set localization either from a dict or a simple locale string.
+
+        Args:
+          localization: Full localization dict (if provided, locale is ignored)
+          locale: Simple locale string for defaultLocale only
+
+        Returns:
+            self for method chaining
+        """
+
+        if localization is not None:
+            self._localization = localization
+        elif locale is not None:
+            self._localization = {
+                "defaultLocale": locale,
+                "inputLocales": [{"locale": locale, "probability": 1.0}],
+                "expectedLocales": [locale],
+            }
+        else:
+            raise ValueError("Must provide either localization or locale")
+
+        return self
+
+    def set_tags(self, tags: Dict[str, Any]) -> "GenerateTextRequestBuilder":
+        self._tags = tags
+        return self
+
+    def build(self) -> GenerateTextRequest:
+
+        request = GenerateTextRequest(
+            prompt=self._prompt,
+            model_name=self._model_name,
+            localization=self._localization,
+            tags=self._tags,
+        )
+
+        return request

salesforce_data_customcode-3.0.1/src/datacustomcode/llm_gateway/types/generate_text_response.py

@@ -0,0 +1,58 @@
+# 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.
+
+from typing import (
+    Any,
+    Dict,
+    Optional,
+)
+
+from pydantic import BaseModel, Field
+
+
+class GenerateTextResponse(BaseModel):
+    """Response from LLM text generation"""
+
+    version: str = Field(default="v1", description="API version")
+    status_code: int = Field(..., description="HTTP status code", ge=0)
+    data: Optional[Dict[str, Any]] = Field(default=None, description="Response data")
+
+    @property
+    def is_success(self) -> bool:
+        """Check if request succeeded."""
+        return self.status_code == 200
+
+    @property
+    def is_error(self) -> bool:
+        """Check if request failed."""
+        return not self.is_success
+
+    @property
+    def text(self) -> str:
+        """Generated text (convenience property)."""
+        if self.is_success and self.data:
+            generation = self.data.get("generation", {})
+            if isinstance(generation, dict):
+                text = generation.get("generatedText", "")
+                return str(text) if text else ""
+        return ""
+
+    @property
+    def error_code(self) -> str:
+        """Generated text (convenience property)."""
+        if self.is_error and self.data:
+            error_code = self.data.get("errorCode", str(self.status_code))
+            return str(error_code)
+        return ""

salesforce_data_customcode-3.0.1/src/datacustomcode/llm_gateway/types/generate_text_response_builder.py

@@ -0,0 +1,37 @@
+# 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.
+
+from typing import Any, Dict
+
+from datacustomcode.llm_gateway.types.generate_text_response import GenerateTextResponse
+
+
+class GenerateTextResponseBuilder:
+    def __init__(self):
+        self._version = "v1"  # Hardcoded default for your SDK
+        self._status_code = None
+        self._data = None
+
+    def set_status_code(self, status_code: int):
+        self._status_code = status_code
+        return self
+
+    def set_data(self, data: dict):
+        self._data = data
+        return self
+
+    @staticmethod
+    def build(response_dict: Dict[str, Any]) -> GenerateTextResponse:
+        return GenerateTextResponse.model_validate(response_dict)

salesforce_data_customcode-3.0.1/src/datacustomcode/proxy/__init__.py

@@ -0,0 +1,14 @@
+# 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.

salesforce_data_customcode-3.0.1/src/datacustomcode/proxy/client/__init__.py

@@ -0,0 +1,14 @@
+# 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.

{salesforce_data_customcode-2.0.0 → salesforce_data_customcode-3.0.1}/src/datacustomcode/templates/function/payload/entrypoint.py

@@ -2,6 +2,11 @@ import logging
 from typing import List
 from uuid import uuid4
 
+from datacustomcode.function import Runtime
+from datacustomcode.llm_gateway.types.generate_text_request_builder import (
+    GenerateTextRequestBuilder,
+)
+
 logger = logging.getLogger(__name__)
 
 
@@ -33,7 +38,7 @@ def chunk_text(text: str, chunk_size: int = 1000) -> List[str]:
     return chunks
 
 
-def function(request: dict) -> dict:
+def function(request: dict, runtime: Runtime) -> dict:
     logger.info("Inside Function")
     logger.info(request)
 
@@ -41,6 +46,15 @@ def function(request: dict) -> dict:
     output_chunks = []
     current_seq_no = 1  # Start sequence number from 1
 
+    builder = GenerateTextRequestBuilder()
+    llm_request = builder.set_prompt("Hello").set_model("modelName").build()
+    llm_response = runtime.llm_gateway.generate_text(llm_request)
+
+    if llm_response.is_success:
+        print(llm_response.text)
+    else:
+        print(llm_response.error_code)
+
     for item in items:
         # Item is DocElement as dict
         logger.info(f"Processing item: {item}")
@@ -107,7 +121,7 @@ if __name__ == "__main__":
     }
 
     # Run the function
-    result = function(test_request)
+    result = function(test_request, Runtime())
 
     # Print the results in a more readable format
     print("\nChunking Results:")