kind of working

Author: joroscoSF

src/datacustomcode/__init__.py

@@ -17,11 +17,13 @@
 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.local_proxy_client import LocalProxyClientProvider
 
 __all__ = [
     "AuthType",
     "Client",
     "Credentials",
+    "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,12 +107,14 @@ class Client:
     _reader: BaseDataCloudReader
     _writer: BaseDataCloudWriter
     _file: DefaultFindFilePath
+    _proxy: BaseProxyClient
     _data_layer_history: dict[DataCloudObjectType, set[str]]
 
     def __new__(
         cls,
         reader: Optional[BaseDataCloudReader] = None,
         writer: Optional["BaseDataCloudWriter"] = None,
+        proxy: Optional[BaseProxyClient] = None,
         spark_provider: Optional["BaseSparkSessionProvider"] = None,
     ) -> Client:
         if cls._instance is None:
@@ -143,9 +146,21 @@ 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"
+                    )
+                proxy_init = config.proxy_config.to_object(spark)
+                print(f"chuy1 reader, proxy_init: {proxy_init}")
+
                 reader_init = config.reader_config.to_object(spark)  # type: ignore
+                print(f"chuy2 reader, reader_init: {reader_init}, spark: {spark}")
             else:
+                print("chuy2 reader")
                 reader_init = reader
+                if config.proxy_config is None:
+                    raise ValueError("Proxy config is required when reader is provided")
+                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"
@@ -159,6 +174,7 @@ def __new__(
             cls._instance._reader = reader_init
             cls._instance._writer = writer_init
             cls._instance._file = DefaultFindFilePath()
+            cls._instance._proxy = proxy_init
             cls._instance._data_layer_history = {
                 DataCloudObjectType.DLO: set(),
                 DataCloudObjectType.DMO: set(),
@@ -217,6 +233,9 @@ 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):
+        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,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.client.base import BaseProxyClient  # noqa: TCH001
 from datacustomcode.spark.base import BaseSparkSessionProvider
 
 DEFAULT_CONFIG_NAME = "config.yaml"
@@ -109,6 +110,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
     spark_config: Union[SparkConfig, None] = None
     spark_provider_config: Union[
         SparkProviderConfig[BaseSparkSessionProvider], None
@@ -136,6 +138,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/config.yaml

@@ -17,3 +17,8 @@ spark_config:
     spark.submit.deployMode: client
     spark.sql.execution.arrow.pyspark.enabled: 'true'
     spark.driver.extraJavaOptions: -Djava.security.manager=allow
+
+proxy_config:
+  type_config_name: LocalProxyClientProvider
+  options:
+    credentials_profile: default

src/datacustomcode/credentials.py

@@ -192,6 +192,8 @@ def from_available(cls, profile: str = "default") -> Credentials:
         Raises:
             ValueError: If no credentials are found in any source
         """
+        # print(f"Chuy profile: {profile}")
+
         # Check environment variables first
         if os.environ.get("SFDC_LOGIN_URL"):
             logger.debug("Loading credentials from environment variables")

src/datacustomcode/run.py

@@ -21,6 +21,7 @@
 from typing import List, Union
 
 from datacustomcode.config import config
+from datacustomcode.scan import get_package_type
 
 
 def _set_config_option(config_obj, key: str, value: str) -> None:
@@ -60,6 +61,12 @@ def run_entrypoint(
             f"config.json not found at {config_json_path}. config.json is required."
         )
 
+    package_type = get_package_type(entrypoint_dir)
+    print(
+        f"Chuy entrypoint: {entrypoint} directory: {entrypoint_dir} "
+        f"package type: {package_type}"
+    )
+
     try:
         with open(config_json_path, "r") as f:
             config_json = json.load(f)
@@ -68,21 +75,23 @@ def run_entrypoint(
             f"config.json at {config_json_path} is not valid JSON"
         ) from err
 
-    # Require dataspace to be present in config.json
-    dataspace = config_json.get("dataspace")
-    if not dataspace:
-        raise ValueError(
-            f"config.json at {config_json_path} is missing required field 'dataspace'. "
-            f"Please ensure config.json contains a 'dataspace' field."
-        )
-
-    # Load config file first
-    if config_file:
-        config.load(config_file)
-
-    # Add dataspace to reader and writer config options
-    _set_config_option(config.reader_config, "dataspace", dataspace)
-    _set_config_option(config.writer_config, "dataspace", dataspace)
+    if package_type == "script":
+        # Require dataspace to be present in config.json
+        dataspace = config_json.get("dataspace")
+        if not dataspace:
+            raise ValueError(
+                f"config.json at {config_json_path} is missing required "
+                f"field 'dataspace'. "
+                f"Please ensure config.json contains a 'dataspace' field."
+            )
+
+        # Load config file first
+        if config_file:
+            config.load(config_file)
+
+        # Add dataspace to reader and writer config options
+        _set_config_option(config.reader_config, "dataspace", dataspace)
+        _set_config_option(config.writer_config, "dataspace", dataspace)
 
     if profile != "default":
         _set_config_option(config.reader_config, "credentials_profile", profile)