Call Generative AI Models in Batch Transform Scripts

Call generative AI models from your code extension batch transform scripts to enrich data at scale. 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.

• 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 external client app access for model API authentication. Complete the steps from Get Started with 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 every record in your dataset. This built-in user-defined function (UDF) applies directly to DataFrame columns, so you can process thousands of rows efficiently. 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 deploying 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. Set up an external client app for authentication.

   Create an external client app with access to the Agent API by following Get Started with Agent API.

   Use http://localhost:1717/OauthRedirect as the callback URL.

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

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