AIPrompt: one Spark transformer for AI inference

AIPrompt extends the familiar OpenAIPrompt interface with explicit API selection, JSON request options and native Microsoft Decision-1 results. This notebook demonstrates the supported public workflows using small, synthetic datasets. It complements the OpenAI notebook and the Unified AI prompting guide.

Choose the API your endpoint implements, not just a model name. The five routes share an interface, but preserve their own request and response contracts. API selection does not imply provider or model availability.

API	Examples in this notebook
chat_completions	Templates, messages, generation options and structured output
responses	Text, request-body extensions, structured output and response chaining
embeddings	Vectors, dimensions and per-row options
systemone	Decision-1 through AIPrompt and the direct AISystemOne class
decisions	Foundry gpt-6-luna through AIPrompt and the direct OpenAIDecisions class

The examples also cover output/error/usage columns, parameter inspection and copying, SparkML persistence, pipelines and pandas interoperability. Recovery and durable Delta output are a separate feature, not part of this notebook.

Compatibility and migration

AIPrompt preserves supported scalar/column parameter bindings (the last setter wins), copying and save/load, but it is not a drop-in replacement for every OpenAIPrompt behavior. Existing legacy stages keep their own contracts. Inheriting a setter does not make that option valid for every API.

Area	What to expect from AIPrompt
Request options	Explicitly unsupported typed options fail before HTTP. Chat-only options include n, stop, penalties and seed; Responses-only options include store and previousResponseId; embeddings supports dimensions and user. The provider still decides which options its model accepts.
Wire mapping	Token-limit aliases map to max_completion_tokens for Chat and max_output_tokens for Responses. Responses maps responseFormat to text.format, verbosity to text.verbosity, and reasoningEffort to reasoning.effort.
API versions	Chat, Responses and OpenAIDecisions default to v1, without an api-version query. For Chat and Responses, the default (or explicit setApiVersion("v1")) selects the v1 route from a resource root. Configured version defaults and explicit overrides remain supported; other strings such as preview or 2025-04-01 are forwarded as api-version, without a version allowlist. setApiVersionCol supports per-row selection. An existing v1 URL keeps its path and accepts explicit non-v1 query versions. OpenAIDecisions forwards other explicit versions on its native route.
Messages	messagesCol is input-only, unlike legacy usage that could name generated messages. Choose exactly one input binding; do not combine it with a prompt template.
Responses sampling	Configured temperature and topP are forwarded even for GPT-5-named deployments by both the facade and legacy OpenAI stages. Configure only options your model supports; seed is not supported on Responses.
Native APIs	AISystemOne and OpenAIDecisions lock their respective API selectors and reuse AIPrompt's implementation. SystemOne uses named questions and a lossless answer map; Decisions uses ordered questions and typed ordered answers. Their outputs remain native structs, not text postprocessing results.
Native configuration	Provide an explicit native endpoint, model and authentication. There is no default model or implicit reuse of global Chat routing/credentials. Foundry Microsoft Decision-1 uses SystemOne; supported gpt-6-luna deployments use the separate Decisions API. Check provider and region availability for each API.
Transport	The default transport bounds requests and retries and does not automatically follow redirects or retain cookies. An explicit custom handler opts out of this default policy.

See the guide's full parameter capability table and migration notes before porting an existing stage. extraBody supports provider extensions but cannot overwrite explicitly configured typed fields or request identity.

Prerequisites and service configuration

Use a Spark environment with matching SynapseML Python and JVM artifacts that include AIPrompt and AISystemOne. See the installation guide. Import this notebook into your Spark platform, as with OpenAI.ipynb.

Provide the following environment variables through your platform's secret/configuration mechanism. No credentials are embedded or printed. Missing configuration is an error, not a skipped example.

• AZURE_OPENAI_ENDPOINT: your resource root, such as https://<resource>.openai.azure.com.
• AZURE_OPENAI_CHAT_DEPLOYMENT: a deployment supporting Chat, Responses, structured output and vision; for example, a GPT-4.1-family deployment. The generation examples use sampling parameters supported by that family. Reasoning models may require different options.
• AZURE_OPENAI_EMBEDDING_DEPLOYMENT: a text-embedding-3 deployment supporting the dimensions option.
• AZURE_OPENAI_TOKEN: a current Entra bearer token for Azure AI services, with inference permissions.
• SYSTEMONE_URL: the complete Foundry route ending in /providers/microsoft/v1/systemone.
• SYSTEMONE_DEPLOYMENT: your Microsoft Decision-1 deployment name.
• SYSTEMONE_TOKEN: a current Entra token authorized for that resource.
• DECISIONS_ENDPOINT: a Foundry resource root exposing /openai/v1/decisions in a supported region.
• DECISIONS_DEPLOYMENT: your gpt-6-luna deployment name.
• DECISIONS_TOKEN: a current Entra token authorized for that resource.

The examples use setAADToken. API-key callers can use setSubscriptionKey instead. Do not put credentials in a prompt, extraBody, notebook output or a saved stage. No OpenAI Python SDK is required.

Cost and replay: these are real model calls. transform is lazy; an action such as collect, show or toPandas sends requests. Repeated actions, task retries or cache eviction can repeat paid requests. Each example below materializes its inference result once. Neither Spark caching nor the provider's store option guarantees exactly-once execution.

import json
import os
from pathlib import Path
from tempfile import TemporaryDirectory

from pyspark.ml import Pipeline, PipelineModel
from pyspark.ml.linalg import Vector
from pyspark.sql import Row, SparkSession, functions as F
from synapse.ml.services.ai import AIPrompt, AISystemOne
from synapse.ml.services.openai import OpenAIDecisions

spark = SparkSession.builder.getOrCreate()


def required_setting(name):
    value = os.environ.get(name)
    if not value or not value.strip():
        raise ValueError(f"Set {name} before running this notebook")
    return value.strip()


openai_endpoint = required_setting("AZURE_OPENAI_ENDPOINT").rstrip("/")
chat_deployment = required_setting("AZURE_OPENAI_CHAT_DEPLOYMENT")
embedding_deployment = required_setting("AZURE_OPENAI_EMBEDDING_DEPLOYMENT")
openai_token = required_setting("AZURE_OPENAI_TOKEN")
systemone_url = required_setting("SYSTEMONE_URL")
systemone_deployment = required_setting("SYSTEMONE_DEPLOYMENT")
systemone_token = required_setting("SYSTEMONE_TOKEN")
decisions_endpoint = required_setting("DECISIONS_ENDPOINT")
decisions_deployment = required_setting("DECISIONS_DEPLOYMENT")
decisions_token = required_setting("DECISIONS_TOKEN")

completed_examples = []


def prompt_for(api_type):
    route = {
        "chat_completions": "chat/completions",
        "responses": "responses",
        "embeddings": "embeddings",
    }[api_type]
    deployment = embedding_deployment if api_type == "embeddings" else chat_deployment
    stage = (
        AIPrompt()
        .setApiType(api_type)
        .setUrl(f"{openai_endpoint}/openai/v1/{route}")
        .setDeploymentName(deployment)
        .setAADToken(openai_token)
        .setOutputCol("answer")
        .setErrorCol("error")
        .setUsageCol("usage")
        .setRawResponseCol("raw_response")
        .setConcurrency(1)
        .setTimeout(120)
    )
    if api_type == "responses":
        stage.setStore(False)
    return stage


def validate_example(name, rows):
    assert rows, f"{name}: expected at least one result"
    errors = [row["error"] for row in rows if row["error"] is not None]
    assert not errors, f"{name}: {errors}"
    assert all(row["answer"] is not None for row in rows), f"{name}: missing answer"
    assert all(
        json.loads(row["raw_response"]) for row in rows
    ), f"{name}: missing raw JSON"
    completed_examples.append(name)
    print(f"Passed: {name} ({len(rows)} rows)")
    return rows


def run_example(name, stage, data):
    return validate_example(name, stage.transform(data).collect())


countries = spark.createDataFrame([("France",)], ["country"])
tickets = spark.createDataFrame(
    [("Checkout is unavailable. Every customer receives an error.",)], ["text"]
)

Prompt templates and generation options

Select the API explicitly. Template placeholders bind to input columns. A complete URL is used as supplied rather than having another route appended. For Chat, the normal output is the first choice's text; rawResponseCol retains the full response, including all choices.

Set only options supported by your chosen API and model. Explicit options are not silently dropped.

chat = (
    prompt_for("chat_completions")
    .setSystemPrompt("Answer factual questions briefly.")
    .setPromptTemplate("What is the capital of {country}?")
    .setMaxTokens(128)
    .setTemperature(0.0)
    .setTopP(1.0)
    .setUser("aiprompt-notebook")
)
planned = chat.transform(countries)
planned.printSchema()  # Schema planning does not call the model.
chat_rows = run_example("chat/template", chat, countries)
print(chat_rows[0]["answer"])
print(chat_rows[0]["usage"])

Column-driven settings, constructor parameters and setParams

Service parameters support scalar and column bindings. The last setter for the same parameter wins. An explicit deployment or deployment column takes precedence over applicable global defaults; a missing per-row value is an error, not an instruction to switch models.

per_row = countries.withColumn("deployment", F.lit(chat_deployment)).withColumn(
    "temperature", F.lit(0.0)
)
column_chat = (
    chat.copy({})
    .setDeploymentNameCol("deployment")
    .setTemperatureCol("temperature")
    .setParams(promptTemplate="Name the capital of {country} in one word.")
)
run_example("chat/column-parameters", column_chat, per_row)

constructed = AIPrompt(apiType="responses", promptTemplate="Say hello to {country}.")
constructed.setParams(extraBody={"metadata": {"example": "constructor"}})
assert constructed.getApiType() == "responses"
assert constructed.getExtraBody()["metadata"]["example"] == "constructor"
assert constructed.hasParam("extraBody")
assert constructed.isSet(constructed.apiType)

Supply messages directly

Use setMessagesCol when you already have role/content messages. Use an explicit Spark schema. Do not combine raw messages with a prompt template, inputCol, systemPrompt or template-attachment settings; put system messages in the array instead.

messages = spark.createDataFrame(
    [([{"role": "user", "content": "Explain Spark laziness in one sentence."}],)],
    "messages ARRAY<STRUCT<role:STRING,content:STRING>>",
)
for api_type in ("chat_completions", "responses"):
    message_prompt = prompt_for(api_type).setMessagesCol("messages")
    run_example(f"{api_type}/raw-messages", message_prompt, messages)

Responses and request-body extensions

setExtraBody accepts a Python dictionary or a JSON-object string. Keys use the selected provider API's wire names, not Spark setter names. A typed option that you explicitly set owns its JSON subtree: a conflicting body extension is rejected. Identity, credentials, routing and asynchronous/streaming controls cannot be overridden through this parameter.

Responses text concatenates assistant text parts in response order. Raw JSON preserves non-text outputs. AIPrompt does not execute tools or run a tool loop.

responses = (
    prompt_for("responses")
    .setPromptTemplate(
        "Explain the customer impact of this ticket in one sentence: {text}"
    )
    .setMaxTokens(256)
    .setStore(False)
    .setExtraBody({"metadata": {"example": "aiprompt-notebook"}})
)
response_rows = run_example("responses/dictionary-extra-body", responses, tickets)
print(response_rows[0]["answer"])

json_extension = responses.copy({}).setExtraBody(
    '{"metadata":{"example":"json-string"}}'
)
run_example("responses/json-string-extra-body", json_extension, tickets)

row_extensions = tickets.withColumn(
    "body_options", F.lit('{"metadata":{"example":"per-row"}}')
)
column_extension = responses.copy({}).setExtraBodyCol("body_options")
run_example("responses/column-extra-body", column_extension, row_extensions)

Structured output and Spark postprocessing

setResponseSchema configures the model's JSON Schema. setPostProcessing("json") and its jsonSchema option configure the separate Spark SQL output type. Neither setting substitutes for the other. Continue checking errors, refusals and incomplete output.

Use setResponseFormat for an existing named envelope or the text/json_object selectors. Do not pass the bare string json_schema.

answer_schema = {
    "type": "object",
    "properties": {"answer": {"type": "string"}},
    "required": ["answer"],
    "additionalProperties": False,
}

for api_type in ("chat_completions", "responses"):
    structured = (
        prompt_for(api_type)
        .setPromptTemplate("What is the capital of {country}?")
        .setResponseSchema(answer_schema, name="capital", strict=True)
        .setPostProcessing("json")
        .setPostProcessingOptions({"jsonSchema": "answer STRING"})
    )
    structured_rows = run_example(
        f"{api_type}/schema-and-postprocessing", structured, countries
    )
    assert isinstance(structured_rows[0]["answer"]["answer"], str)

formatted = (
    prompt_for("chat_completions")
    .setPromptTemplate(
        "Return a JSON object with an answer naming the capital of {country}."
    )
    .setResponseFormat({"type": "json_object"})
    .setPostProcessing("json")
    .setPostProcessingOptions({"jsonSchema": "answer STRING"})
)
run_example("chat/json-object-format", formatted, countries)

envelope = formatted.copy({}).setResponseFormat(
    {
        "type": "json_schema",
        "json_schema": {"name": "capital", "strict": True, "schema": answer_schema},
    }
)
run_example("chat/explicit-schema-envelope", envelope, countries)

Delimiter and regular-expression postprocessing

csv splits the trimmed response using a Spark regular-expression delimiter; it is not an RFC CSV parser and does not interpret quoted fields. regex extracts the requested capture group. regexGroup is required. Supplying only delimiter, regex or jsonSchema options infers the corresponding mode; do not combine different option families.

list_input = spark.createDataFrame(
    [("Return exactly one;two;three with no extra text or formatting.",)],
    ["instruction"],
)
delimited = (
    prompt_for("responses")
    .setInputCol("instruction")
    .setPostProcessing("csv")
    .setPostProcessingOptions({"delimiter": ";"})
)
delimited_rows = run_example("postprocessing/delimiter", delimited, list_input)
assert delimited_rows[0]["answer"] == ["one", "two", "three"]

code_input = spark.createDataFrame(
    [("Return exactly CODE: 42 with no other text.",)], ["instruction"]
)
regex_prompt = (
    prompt_for("responses")
    .setInputCol("instruction")
    .setPostProcessingOptions({"regex": r"CODE:\s*(\d+)", "regexGroup": "1"})
)
regex_rows = run_example("postprocessing/regex", regex_prompt, code_input)
assert regex_rows[0]["answer"] == "42"

Images and files

For template attachments, mark a column as path with setColumnTypes and include it in the prompt template. setFileSizeLimitMB bounds attachment reads, not model response size. Paths can be files, HTTPS URLs or self-contained data URLs. These examples use a synthetic blue square and a one-page PDF containing Invoice total: 42 dollars.; no files are downloaded or uploaded to a provider file store.

Raw message parts are API-specific: Chat uses an image_url object with a url, whereas Responses uses an image_url string and supports input_file parts. Use a vision/file-capable deployment. Do not send arbitrary binary files to Chat.

image_data_url = "data:image/png;base64," + (
    "iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAeklEQVR4nO3PUQkAIBTAwJfEYGY3"
    "jCH8OITBAtxm7fN1wwUNaEEDWtCAFjSgBQ1oQQNa0IAWNKAFDWhBA1rQgBY0oAUNaEEDWtCA"
    "FjSgBQ1oQQNa0IAWNKAFDWhBA1rQgBY0oAUNaEEDWtCAFjSgBQ1oQQNa0IAWPHYBWyZBPPuE"
    "Fn4AAAAASUVORK5CYII="
)
file_data_url = "data:application/pdf;base64," + (
    "JVBERi0xLjQKMSAwIG9iago8PCAvVHlwZSAvQ2F0YWxvZyAvUGFnZXMgMiAwIFIgPj4KZW5kb2Jq"
    "CjIgMCBvYmoKPDwgL1R5cGUgL1BhZ2VzIC9LaWRzIFszIDAgUl0gL0NvdW50IDEgPj4KZW5kb2Jq"
    "CjMgMCBvYmoKPDwgL1R5cGUgL1BhZ2UgL1BhcmVudCAyIDAgUiAvTWVkaWFCb3ggWzAgMCA2MTIg"
    "NzkyXSAvUmVzb3VyY2VzIDw8IC9Gb250IDw8IC9GMSA0IDAgUiA+PiA+PiAvQ29udGVudHMgNSAw"
    "IFIgPj4KZW5kb2JqCjQgMCBvYmoKPDwgL1R5cGUgL0ZvbnQgL1N1YnR5cGUgL1R5cGUxIC9C"
    "YXNlRm9udCAvSGVsdmV0aWNhID4+CmVuZG9iago1IDAgb2JqCjw8IC9MZW5ndGggNTcgPj4Kc3Ry"
    "ZWFtCkJUIC9GMSAxOCBUZiA1MCA3NTAgVGQgKEludm9pY2UgdG90YWw6IDQyIGRvbGxhcnMuKSBU"
    "aiBFVAplbmRzdHJlYW0KZW5kb2JqCnhyZWYKMCA2CjAwMDAwMDAwMDAgNjU1MzUgZiAKMDAwMDAw"
    "MDAwOSAwMDAwMCBuIAowMDAwMDAwMDU4IDAwMDAwIG4gCjAwMDAwMDAxMTUgMDAwMDAgbiAKMDAw"
    "MDAwMDI0MSAwMDAwMCBuIAowMDAwMDAwMzExIDAwMDAwIG4gCnRyYWlsZXIKPDwgL1NpemUgNiAv"
    "Um9vdCAxIDAgUiA+PgpzdGFydHhyZWYKNDE4CiUlRU9GCg=="
)

for name, instruction, asset, expected in (
    ("image", "Name the main color of this image.", image_data_url, "blue"),
    ("file", "What is the invoice total in this file?", file_data_url, "42"),
):
    assets = spark.createDataFrame(
        [(instruction, asset)], "instruction STRING, asset STRING"
    )
    attachment_prompt = (
        prompt_for("responses")
        .setPromptTemplate("{instruction} {asset}")
        .setColumnTypes({"asset": "path"})
        .setFileSizeLimitMB(2.0)
    )
    asset_rows = run_example(f"attachments/template-{name}", attachment_prompt, assets)
    assert expected in asset_rows[0]["answer"].lower()

chat_image = spark.createDataFrame(
    [
        (
            [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "Name the main color of this image."},
                        {
                            "type": "image_url",
                            "image_url": {"url": image_data_url, "detail": "low"},
                        },
                    ],
                }
            ],
        )
    ],
    "messages ARRAY<STRUCT<role:STRING,content:ARRAY<STRUCT<"
    "type:STRING,text:STRING,image_url:STRUCT<url:STRING,detail:STRING>>>>>",
)
chat_image_rows = run_example(
    "attachments/raw-chat-image",
    prompt_for("chat_completions").setMessagesCol("messages"),
    chat_image,
)
assert "blue" in chat_image_rows[0]["answer"].lower()

response_assets = spark.createDataFrame(
    [
        (
            [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "input_text",
                            "text": "Name the image's main color and the invoice total.",
                        },
                        {"type": "input_image", "image_url": image_data_url},
                        {
                            "type": "input_file",
                            "file_data": file_data_url,
                            "filename": "invoice.pdf",
                        },
                    ],
                }
            ],
        )
    ],
    "messages ARRAY<STRUCT<role:STRING,content:ARRAY<STRUCT<"
    "type:STRING,text:STRING,image_url:STRING,detail:STRING,"
    "file_data:STRING,file_id:STRING,filename:STRING>>>>",
)
response_asset_rows = run_example(
    "attachments/raw-responses-image-and-file",
    prompt_for("responses").setMessagesCol("messages"),
    response_assets,
)
assert "blue" in response_asset_rows[0]["answer"].lower()
assert "42" in response_asset_rows[0]["answer"]

Responses storage and conversation chaining

setStore(True) stores a response at the provider, not a Spark result table. setResponseIdCol exposes its identifier. Supply that identifier with setPreviousResponseId or the per-row setPreviousResponseIdCol to continue a conversation. This requires provider-side storage to be enabled for your resource.

This example stores only synthetic content and deletes its own stored response in finally. It never lists or deletes other responses. The continuation calls use store=False.

from urllib.parse import quote
import requests
import sys

stored = (
    prompt_for("responses")
    .setPromptTemplate(
        "Remember the code name ALPHA. Reply with a brief acknowledgement."
    )
    .setStore(True)
    .setResponseIdCol("response_id")
)
created_response_ids = []
deleted_response_ids = []
try:
    first_rows = stored.transform(countries).collect()
    created_response_ids.extend(
        row[stored.getResponseIdCol()]
        for row in first_rows
        if row[stored.getResponseIdCol()]
    )
    validate_example("responses/store-and-response-id", first_rows)
    response_id = first_rows[0][stored.getResponseIdCol()]
    assert response_id

    continued = (
        prompt_for("responses")
        .setStore(False)
        .setPreviousResponseId(response_id)
        .setPromptTemplate(
            "What code name did I give you? Reply with just the code name."
        )
    )
    continued_rows = run_example("responses/scalar-previous-id", continued, countries)
    assert "ALPHA" in continued_rows[0]["answer"].upper()

    with_ids = countries.withColumn("previous_id", F.lit(response_id))
    column_continuation = continued.copy({}).setPreviousResponseIdCol("previous_id")
    column_rows = run_example(
        "responses/column-previous-id", column_continuation, with_ids
    )
    assert "ALPHA" in column_rows[0]["answer"].upper()
finally:
    original_error = sys.exc_info()[1]
    cleanup_errors = []
    for response_id in created_response_ids:
        try:
            deleted = requests.delete(
                f"{openai_endpoint}/openai/v1/responses/{quote(response_id, safe='')}",
                headers={"Authorization": f"Bearer {openai_token}"},
                timeout=60,
            )
            deleted.raise_for_status()
            if deleted.status_code != 204:
                confirmation = deleted.json()
                if (
                    not isinstance(confirmation, dict)
                    or confirmation.get("deleted") is not True
                ):
                    raise ValueError("Provider did not confirm response deletion")
            deleted_response_ids.append(response_id)
        except (requests.RequestException, ValueError) as error:
            cleanup_errors.append((response_id, error))
    print(f"Deleted {len(deleted_response_ids)} response(s) created by this example.")
    if cleanup_errors:
        failed_ids = [identifier for identifier, _ in cleanup_errors]
        message = (
            f"Could not delete stored response(s) {failed_ids}; retry these deletions."
        )
        if original_error is None:
            raise RuntimeError(message) from cleanup_errors[0][1]
        print(message, file=sys.stderr)

Embeddings

Select embeddings and bind text with setTextCol. Results are Spark ML vectors, not JSON strings. The dimensions option requires a model that supports it. AIPrompt supports float embeddings; it rejects unsupported encodings before sending a request. Both scalar and per-row user values are sent to the embedding API.

embedding = (
    prompt_for("embeddings")
    .setTextCol("text")
    .setDimensions(64)
    .setUser("aiprompt-notebook")
    .setExtraBody({"encoding_format": "float"})
)
embedding_rows = run_example("embeddings/scalar-options", embedding, tickets)
assert isinstance(embedding_rows[0]["answer"], Vector)
assert len(embedding_rows[0]["answer"]) == 64

embedding_input = tickets.withColumn("dimensions", F.lit(32)).withColumn(
    "request_user", F.lit("notebook-row")
)
row_embedding = (
    embedding.copy({}).setDimensionsCol("dimensions").setUserCol("request_user")
)
row_embedding_rows = run_example(
    "embeddings/column-options", row_embedding, embedding_input
)
assert len(row_embedding_rows[0]["answer"]) == 32

Microsoft Decision-1 with AISystemOne and AIPrompt

AISystemOne is the directly usable native transformer. The facade uses the same protocol with apiType="systemone". This is not a Chat emulation and is distinct from OpenAIDecisions.

Use the complete Foundry System One route without a dated Azure OpenAI API-version parameter. Supply your deployment explicitly; the native class does not invent a default model. Questions form a named object with noul, choice and score entries.

The result's answers field is a map of complete JSON-encoded answer values. Decode the values without assuming that every answer has the same undocumented probability or confidence fields.

questions = {
    "urgent": {"type": "noul", "instructions": "Is checkout unavailable?"},
    "team": {
        "type": "choice",
        "instructions": "Which team should handle this request?",
        "criteria": {"billing": "Billing problems", "engineering": "Software outages"},
    },
    "severity": {
        "type": "score",
        "instructions": "How severe is the issue?",
        "criteria": ["Cosmetic", "Workaround available", "Service unavailable"],
    },
}


def configure_systemone(stage):
    return (
        stage.setUrl(systemone_url)
        .setDeploymentName(systemone_deployment)
        .setAADToken(systemone_token)
        .setStateCol("text")
        .setQuestions(questions)
        .setOutputCol("answer")
        .setErrorCol("error")
        .setUsageCol("usage")
        .setRawResponseCol("raw_response")
        .setConcurrency(1)
        .setTimeout(120)
    )


systemone = configure_systemone(AISystemOne())
native_rows = run_example("systemone/direct", systemone, tickets)
decoded_answers = {
    name: json.loads(value)
    for name, value in native_rows[0]["answer"]["answers"].items()
}
assert set(decoded_answers) == set(questions)
print(decoded_answers)

systemone_facade = configure_systemone(AIPrompt().setApiType("systemone"))
run_example("systemone/facade", systemone_facade, tickets)

JSON state and per-row questions

setStateMode("json") interprets the bound string as JSON rather than literal text. setQuestionsCol takes a JSON-string column. Scalar question bindings also accept JSON strings. As with other service parameters, the last scalar/column setter wins.

json_tickets = spark.createDataFrame(
    [(json.dumps({"service": "checkout", "available": False}), json.dumps(questions))],
    ["state_json", "questions_json"],
)
json_systemone = (
    systemone.copy({})
    .setStateCol("state_json")
    .setStateMode("json")
    .setQuestionsCol("questions_json")
)
run_example("systemone/json-state-and-column-questions", json_systemone, json_tickets)

string_questions = systemone.copy({}).setQuestions(json.dumps(questions))
run_example("systemone/json-string-questions", string_questions, tickets)

Row errors and lossless observations

Malformed per-row body extensions become row errors. Valid rows can still succeed. Invalid stage-wide configuration instead fails preflight. HTTP errors, provider failures, refusals, incomplete output and parse failures must not be interpreted as successful answers.

rawResponseCol contains response JSON, not authentication headers. It can contain sensitive model output: apply your normal access and retention controls.

mixed_input = spark.createDataFrame(
    [("valid", "Say hello.", "{}"), ("invalid", "Say hello.", "[]")],
    ["case", "text", "body_options"],
)
row_validation = (
    prompt_for("chat_completions")
    .setPromptTemplate("{text}")
    .setExtraBodyCol("body_options")
)
mixed_rows = {
    row["case"]: row for row in row_validation.transform(mixed_input).collect()
}
assert mixed_rows["valid"]["error"] is None
assert mixed_rows["valid"]["answer"] is not None
assert mixed_rows["invalid"]["error"] is not None
assert mixed_rows["invalid"]["answer"] is None
completed_examples.append("validation/per-row-error")
print("Passed: valid row succeeds; malformed extension is an explicit row error")

Copy, clear and inspect parameters

These are ordinary SparkML stages. copy keeps configuration without mutating the original. clear removes an explicit binding. Do not print extractParamMap() on an authenticated stage because its parameters can include credentials.

copied = chat.copy({chat.temperature: 0.1})
assert copied.getTemperature() == 0.1
assert chat.getTemperature() == 0.0
copied.clear(copied.temperature)
copied.setTemperature(0.0)
run_example("sparkml/copy-and-clear", copied, countries)

Save/load and pipelines

Save a credential-free copy, then rebind authentication after loading with the same concrete class. This saves the stage configuration, not model results. The example uses a temporary local directory; on a distributed cluster, use a fresh caller-owned path accessible to all Spark nodes instead.

A Pipeline composes AIPrompt with other SparkML stages. Fitting this transformer-only pipeline does not perform inference; transforming and materializing its result does.

saveable = chat.copy({})
saveable.clear(saveable.AADToken)

with TemporaryDirectory(prefix="aiprompt-example-") as temporary:
    stage_path = str(Path(temporary) / "stage")
    saveable.write().save(stage_path)
    loaded = AIPrompt.load(stage_path).setAADToken(openai_token)
    run_example("sparkml/save-load", loaded, countries)

    read_loaded = AIPrompt.read().load(stage_path).setAADToken(openai_token)
    assert read_loaded.getPromptTemplate() == chat.getPromptTemplate()

    pipeline_path = str(Path(temporary) / "pipeline")
    pipeline = Pipeline(stages=[saveable]).fit(countries)
    pipeline.write().save(pipeline_path)
    restored_pipeline = PipelineModel.load(pipeline_path)
    restored_pipeline.stages[0].setAADToken(openai_token)
    pipeline_rows = restored_pipeline.transform(countries).collect()
    assert len(pipeline_rows) == 1 and pipeline_rows[0]["error"] is None
    assert pipeline_rows[0]["answer"] is not None
    completed_examples.append("sparkml/pipeline-save-load")

    native_path = str(Path(temporary) / "systemone")
    native_saveable = systemone.copy({})
    native_saveable.clear(native_saveable.AADToken)
    native_saveable.save(native_path)
    native_loaded = AISystemOne.load(native_path).setAADToken(systemone_token)
    run_example("systemone/save-load", native_loaded, tickets)

pandas interoperability

AIPrompt transforms Spark DataFrames. Convert a small pandas input to Spark and collect the small result back to pandas. Do not collect a large distributed dataset to the driver.

import pandas as pd

pandas_input = pd.DataFrame({"country": ["France"]})
pandas_result = chat.transform(spark.createDataFrame(pandas_input)).toPandas()
assert len(pandas_result) == 1
assert pandas_result["error"].isna().all()
assert pandas_result["answer"].str.len().gt(0).all()
completed_examples.append("pandas/round-trip")
print(pandas_result[["country", "answer"]])

OpenAIDecisions on Foundry

Foundry exposes /openai/v1/decisions for supported gpt-6-luna deployments. See the provider documentation for availability and region restrictions, including the East US 2 exclusion. The Decision-1 examples above use AISystemOne, a different API.

These examples exercise the direct class and facade with an ordered predicate question and a typed probability answer. Model discovery alone does not establish live endpoint support; each example checks the inference result.

decision_questions = [
    {"type": "predicate", "name": "urgent", "instructions": "Is checkout unavailable?"}
]
for name, stage in (
    ("decisions/direct", OpenAIDecisions()),
    ("decisions/facade", AIPrompt().setApiType("decisions")),
):
    stage.setUrl(decisions_endpoint).setDeploymentName(decisions_deployment)
    stage.setAADToken(decisions_token).setInputCol("text")
    stage.setQuestions(decision_questions).setOutputCol("answer")
    stage.setErrorCol("error").setRawResponseCol("raw_response")
    stage.setConcurrency(1).setTimeout(120)
    rows = run_example(name, stage, tickets)
    for row in rows:
        assert len(row.answer.answers) == 1
        answer = row.answer.answers[0]
        assert answer.name == "urgent" and answer.type == "predicate"
        assert 0 <= answer.probability <= 1

Execution summary

Every inference example above checks its result, including both native APIs. Notebook outputs are intentionally cleared in source control so credentials, endpoint details and model responses are not published.

assert len(completed_examples) == len(set(completed_examples))
print(f"Completed {len(completed_examples)} supported examples.")
for name in completed_examples:
    print(name)