SDK refactoring

.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

pyproject.toml

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

src/datacustomcode/client.py

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

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

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

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

src/datacustomcode/io/reader/query_api.py

@@ -197,7 +197,7 @@ def read_dlo(
             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 @@ def read_dmo(
             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)
 

src/datacustomcode/llm_gateway/__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.