# Use Predictive Model Calls in Custom Chunking Functions

Call Einstein Studio predictive models from custom chunking functions when you want to enrich chunks with model predictions, such as regression values, scores, or classifications. This approach passes structured feature columns to a deployed model and maps the prediction into your chunk output.

In a custom chunking function, your code controls how content is split into chunks, either by splitting element text into smaller pieces or by passing each element through as a single chunk. See [Write a Custom Chunking Function](/docs/data/data-cloud-code-ext/guide/write-custom-chunking-function.md). Predictive calls take structured feature fields and return a prediction value that you can use anywhere in your chunking logic, such as to enrich `citations`, filter chunks, or classify chunk types. Build the feature fields from any data available in your function, such as fields in the `SearchIndexChunkingV1Request` input or values that you compute in your code.

| 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 use predictive model call in custom chunking logic: | Permission set: <ul><li>Data Cloud Architect</li></ul> |

## Prerequisites

- Set up Salesforce CLI, the Data Code Extension plugin, and the Python SDK. See [Set Up Salesforce CLI for Custom Functions](/docs/data/data-cloud-code-ext/guide/set-up-sdk-for-functions.md).
- Initialize a function package and make sure you can run local function validation commands. See [Write a Custom Chunking Function](/docs/data/data-cloud-code-ext/guide/write-custom-chunking-function.md).
- Use a Salesforce org with at least one Einstein Studio predictive model deployed and activated. Note the model API name.
- Configure an external client app to authenticate calls to predictive 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 Predictive Model Calls to Enrich Chunks

Implement this flow inside the `function(request, runtime)` callable in `payload/entrypoint.py`: feature construction, prediction request, prediction invocation, response parsing, and chunk output mapping.

1. Read the document elements to chunk from `request.input`. For each element, run steps 2-5 to generate its enriched chunk.
2. Build the feature columns for the model. For each feature, use `PredictionColumBuilder` to set the column name and typed values, and use `set_string_values` for text features or `set_double_values` for numeric features. Read feature values from the element's metadata (for example, `source_dmo_fields`) or from other content, such as an LLM-generated summary.
3. Build a prediction request with `PredictionRequestBuilder`: set the prediction type (for example, `PredictionType.REGRESSION`), the model API name, and the prediction columns.
4. Call `runtime.einstein_predictions.predict(...)`.
5. Check `prediction_response.is_success` and that the result type indicates success (for example, `RegressionPredictionSuccess`), then read the value from `prediction_response.data` (for example, `results[0]["prediction"]["predictedValue"]`).
6. Map the prediction into the chunk (for example, add it to `citations`), and build a chunk (`text`, `seq_no`, `chunk_type`). Aggregate chunks across all elements with a continuous `seq_no`, and return `SearchIndexChunkingV1Response(output=chunks)`.

For a reference implementation, see `example/chunking_with_prediction/entrypoint.py` in the initialized function package.

Handle failed predictions explicitly. Log failures, decide whether to continue with a fallback value or raise a controlled error, and avoid placing sensitive data in logs.

## Validate Predictive Model Calls Locally

1. From your function package root, confirm that your function entrypoint and test payload files are ready for local validation.

2. Authenticate with your org.

   ```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's domain and `{CONSUMER_KEY}` with your external client app consumer key.

3. Run local validation against your function entrypoint and test payload.
   - `sf data-code-extension function run -e payload/entrypoint.py -t tests/test.json`

## See Also

- [Write a Custom Chunking Function](/docs/data/data-cloud-code-ext/guide/write-custom-chunking-function.md)
- [Considerations When Writing Custom Chunking Functions](/docs/data/data-cloud-code-ext/guide/considerations-when-writing-custom-function.md)
- [Use LLM Gateway Calls in Custom Chunking Functions](/docs/data/data-cloud-code-ext/guide/use-llm-gateway-in-custom-function.md)
- [Use Custom Functions in Data 360](/docs/data/data-cloud-code-ext/guide/use-custom-function.md)