Remote proxy phase 1

Author: atulikumwenayo

src/datacustomcode/__init__.py

@@ -17,15 +17,11 @@
 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,
-)
 
 __all__ = [
     "AuthType",
     "Client",
     "Credentials",
-    "LocalProxyClientProvider",
     "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:
@@ -223,11 +220,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)  # 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)  # type: ignore[no-any-return]
-
     def find_file_path(self, file_name: str) -> Path:
         """Return a file path"""
 

src/datacustomcode/config.py

@@ -39,8 +39,6 @@
 from datacustomcode.io.base import BaseDataAccessLayer
 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: TCH002
 from datacustomcode.spark.base import BaseSparkSessionProvider
 
 if TYPE_CHECKING:
@@ -74,18 +72,6 @@ class SparkConfig(ForceableConfig):
 
 _P = TypeVar("_P", bound=BaseSparkSessionProvider)
 
-_PX = TypeVar("_PX", bound=BaseProxyAccessLayer)
-
-
-class ProxyAccessLayerObjectConfig(BaseObjectConfig, Generic[_PX]):
-    """Config for proxy clients that take no constructor args (e.g. no spark)."""
-
-    type_base: ClassVar[Type[BaseProxyAccessLayer]] = BaseProxyAccessLayer
-
-    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(BaseObjectConfig, Generic[_P]):
     type_base: ClassVar[Type[BaseSparkSessionProvider]] = BaseSparkSessionProvider
@@ -98,7 +84,6 @@ def to_object(self) -> _P:
 class ClientConfig(BaseConfig):
     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
@@ -126,7 +111,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/config.yaml

@@ -19,11 +19,6 @@ spark_config:
     spark.sql.execution.arrow.pyspark.enabled: 'true'
     spark.driver.extraJavaOptions: -Djava.security.manager=allow
 
-proxy_config:
-  type_config_name: LocalProxyClientProvider
-  options:
-    credentials_profile: default
-
 einstein_predictions_config:
   type_config_name: DefaultEinsteinPredictions
   options:

src/datacustomcode/proxy/client/LocalProxyClientProvider.py

@@ -65,6 +65,16 @@ def make_einstein_prediction(runtime: Runtime) -> None:
     )
 
 
+def generate_text(runtime: Runtime):
+    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)
+
 def function(request: dict, runtime: Runtime) -> dict:
     logger.info("Inside Function")
     logger.info(request)
@@ -73,17 +83,9 @@ def function(request: dict, runtime: Runtime) -> dict:
     output_chunks = []
     current_seq_no = 1  # Start sequence number from 1
 
+    generate_text(runtime)
     make_einstein_prediction(runtime)
 
-    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}")

src/datacustomcode/templates/function/payload/entrypoint.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")
         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")
         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")