Restoring changes done to script client

Author: ritaagarwala-sf

src/datacustomcode/__init__.py

@@ -17,16 +17,15 @@
 from datacustomcode.credentials import AuthType, Credentials
 from datacustomcode.io.reader.query_api import QueryAPIDataCloudReader
 from datacustomcode.io.writer.print import PrintDataCloudWriter
-# TODO: Restore proxy/LLM gateway integration
-# from datacustomcode.proxy.client.LocalProxyClientProvider import (
-#     LocalProxyClientProvider,
-# )
+from datacustomcode.proxy.client.LocalProxyClientProvider import (
+    LocalProxyClientProvider,
+)
 
 __all__ = [
     "AuthType",
     "Client",
     "Credentials",
-    # "LocalProxyClientProvider",  # TODO: Restore
+    "LocalProxyClientProvider",
     "PrintDataCloudWriter",
     "QueryAPIDataCloudReader",
 ]

src/datacustomcode/client.py

@@ -33,6 +33,7 @@
 
     from datacustomcode.io.reader.base import BaseDataCloudReader
     from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
+    from datacustomcode.proxy.client.base import BaseProxyClient
     from datacustomcode.spark.base import BaseSparkSessionProvider
 
 
@@ -106,13 +107,15 @@ class Client:
     _reader: BaseDataCloudReader
     _writer: BaseDataCloudWriter
     _file: DefaultFindFilePath
+    _proxy: Optional[BaseProxyClient]
     _data_layer_history: dict[DataCloudObjectType, set[str]]
     _code_type: str
 
     def __new__(
         cls,
         reader: Optional[BaseDataCloudReader] = None,
         writer: Optional["BaseDataCloudWriter"] = None,
+        proxy: Optional[BaseProxyClient] = None,
         spark_provider: Optional["BaseSparkSessionProvider"] = None,
         code_type: str = "script",
     ) -> Client:
@@ -175,6 +178,11 @@ def __new__(
     @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, row_limit: int = 1000) -> PySparkDataFrame:
@@ -229,6 +237,11 @@ def write_to_dmo(
         self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO)
         return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs)
 
+    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)
+
     def find_file_path(self, file_name: str) -> Path:
         """Return a file path"""
 

src/datacustomcode/config.py

@@ -38,6 +38,8 @@
 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.proxy.base import BaseProxyAccessLayer
+from datacustomcode.proxy.client.base import BaseProxyClient  # noqa: TCH001
 from datacustomcode.spark.base import BaseSparkSessionProvider
 
 DEFAULT_CONFIG_NAME = "config.yaml"
@@ -92,6 +94,23 @@ class SparkConfig(ForceableConfig):
 
 _P = TypeVar("_P", bound=BaseSparkSessionProvider)
 
+_PX = TypeVar("_PX", bound=BaseProxyAccessLayer)
+
+
+class ProxyAccessLayerObjectConfig(ForceableConfig, Generic[_PX]):
+    """Config for proxy clients that take no constructor args (e.g. no spark)."""
+
+    model_config = ConfigDict(validate_default=True, extra="forbid")
+    type_base: ClassVar[Type[BaseProxyAccessLayer]] = BaseProxyAccessLayer
+    type_config_name: str = Field(
+        description="CONFIG_NAME of the proxy client (e.g. 'LocalProxyClient').",
+    )
+    options: dict[str, Any] = Field(default_factory=dict)
+
+    def to_object(self) -> _PX:
+        type_ = self.type_base.subclass_from_config_name(self.type_config_name)
+        return cast(_PX, type_(**self.options))
+
 
 class SparkProviderConfig(ForceableConfig, Generic[_P]):
     model_config = ConfigDict(validate_default=True, extra="forbid")
@@ -109,6 +128,7 @@ def to_object(self) -> _P:
 class ClientConfig(BaseModel):
     reader_config: Union[AccessLayerObjectConfig[BaseDataCloudReader], None] = None
     writer_config: Union[AccessLayerObjectConfig[BaseDataCloudWriter], None] = None
+    proxy_config: Union[ProxyAccessLayerObjectConfig[BaseProxyClient], None] = None
     spark_config: Union[SparkConfig, None] = None
     spark_provider_config: Union[
         SparkProviderConfig[BaseSparkSessionProvider], None
@@ -136,6 +156,7 @@ def merge(
 
         self.reader_config = merge(self.reader_config, other.reader_config)
         self.writer_config = merge(self.writer_config, other.writer_config)
+        self.proxy_config = merge(self.proxy_config, other.proxy_config)
         self.spark_config = merge(self.spark_config, other.spark_config)
         self.spark_provider_config = merge(
             self.spark_provider_config, other.spark_provider_config

src/datacustomcode/file/base.py

@@ -13,13 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from __future__ import annotations
-from abc import abstractmethod
-from pathlib import Path
 
 
 class BaseDataAccessLayer:
-    def __init__(self):
-        pass
-
-    @abstractmethod
-    def find_file_path(self, file_name: str) -> Path: ...
+    """Base class for data access layer implementations."""

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.

src/datacustomcode/proxy/base.py

@@ -0,0 +1,24 @@
+# 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 ABC
+
+from datacustomcode.mixin import UserExtendableNamedConfigMixin
+
+
+class BaseProxyAccessLayer(ABC, UserExtendableNamedConfigMixin):
+    def __init__(self):
+        pass

src/datacustomcode/proxy/client/LocalProxyClientProvider.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 datacustomcode.proxy.client.base import BaseProxyClient
+
+
+class LocalProxyClientProvider(BaseProxyClient):
+    """Default proxy client provider."""
+
+    CONFIG_NAME = "LocalProxyClientProvider"
+
+    def __init__(self, **kwargs: object) -> None:
+        pass
+
+    def call_llm_gateway(self, llmModelId: str, prompt: str, maxTokens: int) -> str:
+        return f"Hello, thanks for using {llmModelId}. So many tokens: {maxTokens}"
+
+    def llm_gateway_generate_text(
+        self, template, values, llmModelId: str, maxTokens: int
+    ):
+        return f"Using Generate Text with {llmModelId} and maxTokens: {maxTokens}"

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.

src/datacustomcode/proxy/client/base.py

@@ -0,0 +1,32 @@
+# 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 datacustomcode.proxy.base import BaseProxyAccessLayer
+
+
+class BaseProxyClient(BaseProxyAccessLayer):
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def call_llm_gateway(self, llmModelId: str, prompt: str, maxTokens: int) -> str: ...
+
+    @abstractmethod
+    def llm_gateway_generate_text(
+        self, template, values, llmModelId: str, maxTokens: int
+    ): ...