Removing reference to proxy for cleanup

Author: ritaagarwala-sf

src/datacustomcode/__init__.py

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

src/datacustomcode/client.py

@@ -33,7 +33,6 @@
 
     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
 
 
@@ -107,15 +106,13 @@ 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:
@@ -178,11 +175,6 @@ 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:
@@ -237,11 +229,6 @@ 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,8 +38,6 @@
 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"
@@ -94,23 +92,6 @@ 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")
@@ -128,7 +109,6 @@ 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
@@ -156,7 +136,6 @@ 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/llm_gateway/base.py

@@ -16,9 +16,6 @@
 
 from abc import abstractmethod
 
-from datacustomcode.proxy.base import BaseProxyAccessLayer
-
-
 class LLMGateway:
     def __init__(self):
         pass

src/datacustomcode/llm_gateway/types/generate_text_request.py

@@ -1,11 +1,6 @@
-# Generated by the protocol buffer compiler.  DO NOT EDIT!
-# sources: llm_gateway.proto
-# plugin: python-betterproto
 from dataclasses import dataclass
-from typing import Optional
 
 import betterproto
-import grpclib
 
 from .google import protobuf
 

src/datacustomcode/llm_gateway/types/generate_text_response.py

@@ -1,9 +1,7 @@
 
 from dataclasses import dataclass
-from typing import Optional
 
 import betterproto
-import grpclib
 
 from .google import protobuf
 

src/datacustomcode/llm_gateway/types/generate_text_response_builder.py

@@ -1,10 +1,3 @@
-
-from dataclasses import dataclass
-from typing import Optional
-
-import betterproto
-import grpclib
-
 from .google import protobuf
 
 

src/datacustomcode/templates/function/payload/entrypoint.py

@@ -50,6 +50,7 @@ def function(request: dict, runTime: Runtime) -> dict:
     builder = GenerateTextRequestBuilder()
     request = builder.set_prompt("Hello").set_model("gpt-4").build()
     response = runTime.llm_gateway.generate_text(request)
+
     if response.is_success:
         print(response.text)
     else:

tests/test_client.py

@@ -17,7 +17,6 @@
 )
 from datacustomcode.io.reader.base import BaseDataCloudReader
 from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
-from datacustomcode.proxy.client.base import BaseProxyClient
 
 
 class MockDataCloudReader(BaseDataCloudReader):
@@ -76,13 +75,6 @@ def mock_config(mock_spark):
     )
 
 
-@pytest.fixture
-def mock_proxy():
-    """Mock proxy client to avoid starting Spark when reader/writer are provided."""
-    proxy = MagicMock(spec=BaseProxyClient)
-    return proxy
-
-
 @pytest.fixture
 def reset_client():
     """Reset the Client singleton between tests."""
@@ -93,12 +85,12 @@ def reset_client():
 
 class TestClient:
 
-    def test_singleton_pattern(self, reset_client, mock_spark, mock_proxy):
+    def test_singleton_pattern(self, reset_client, mock_spark):
         """Test that Client behaves as a singleton."""
         reader = MockDataCloudReader(mock_spark)
         writer = MockDataCloudWriter(mock_spark)
 
-        client1 = Client(reader=reader, writer=writer, proxy=mock_proxy)
+        client1 = Client(reader=reader, writer=writer)
         client2 = Client()
 
         assert client1 is client2
@@ -144,38 +136,38 @@ def test_initialization_with_config(self, mock_config, reset_client, mock_spark)
             assert client._reader is mock_reader
             assert client._writer is mock_writer
 
-    def test_read_dlo(self, reset_client, mock_spark, mock_proxy):
+    def test_read_dlo(self, reset_client, mock_spark):
         reader = MagicMock(spec=BaseDataCloudReader)
         writer = MagicMock(spec=BaseDataCloudWriter)
         mock_df = MagicMock(spec=DataFrame)
         reader.read_dlo.return_value = mock_df
 
-        client = Client(reader=reader, writer=writer, proxy=mock_proxy)
+        client = Client(reader=reader, writer=writer)
         result = client.read_dlo("test_dlo")
 
         reader.read_dlo.assert_called_once_with("test_dlo", row_limit=1000)
         assert result is mock_df
         assert "test_dlo" in client._data_layer_history[DataCloudObjectType.DLO]
 
-    def test_read_dmo(self, reset_client, mock_spark, mock_proxy):
+    def test_read_dmo(self, reset_client, mock_spark):
         reader = MagicMock(spec=BaseDataCloudReader)
         writer = MagicMock(spec=BaseDataCloudWriter)
         mock_df = MagicMock(spec=DataFrame)
         reader.read_dmo.return_value = mock_df
 
-        client = Client(reader=reader, writer=writer, proxy=mock_proxy)
+        client = Client(reader=reader, writer=writer)
         result = client.read_dmo("test_dmo")
 
         reader.read_dmo.assert_called_once_with("test_dmo", row_limit=1000)
         assert result is mock_df
         assert "test_dmo" in client._data_layer_history[DataCloudObjectType.DMO]
 
-    def test_write_to_dlo(self, reset_client, mock_spark, mock_proxy):
+    def test_write_to_dlo(self, reset_client, mock_spark):
         reader = MagicMock(spec=BaseDataCloudReader)
         writer = MagicMock(spec=BaseDataCloudWriter)
         mock_df = MagicMock(spec=DataFrame)
 
-        client = Client(reader=reader, writer=writer, proxy=mock_proxy)
+        client = Client(reader=reader, writer=writer)
         client._record_dlo_access("some_dlo")
 
         client.write_to_dlo("test_dlo", mock_df, WriteMode.APPEND, extra_param=True)
@@ -184,12 +176,12 @@ def test_write_to_dlo(self, reset_client, mock_spark, mock_proxy):
             "test_dlo", mock_df, WriteMode.APPEND, extra_param=True
         )
 
-    def test_write_to_dmo(self, reset_client, mock_spark, mock_proxy):
+    def test_write_to_dmo(self, reset_client, mock_spark):
         reader = MagicMock(spec=BaseDataCloudReader)
         writer = MagicMock(spec=BaseDataCloudWriter)
         mock_df = MagicMock(spec=DataFrame)
 
-        client = Client(reader=reader, writer=writer, proxy=mock_proxy)
+        client = Client(reader=reader, writer=writer)
         client._record_dmo_access("some_dmo")
 
         client.write_to_dmo("test_dmo", mock_df, WriteMode.OVERWRITE, extra_param=True)
@@ -198,42 +190,42 @@ def test_write_to_dmo(self, reset_client, mock_spark, mock_proxy):
             "test_dmo", mock_df, WriteMode.OVERWRITE, extra_param=True
         )
 
-    def test_mixed_dlo_dmo_raises_exception(self, reset_client, mock_spark, mock_proxy):
+    def test_mixed_dlo_dmo_raises_exception(self, reset_client, mock_spark):
         """Test that mixing DLOs and DMOs raises an exception."""
         reader = MagicMock(spec=BaseDataCloudReader)
         writer = MagicMock(spec=BaseDataCloudWriter)
         mock_df = MagicMock(spec=DataFrame)
 
-        client = Client(reader=reader, writer=writer, proxy=mock_proxy)
+        client = Client(reader=reader, writer=writer)
         client._record_dlo_access("test_dlo")
 
         with pytest.raises(DataCloudAccessLayerException) as exc_info:
             client.write_to_dmo("test_dmo", mock_df, WriteMode.APPEND)
 
         assert "test_dlo" in str(exc_info.value)
 
-    def test_mixed_dmo_dlo_raises_exception(self, reset_client, mock_spark, mock_proxy):
+    def test_mixed_dmo_dlo_raises_exception(self, reset_client, mock_spark):
         """Test that mixing DMOs and DLOs raises an exception (converse case)."""
         reader = MagicMock(spec=BaseDataCloudReader)
         writer = MagicMock(spec=BaseDataCloudWriter)
         mock_df = MagicMock(spec=DataFrame)
 
-        client = Client(reader=reader, writer=writer, proxy=mock_proxy)
+        client = Client(reader=reader, writer=writer)
         client._record_dmo_access("test_dmo")
 
         with pytest.raises(DataCloudAccessLayerException) as exc_info:
             client.write_to_dlo("test_dlo", mock_df, WriteMode.APPEND)
 
         assert "test_dmo" in str(exc_info.value)
 
-    def test_read_pattern_flow(self, reset_client, mock_spark, mock_proxy):
+    def test_read_pattern_flow(self, reset_client, mock_spark):
         """Test a complete flow of reading and writing within the same object type."""
         reader = MagicMock(spec=BaseDataCloudReader)
         writer = MagicMock(spec=BaseDataCloudWriter)
         mock_df = MagicMock(spec=DataFrame)
         reader.read_dlo.return_value = mock_df
 
-        client = Client(reader=reader, writer=writer, proxy=mock_proxy)
+        client = Client(reader=reader, writer=writer)
 
         df = client.read_dlo("source_dlo")
         client.write_to_dlo("target_dlo", df, WriteMode.APPEND)
@@ -247,7 +239,7 @@ def test_read_pattern_flow(self, reset_client, mock_spark, mock_proxy):
 
         # Reset for DMO test
         Client._instance = None
-        client = Client(reader=reader, writer=writer, proxy=mock_proxy)
+        client = Client(reader=reader, writer=writer)
         reader.read_dmo.return_value = mock_df
 
         df = client.read_dmo("source_dmo")
@@ -260,28 +252,28 @@ def test_read_pattern_flow(self, reset_client, mock_spark, mock_proxy):
 
         assert "source_dmo" in client._data_layer_history[DataCloudObjectType.DMO]
 
-    def test_read_dlo_with_row_limit(self, reset_client, mock_spark, mock_proxy):
+    def test_read_dlo_with_row_limit(self, reset_client, mock_spark):
         """Test that row_limit parameter is passed through to reader."""
         reader = MagicMock(spec=BaseDataCloudReader)
         writer = MagicMock(spec=BaseDataCloudWriter)
         mock_df = MagicMock(spec=DataFrame)
         reader.read_dlo.return_value = mock_df
 
-        client = Client(reader=reader, writer=writer, proxy=mock_proxy)
+        client = Client(reader=reader, writer=writer)
         result = client.read_dlo("test_dlo", row_limit=500)
 
         reader.read_dlo.assert_called_once_with("test_dlo", row_limit=500)
         assert result is mock_df
         assert "test_dlo" in client._data_layer_history[DataCloudObjectType.DLO]
 
-    def test_read_dmo_with_row_limit(self, reset_client, mock_spark, mock_proxy):
+    def test_read_dmo_with_row_limit(self, reset_client, mock_spark):
         """Test that row_limit parameter is passed through to reader."""
         reader = MagicMock(spec=BaseDataCloudReader)
         writer = MagicMock(spec=BaseDataCloudWriter)
         mock_df = MagicMock(spec=DataFrame)
         reader.read_dmo.return_value = mock_df
 
-        client = Client(reader=reader, writer=writer, proxy=mock_proxy)
+        client = Client(reader=reader, writer=writer)
         result = client.read_dmo("test_dmo", row_limit=100)
 
         reader.read_dmo.assert_called_once_with("test_dmo", row_limit=100)