# 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](https://help.salesforce.com/s/articleView?id=ai.generative_ai_llm_limits.htm\&type=5). To increase throughput for larger datasets:

- Repartition your DataFrame with [`.repartition()`](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.repartition.html) 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](https://help.salesforce.com/s/articleView?id=sf.c360_a_dc_editions.htm&type=5). |

| 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](/docs/data/data-cloud-code-ext/guide/set-up-sdk.md).
- Initialize a script package and make sure that you can run local validation commands. See [Write and Validate Custom Scripts](/docs/data/data-cloud-code-ext/guide/author-custom-script.md).
- 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](https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.md#create-a-salesforce-app) and [Obtain Credentials](https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.md#obtain-credentials) sections of [Get Started with the Agent API](https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.md).

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

```python
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.

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

   ```python
   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")
   ```

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.

   ```python
       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"]).

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

   ```python
       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.

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.

   ```bash
   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.

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

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

## See Also

- [Write and Validate Custom Scripts](/docs/data/data-cloud-code-ext/guide/author-custom-script.md)
- [Considerations When Writing Code Extension](/docs/data/data-cloud-code-ext/guide/considerations-when-writing-custom-code.md)
- [Invoke Code Extension by Creating a Batch Data Transform](/docs/data/data-cloud-code-ext/guide/create-batch-transform.md)
- [Get Started with Agent API](https://developer.salesforce.com/docs/ai/agentforce/guide/agent-api-get-started.md)