provider

Author: joroscoSF

src/datacustomcode/client.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 from enum import Enum
+import pprint
 from typing import (
     TYPE_CHECKING,
     ClassVar,
@@ -107,7 +108,7 @@ class Client:
     _reader: BaseDataCloudReader
     _writer: BaseDataCloudWriter
     _file: DefaultFindFilePath
-    _proxy: BaseProxyClient
+    _proxy: Optional[BaseProxyClient]
     _data_layer_history: dict[DataCloudObjectType, set[str]]
 
     def __new__(
@@ -117,18 +118,30 @@ def __new__(
         proxy: Optional[BaseProxyClient] = None,
         spark_provider: Optional["BaseSparkSessionProvider"] = None,
     ) -> Client:
+        print("Chuy client start 2 config:")
+        pprint.pprint(str(config), indent=4)
+
         if cls._instance is None:
             cls._instance = super().__new__(cls)
 
-            spark = None
+            print("Chuy client here")
+
             # Initialize Readers and Writers from config
             # and/or provided reader and writer
             if reader is None or writer is None:
                 # We need a spark because we will initialize readers and writers
                 if config.spark_config is None:
-                    raise ValueError(
-                        "Spark config is required when reader/writer is not provided"
+                    # Assume BYOC Function
+                    # cls._instance._reader = None
+                    # cls._instance._writer = None
+                    cls._instance._file = DefaultFindFilePath()
+                    # cls._instance._data_layer_history = None
+                    cls._instance._proxy = (
+                        config.proxy_config.to_object()  # type: ignore
+                        if config.proxy_config is not None
+                        else None
                     )
+                    return cls._instance
 
                 provider: BaseSparkSessionProvider
                 if spark_provider is not None:
@@ -139,22 +152,6 @@ def __new__(
                     provider = DefaultSparkSessionProvider()
 
                 spark = provider.get_session(config.spark_config)
-            elif (
-                proxy is None
-                and config.proxy_config is not None
-                and config.spark_config is not None
-            ):
-                # Both reader and writer provided; we still need spark for proxy init
-                provider = (
-                    spark_provider
-                    if spark_provider is not None
-                    else (
-                        config.spark_provider_config.to_object()
-                        if config.spark_provider_config is not None
-                        else DefaultSparkSessionProvider()
-                    )
-                )
-                spark = provider.get_session(config.spark_config)
 
             if config.reader_config is None and reader is None:
                 raise ValueError(
@@ -163,44 +160,29 @@ def __new__(
             elif reader is None or (
                 config.reader_config is not None and config.reader_config.force
             ):
-                if config.proxy_config is None:
-                    raise ValueError(
-                        "Proxy config is required when reader is built from config"
-                    )
-                assert (
-                    spark is not None
-                )  # set in "reader is None or writer is None" branch
-                assert config.reader_config is not None  # ensured by branch condition
-                proxy_init = config.proxy_config.to_object(spark)
-
-                reader_init = config.reader_config.to_object(spark)
+                reader_init = config.reader_config.to_object(spark)  # type: ignore
             else:
                 reader_init = reader
-                if proxy is not None:
-                    proxy_init = proxy
-                elif config.proxy_config is None:
-                    raise ValueError("Proxy config is required when reader is provided")
-                else:
-                    assert (
-                        spark is not None
-                    )  # set in "both provided; proxy from config" branch
-                    proxy_init = config.proxy_config.to_object(spark)
             if config.writer_config is None and writer is None:
                 raise ValueError(
                     "Writer config is required when writer is not provided"
                 )
             elif writer is None or (
                 config.writer_config is not None and config.writer_config.force
             ):
-                assert spark is not None  # set when reader or writer from config
-                assert config.writer_config is not None  # ensured by branch condition
-                writer_init = config.writer_config.to_object(spark)
+                writer_init = config.writer_config.to_object(spark)  # type: ignore
             else:
                 writer_init = writer
+            proxy_init: Optional["BaseProxyClient"] = None
+            if proxy is not None:
+                proxy_init = proxy
+            elif config.proxy_config is not None:
+                proxy_init = config.proxy_config.to_object()  # type: ignore
+
             cls._instance._reader = reader_init
             cls._instance._writer = writer_init
-            cls._instance._file = DefaultFindFilePath()
             cls._instance._proxy = proxy_init
+            cls._instance._file = DefaultFindFilePath()
             cls._instance._data_layer_history = {
                 DataCloudObjectType.DLO: set(),
                 DataCloudObjectType.DMO: set(),
@@ -260,6 +242,8 @@ def write_to_dmo(
         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:

src/datacustomcode/config.py

@@ -38,6 +38,7 @@
 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
 
@@ -93,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")
@@ -110,7 +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[AccessLayerObjectConfig[BaseProxyClient], None] = None
+    proxy_config: Union[ProxyAccessLayerObjectConfig[BaseProxyClient], None] = None
     spark_config: Union[SparkConfig, None] = None
     spark_provider_config: Union[
         SparkProviderConfig[BaseSparkSessionProvider], None

src/datacustomcode/config.yaml

@@ -19,6 +19,6 @@ spark_config:
     spark.driver.extraJavaOptions: -Djava.security.manager=allow
 
 proxy_config:
-  type_config_name: LocalProxyClientProvider
+  type_config_name: LocalProxyClient
   options:
     credentials_profile: default

src/datacustomcode/proxy/base.py

@@ -19,6 +19,6 @@
 from datacustomcode.mixin import UserExtendableNamedConfigMixin
 
 
-class BaseDataAccessLayer(ABC, UserExtendableNamedConfigMixin):
+class BaseProxyAccessLayer(ABC, UserExtendableNamedConfigMixin):
     def __init__(self):
         pass

src/datacustomcode/proxy/client/base.py

@@ -16,13 +16,12 @@
 
 from abc import abstractmethod
 
-from datacustomcode.io.base import BaseDataAccessLayer
+from datacustomcode.proxy.base import BaseProxyAccessLayer
 
 
-class BaseProxyClient(BaseDataAccessLayer):
-    def __init__(self, spark=None, **kwargs):
-        if spark is not None:
-            super().__init__(spark)
+class BaseProxyClient(BaseProxyAccessLayer):
+    def __init__(self):
+        pass
 
     @abstractmethod
     def call_llm_gateway(self, llmModelId: str, prompt: str, maxTokens: int) -> str: ...

src/datacustomcode/proxy/client/local_proxy_client.py

@@ -20,7 +20,11 @@
 class LocalProxyClientProvider(BaseProxyClient):
     """Default proxy client provider."""
 
-    CONFIG_NAME = "LocalProxyClientProvider"
+    CONFIG_NAME = "LocalProxyClient"
+
+    def __init__(self, credentials_profile: str = "default", **kwargs: object) -> None:
+        super().__init__()
+        self.credentials_profile = credentials_profile
 
     def call_llm_gateway(self, llmModelId: str, prompt: str, maxTokens: int) -> str:
         return f"Hello, thanks for using {llmModelId}. So many tokens: {maxTokens}"