Call Generative AI Models in Batch Transform Scripts

Call generative AI models from your code extension batch transform scripts to enrich your data. Send prompts and data from a data lake object (DLO) or data model object (DMO) to a configured AI model, then write the generated content to a target DLO or DMO. Use this capability to summarize, classify, or enhance records during batch transformations.

Rate limits for LLM calls apply as described in Large Language Model Limits. To increase throughput for larger datasets:

Repartition your DataFrame with .repartition() to spread the workload across more executors. Without repartitioning, requests can run serially on a few executors and create a bottleneck.
Increase the compute type to a larger size so that more CPUs are available to share the load.

Edition Table
Available in: Developer, Enterprise, Performance, and Unlimited Editions. See Data 360 edition availability.

Permission Sets Needed
To call generative AI models from custom scripts:	Permission set: <ul><li>Data Cloud Architect</li></ul>

Prerequisites

Set up Salesforce CLI and the Data Code Extension plugin. See Set Up Salesforce CLI for Code Extension.
Initialize a script package and make sure that you can run local validation commands. See Write and Validate Custom Scripts.
Use a Salesforce org with at least one generative AI model configured.
Configure an external client app to authenticate calls to LLM models. Complete only the Create a Salesforce App and Obtain Credentials sections of Get Started with the Agent API.

Use LLM Gateway in Your Script

For a reference implementation, see example/payload/entrypoint.py in the initialized script package.

Run Single Queries

Use llm_gateway_generate_text() to send a single prompt to a configured AI model for one-off text generation tasks that don't require iterating over DataFrame rows.

response = client.llm_gateway_generate_text(
    "Summarize the key benefits of Data 360 in one paragraph.",
    "sfdc_ai__DefaultGPT54"
)
print(response)

Generate Text for Each Record

Use llm_gateway_generate_text_col() to generate AI content for each record in your DataFrame. This built-in user-defined function (UDF) applies directly to DataFrame columns to process many rows in a single transformation. The function returns a struct containing the response and status information for each row.

Read your source data from a DLO or DMO.

from datacustomcode.client import Client, llm_gateway_generate_text_col
from datacustomcode.io.writer.base import WriteMode
from pyspark.sql.functions import col

def main():
    client = Client()
    df = client.read_dlo("Customer_Data__dll")

Use llm_gateway_generate_text_col() to generate text for each record. Build your prompt template with placeholders in curly braces ({field_name}), then map those placeholders to your DataFrame columns.
```
    df_enriched = df.withColumn(
        "personalized_greeting__c",
        llm_gateway_generate_text_col(
            "Greet {first} {last} welcoming them to the great city of {city}. Tell them what this city has to offer.",
            {"first": col("first_name__c"), "last": col("last_name__c"), "city": col("city__c")},
            model_id="sfdc_ai__DefaultGPT54",
        )["response"],
    )
```
The method returns a column with these fields: status, response, error_code, and error_message. Per-row failures don't stop the entire job. To extract the field you want, use bracket notation (for example, ["response"]).

Write the enriched data back to your target DLO or DMO.

    client.write_to_dlo("Customer_Data_Enriched__dll", df_enriched, write_mode=WriteMode.APPEND)

Validate LLM Gateway Calls Locally

Before you deploy your script to Data 360, test it locally against your sandbox.

From your script package root, confirm that your entrypoint file is ready for local validation.
Log in to your org using the external client app credentials.
```
sf org login web \
  --alias myorg \
  --instance-url https://{MY_DOMAIN_URL} \
  --client-id {CONSUMER_KEY} \
  --scopes "sfap_api api"
```
Replace {MY_DOMAIN_URL} with your org domain and {CONSUMER_KEY} with your external client app consumer key.

Run the script locally to test it against data in Data 360.

sf data-code-extension script run --entrypoint ./payload/entrypoint.py --target-org myorg