fix ci build failures, ruff, black, isort, etc

Author: atulikumwenayo

src/datacustomcode/client.py

@@ -66,10 +66,10 @@ def llm_gateway_generate_text_col(
     Example:
 
         >>> df.withColumn(
-        ...     "Greeting__c",
+        ...     "greeting__c",
         ...     llm_gateway_generate_text_col(
         ...         "In one sentence, greet {name} from {city}.",
-        ...         {"name": col("Name__c"), "city": col("HomeCity__c")},
+        ...         {"name": col("name__c"), "city": col("homecity__c")},
         ...         model_id="sfdc_ai__DefaultGPT4Omni",
         ...         max_tokens=100,
         ...     ),

src/datacustomcode/einstein_platform_config.py

@@ -15,20 +15,23 @@
 
 from typing import (
     ClassVar,
+    Generic,
     Optional,
     Type,
-    cast,
+    TypeVar,
 )
 
 from datacustomcode.common_config import BaseObjectConfig
 
+_T = TypeVar("_T")
 
-class CredentialsObjectConfig(BaseObjectConfig):
+
+class CredentialsObjectConfig(BaseObjectConfig, Generic[_T]):
     type_to_create: ClassVar[Type]
     credentials_profile: Optional[str] = None
     sf_cli_org: Optional[str] = None
 
-    def to_object(self):
+    def to_object(self) -> _T:
         """Create an object instance, automatically including credentials in options"""
 
         options = self.options.copy()
@@ -38,4 +41,5 @@ def to_object(self):
             options["sf_cli_org"] = self.sf_cli_org
 
         type_ = self.type_to_create.subclass_from_config_name(self.type_config_name)
-        return cast("type_", type_(**options))
+        instance: _T = type_(**options)
+        return instance

src/datacustomcode/einstein_predictions_config.py

@@ -28,7 +28,7 @@
 _E = TypeVar("_E", bound=EinsteinPredictions)
 
 
-class EinsteinPredictionsObjectConfig(CredentialsObjectConfig, Generic[_E]):
+class EinsteinPredictionsObjectConfig(CredentialsObjectConfig[_E], Generic[_E]):
     type_to_create: ClassVar[Type[EinsteinPredictions]] = EinsteinPredictions  # type: ignore[type-abstract]
 
 

src/datacustomcode/llm_gateway_config.py

@@ -19,7 +19,6 @@
     Type,
     TypeVar,
     Union,
-    cast,
 )
 
 from datacustomcode.common_config import (
@@ -35,7 +34,7 @@
 _S = TypeVar("_S", bound=SparkLLMGateway)
 
 
-class LLMGatewayObjectConfig(CredentialsObjectConfig, Generic[_E]):
+class LLMGatewayObjectConfig(CredentialsObjectConfig[_E], Generic[_E]):
     type_to_create: ClassVar[Type[LLMGateway]] = LLMGateway  # type: ignore[type-abstract]
 
 
@@ -64,7 +63,7 @@ class SparkLLMGatewayObjectConfig(BaseObjectConfig, Generic[_S]):
 
     def to_object(self) -> SparkLLMGateway:
         type_ = self.type_to_create.subclass_from_config_name(self.type_config_name)
-        return cast("SparkLLMGateway", type_(**self.options))
+        return type_(**self.options)
 
 
 class SparkLLMGatewayConfig(BaseConfig):

src/datacustomcode/templates/function/example/chunking_with_llm/entrypoint.py

@@ -7,16 +7,14 @@
 - Requires Runtime parameter (for agentic capabilities)
 - Type-safe with direct field access (no wrappers)
 - Automatic validation and conversion
-"""
 
-"""
-    You can use your AI models configured in Salesforce to generate texts.
+You can use your AI models configured in Salesforce to generate texts.
+
+For testing locally before deploying your code to Data Cloud
+(``datacustomcode run``), first configure an external client app before using
+LLM functionality, then configure the SDK with your client app credentials.
 
-    For testing locally before deploying your code to Data Cloud (datacustomcode run),
-    first configure an external client app before using LLM functionality, then configure
-    the SDK with your client app credentials.
-    
-    https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html#create-a-salesforce-app
+https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html#create-a-salesforce-app
 """
 
 import logging

src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py

@@ -12,16 +12,14 @@
 Type: Regression
 Input: Year_Built__c (numeric)
 Output: Predicted_SalePrice
-"""
 
-"""
-    You can use your AI models configured in Salesforce to make predictions.
+You can use your AI models configured in Salesforce to make predictions.
+
+For testing locally before deploying your code to Data Cloud
+(``datacustomcode run``), first configure an external client app before using
+LLM functionality, then configure the SDK with your client app credentials.
 
-    For testing locally before deploying your code to Data Cloud (datacustomcode run),
-    first configure an external client app before using LLM functionality, then configure
-    the SDK with your client app credentials.
-    
-    https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html#create-a-salesforce-app
+https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html#create-a-salesforce-app
 """
 
 import logging

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

@@ -84,12 +84,14 @@ def function(request: dict, runtime: Runtime) -> dict:
     current_seq_no = 1  # Start sequence number from 1
 
     """
-    You can use your AI models configured in Salesforce to generate texts or predict an outcome.
+    You can use your AI models configured in Salesforce to generate texts
+    or predict an outcome.
+
+    For testing locally before deploying your code to Data Cloud
+    (``datacustomcode run``), first configure an external client app before
+    using LLM functionality, then configure the SDK with your client app
+    credentials.
 
-    For testing locally before deploying your code to Data Cloud (datacustomcode run),
-    first configure an external client app before using LLM functionality, then configure
-    the SDK with your client app credentials.
-    
     https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html#create-a-salesforce-app
 
     Example:

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

@@ -13,21 +13,22 @@ def main():
     df_upper1 = df.withColumn("description__c", upper(col("description__c")))
 
     """
-    You can use your AI models configured in Salesforce to generate column values.
-    For testing locally before deploying your code to Data Cloud (datacustomcode run),
-    first configure an external client app before using LLM functionality, then configure
-    the SDK with your client app credentials.
-    
+    You can use your AI models configured in Salesforce to generate column
+    values. For testing locally before deploying your code to Data Cloud
+    (``datacustomcode run``), first configure an external client app before
+    using LLM functionality, then configure the SDK with your client app
+    credentials.
+
     https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.html#create-a-salesforce-app
 
     Example:
 
         >>> from datacustomcode.client import llm_gateway_generate_text_col
             df_generated = df.withColumn(
-            ...     "Greeting__c",
+            ...     "greeting__c",
             ...     llm_gateway_generate_text_col(
             ...         "In one sentence, greet {name} from {city}.",
-            ...         {"name": col("Name__c"), "city": col("HomeCity__c")},
+            ...         {"name": col("name__c"), "city": col("homecity__c")},
             ...         model_id="sfdc_ai__DefaultGPT4Omni",
             ...         max_tokens=100,
             ...     ),
@@ -38,7 +39,9 @@ def main():
 
     Example:
 
-        >>> generated_text = client.llm_gateway_generate_text(prompt, model_id, max_tokens)
+        >>> generated_text = client.llm_gateway_generate_text(
+        ...     prompt, model_id, max_tokens
+        ... )
     """
 
     # Drop specific columns related to relationships

tests/test_function_utils.py

@@ -29,7 +29,8 @@
 def sample_entrypoint():
     """Create a temporary entrypoint file with a function."""
     with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as temp_file:
-        entrypoint_content = textwrap.dedent("""
+        entrypoint_content = textwrap.dedent(
+            """
             from typing import List
             from pydantic import BaseModel
 
@@ -45,7 +46,8 @@ class SampleResponse(BaseModel):
 
             def function(request: SampleRequest) -> SampleResponse:
                 return SampleResponse(result=f"Processed {request.message}")
-            """)
+            """
+        )
         temp_file.write(entrypoint_content)
         temp_file_path = temp_file.name
 
@@ -59,10 +61,12 @@ def function(request: SampleRequest) -> SampleResponse:
 def entrypoint_no_annotations():
     """Create an entrypoint with no type annotations."""
     with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as temp_file:
-        entrypoint_content = textwrap.dedent("""
+        entrypoint_content = textwrap.dedent(
+            """
             def function(request):
                 return {"result": "no annotations"}
-            """)
+            """
+        )
         temp_file.write(entrypoint_content)
         temp_file_path = temp_file.name
 
@@ -142,7 +146,8 @@ def test_generate_test_json():
         output_simple = os.path.join(temp_dir, "test_simple.json")
 
         with open(models_file, "w") as f:
-            models_content = textwrap.dedent("""
+            models_content = textwrap.dedent(
+                """
                 from pydantic import BaseModel
                 from typing import List
 
@@ -162,16 +167,19 @@ class ComplexRequest(BaseModel):
                     max_items: int = 100
                     config: NestedConfig
                     metadata: dict = {}
-                """)
+                """
+            )
             f.write(models_content)
 
         with open(entrypoint_simple, "w") as f:
-            entrypoint_content = textwrap.dedent("""
+            entrypoint_content = textwrap.dedent(
+                """
                 from test_models import SimpleRequest
 
                 def function(request: SimpleRequest):
                     return {"result": "ok"}
-                """)
+                """
+            )
             f.write(entrypoint_content)
 
         sys.path.insert(0, temp_dir)
@@ -195,12 +203,14 @@ def function(request: SimpleRequest):
         output_complex = os.path.join(temp_dir, "test_complex.json")
 
         with open(entrypoint_complex, "w") as f:
-            entrypoint_content = textwrap.dedent("""
+            entrypoint_content = textwrap.dedent(
+                """
                 from test_models import ComplexRequest
 
                 def function(request: ComplexRequest):
                     return {"result": "ok"}
-                """)
+                """
+            )
             f.write(entrypoint_content)
 
         function_utils.generate_test_json(entrypoint_complex, output_complex)

tests/test_run.py

@@ -48,7 +48,8 @@ def test_config_file():
 def test_entrypoint_file():
     """Create a temporary test entrypoint script with config.json."""
     with tempfile.NamedTemporaryFile(suffix=".py", delete=False) as temp:
-        entrypoint_content = textwrap.dedent("""
+        entrypoint_content = textwrap.dedent(
+            """
             # Test entrypoint
             import os
             import sys
@@ -71,7 +72,8 @@ def test_entrypoint_file():
                     if config.writer_config else None
                 )
                 f.write(f"Writer type: {writer_type}\\n")
-        """)
+        """
+        )
         temp.write(entrypoint_content.encode("utf-8"))
         temp_name = temp.name
 
@@ -163,14 +165,16 @@ def test_run_entrypoint_with_dependencies():
 
         # Create a temp entrypoint that uses the dependency
         with tempfile.NamedTemporaryFile(suffix=".py", delete=False) as temp:
-            entrypoint_content = textwrap.dedent("""
+            entrypoint_content = textwrap.dedent(
+                """
                 # Test entrypoint with dependency
                 import test_dependency
 
                 # Store dependency value in a file to check later
                 with open("dependency_output.txt", "w") as f:
                     f.write(f"Dependency value: {test_dependency.TEST_CONSTANT}\\n")
-            """)
+            """
+            )
             temp.write(entrypoint_content.encode("utf-8"))
             entrypoint_file = temp.name
 
@@ -239,7 +243,8 @@ class TestDataspaceScenarios:
     def test_run_entrypoint_with_default_dataspace(self):
         """Test that run_entrypoint sets dataspace='default' correctly."""
         with tempfile.NamedTemporaryFile(suffix=".py", delete=False) as temp:
-            entrypoint_content = textwrap.dedent("""
+            entrypoint_content = textwrap.dedent(
+                """
                 # Test entrypoint
                 from datacustomcode.config import config
                 with open("dataspace_output.txt", "w") as f:
@@ -255,7 +260,8 @@ def test_run_entrypoint_with_default_dataspace(self):
                     )
                     f.write(f"Reader dataspace: {rds}\\n")
                     f.write(f"Writer dataspace: {wds}\\n")
-            """)
+            """
+            )
             temp.write(entrypoint_content.encode("utf-8"))
             entrypoint_file = temp.name
 
@@ -290,7 +296,8 @@ def test_run_entrypoint_with_custom_dataspace(self):
         """Test that run_entrypoint sets custom dataspace correctly."""
         custom_dataspace = "dataspace-1"
         with tempfile.NamedTemporaryFile(suffix=".py", delete=False) as temp:
-            entrypoint_content = textwrap.dedent("""
+            entrypoint_content = textwrap.dedent(
+                """
                 # Test entrypoint
                 from datacustomcode.config import config
                 with open("dataspace_output.txt", "w") as f:
@@ -306,7 +313,8 @@ def test_run_entrypoint_with_custom_dataspace(self):
                     )
                     f.write(f"Reader dataspace: {rds}\\n")
                     f.write(f"Writer dataspace: {wds}\\n")
-            """)
+            """
+            )
             temp.write(entrypoint_content.encode("utf-8"))
             entrypoint_file = temp.name