Merge pull request #85 from forcedotcom/SDK_refactoring

SDK refactoring

Author: ritaagarwala-sf

.pre-commit-config.yaml

@@ -53,7 +53,7 @@ repos:
     exclude: \.py$
 
 - repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
-  rev: v2.13.0
+  rev: v2.14.0
   hooks:
   - id: pretty-format-toml
     args: [--autofix]
@@ -69,7 +69,7 @@ repos:
   hooks:
   - id: mypy
     name: mypy
-    entry: mypy src/datacustomcode
+    entry: mypy --explicit-package-bases src/datacustomcode
     language: system
     pass_filenames: false
     types: [python]

CHANGELOG.md

@@ -1,5 +1,28 @@
 # Changelog
 
+## 3.0.0
+
+### Breaking Changes
+
+- **Added  `runtime: datacustomcode.runtime.function.Runtime` to function contract for codeType `function`.
+
+  Function now mandates runtime as arguments.
+
+  **Why:** `runTime` allows access to resources ( llm_gateway / file ) available during function execution.
+
+  **Migration:** use function(request: dict, runTime: Runtime) instead od function(request: dict)
+
+  ```python
+  # Before
+  def function(request: dict):
+    pass
+
+  # After
+  def function(request: dict, runTime: Runtime):
+    pass
+  ```
+
+
 ## 2.0.0
 
 ### Breaking Changes

poetry.lock

@@ -72,7 +72,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"

pyproject.toml

@@ -119,8 +119,6 @@ def __new__(
         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 @@ def __new__(
             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 @@ def read_dlo(self, name: str) -> PySparkDataFrame:
             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 @@ def read_dmo(self, name: str) -> PySparkDataFrame:
             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 @@ def write_to_dlo(
             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 @@ def write_to_dmo(
             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

src/__init__.py

@@ -36,10 +36,10 @@
 # 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"

src/datacustomcode/client.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"]

src/datacustomcode/config.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"""

src/datacustomcode/function/__init__.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")