Calling LLMs from custom functions in flows

In single-tenant environments with advanced view enabled, you can call an LLM from custom function code to get text or structured output—for example, to drive extraction, classification, or validation with AI. Implement Python in the project’s legacy scripts folder or module, register entrypoints, and reference functions by name in the step configuration. Model and provider come from the tenant’s configured LLM provider and the AI runtime model.

Calling LLMs from custom functions works differently in the app editor, see Calling LLMs from custom functions.

Availability

In the flow editor, llm_client is supported in these steps and event hooks:

• Pre-flow UDF

• Post-flow UDF

• Apply classifier

• Apply refiner

• Map UDF

• Reduce UDF

Design considerations

When you use the LLM client in flow steps, consider the following:

• The LLM client bypasses the platform’s built-in guardrails. Use it only where the dedicated agent steps can’t meet the requirement — for example, where custom preprocessing logic must run before an LLM call, or where an external data source must be queried and incorporated into a prompt. In all other cases, use the Agent classifier and Agent extract steps.

• Contain LLM client calls to a single, well-defined step. Don’t scatter custom LLM calls across multiple UDFs throughout the flow.

• When you call an LLM directly through the client rather than through a dedicated agent step, the following don’t apply automatically:

  • Grounding — The model isn’t constrained to return values from the source document. Enforce grounding through careful prompting. For example, if using the LLM client for classification, the prompt must explicitly instruct the model to return only one of the defined class names and handle cases where the document doesn’t match any category.

  • Schema validation — The platform doesn’t validate that LLM output matches an expected structure. Output parsing and error handling are the developer’s responsibility. Error handling must account for both external API failures and LLM output validation failures.

  • Rate limit management — Multiple LLM client calls across a flow at high document volume can quickly exhaust Azure OpenAI or Vertex AI rate limits, causing flows to fail.

Using the LLM client in flow steps

Calling an LLM from a flow step involves the following components and processes.

• Legacy scripts — Add your Python function to the project’s legacy scripts folder or module, then reference the function by name in the step configuration.

• Registration — Decorate the entrypoint with @register_fn(provenance=False) from instabase.provenance.registration so the platform can register and invoke it.

• Clients — Declare CLIENTS as a positional argument in your function signature. The flow runtime passes it in automatically. Use attribute access to reference clients — for example, CLIENTS.llm_client and CLIENTS.ibfile — rather than dict-style access with .get('llm_client').

• LLM usage — Call generate_content on llm_client using the same parameters described in the Call generate_content subsection below.

from instabase.provenance.registration import register_fn
@register_fn(provenance=False)
def llm_step(CLIENTS, *args, **kwargs):
    llm_client = CLIENTS.llm_client
    resp_text = llm_client.generate_content(
        prompt="<Task and source text the model must use.>",
    )
    return resp_text

When you wire the function in the flow—for example for a pre-flow or post-flow event—invoke it by passing CLIENTS according to your step’s formula, such as llm_step(CLIENTS). Use the function name you registered and the variable order your step expects (see Custom functions in flow).

generate_content(
    prompt='str'
    file_data=None,
    mime_type=None,
    file_path=None,
    response_schema=None,
    enable_thinking=True,
    enable_logprobs=False,
)

Example: Map UDF with structured output

The following example uses @register_fn, takes CLIENTS as a positional argument together with Map UDF variables, reads a file with CLIENTS.ibfile, and writes model output into out_files. Set the Map UDF step formula to match your declared parameters—for example map_with_llm(INPUT_RECORD, STEP_FOLDER, CLIENTS).

import json
from instabase.provenance.registration import register_fn
@register_fn(provenance=False)
def map_with_llm(input_record, step_folder, CLIENTS, *args, **kwargs):
    if CLIENTS.llm_client is None:
        return {"out_files": []}
    input_filepath = input_record["input_filepath"]
    file_content, err = CLIENTS.ibfile.read_file(input_filepath)
    if err:
        raise IOError("Could not read file {}".format(input_filepath))
    resp = CLIENTS.llm_client.generate_content(
        prompt="Return the deductions table. Remove any extra symbols from the amount deducted.",
        file_data=file_content,
        mime_type="application/pdf",
        file_path=input_filepath,
        response_schema={
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "description": {"type": "string"},
                    "amount deducted": {"type": "float"},
                },
            },
            "description": "deductions table",
        },
    )
    if isinstance(resp, (dict, list)):
        payload = json.dumps(resp).encode("utf-8")
    else:
        payload = str(resp).encode("utf-8")
    return {
        "out_files": [
            {
                "filename": "deductions.json",
                "content": payload,
            }
        ]
    }

Adjust filenames, MIME type, schema, formula, and return shape to match your step.