Classify reviews in bulk at half the price

A day of product reviews needs sentiment labels, nobody needs them in the next minute, and a hundred thousand full-price calls is hard to justify. This Dag sends them all to the vendor’s batch API in one job at about half the per-request price, with up to 24 hours’ turnaround. Airflow defers while the batch runs so no worker is held, re-attaches to the same batch on retry instead of paying twice, and lands results as JSONL on object storage with only a manifest in XCom.

What this demonstrates

  • Batch processing: LLMBatchOperator – LLMBatchOperator with deferrable=True and a result_path keyed by run_id.

  • Structured output and XCom – every request in the batch asks for the same Sentiment schema, and rows that fail validation are kept with their raw text.

  • ObjectStoragePath – the downstream task reads the landed rows back from storage rather than through XCom.

Run it

  1. Install the provider with the OpenAI extra. Change model_id to run the same batch on Anthropic:

    pip install "apache-airflow-providers-common-ai[openai]"
    
  2. Change RESULT_ROOT in the file to a bucket you can write to, with an object storage connection for it.

  3. Trigger the Dag. It defers until the vendor finishes the batch, which can take hours:

    airflow dags test example_llm_batch_operator
    

The classify_reviews XCom is a manifest with the result URI and counts by status. summarize_manifest reads the JSONL and logs a count per label.

The Dag

The output schema and the reader task, then the Dag body:

airflow/providers/common/ai/example_dags/example_llm_batch.py[source]

class Sentiment(BaseModel):
    """Structured output requested from every request in the batch."""

    label: Literal["positive", "negative", "neutral"]
    confidence: float = Field(ge=0, le=1)
    reason: str


airflow/providers/common/ai/example_dags/example_llm_batch.py[source]

@task
def summarize_manifest(manifest: dict) -> dict[str, int]:
    """Read the landed JSONL rows back; the XCom value is only the manifest."""
    rows = [
        json.loads(line)
        for line in ObjectStoragePath(manifest["result_uri"]).read_text().splitlines()
        if line
    ]
    by_label: dict[str, int] = {}
    for row in rows:
        if row["status"] == "success":
            by_label[row["output"]["label"]] = by_label.get(row["output"]["label"], 0) + 1
        else:
            # ``error`` (provider-side failure) or ``invalid_output`` (schema mismatch,
            # original text kept in ``raw_output``); the manifest's ``counts`` has the totals.
            print(f"request {row['index']} -> {row['status']}: {row['error'] or row['raw_output']}")
    return by_label


airflow/providers/common/ai/example_dags/example_llm_batch.py[source]

@dag(tags=["example"])
def example_llm_batch_operator():
    classify = LLMBatchOperator(
        task_id="classify_reviews",
        requests=[f"Review: {review!r}" for review in REVIEWS],
        result_path=f"{RESULT_ROOT}/classify",
        llm_conn_id="pydanticai_default",
        model_id="openai:gpt-5-mini",
        system_prompt="Classify the sentiment of this product review.",
        output_type=Sentiment,
        deferrable=True,
    )
    summarize_manifest(classify.output)


Adapting it

  • Replace REVIEWS with a task that pulls the day’s reviews from your store and pass its output as requests. The operator accepts a list of prompts or of per-request overrides; see Batch processing: LLMBatchOperator.

  • Keep run_id in result_path. A timestamp would break re-attachment on retry and let a flaky worker submit the batch twice; Batch processing: LLMBatchOperator explains why.

  • Load the JSONL into your warehouse from a downstream task instead of counting labels, and emit an Asset so reporting Dags can run when the day’s labels land.

Was this entry helpful?