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.

• 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.

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

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.

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.

1. Read your source data from a DLO or DMO.

2. 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.

   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”]).

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

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

1. From your script package root, confirm that your entrypoint file is ready for local validation.

2. Log in to your org using the external client app credentials.

   Replace {MY_DOMAIN_URL} with your org domain and {CONSUMER_KEY} with your external client app consumer key.

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

• Write and Validate Custom Scripts
• Considerations When Writing Code Extension
• Invoke Code Extension by Creating a Batch Data Transform
• Get Started with Agent API