Classifier models

Every other model in this provider writes text. A classifier model does not: you give it some text and a typed question, and it answers with a value from a set you named in advance, plus a confidence. Ask it for a string and the request is refused before it leaves your process.

TypeSafe’s Jev is the one pydantic-ai supports, as the typesafe: provider. Nothing in this provider is specific to it – it arrives through the same PydanticAIHook as every other model, so a model id is the whole integration.

Setup

  1. Install the extra:

    pip install 'apache-airflow-providers-common-ai[typesafe]'

    The extra installs the TypeSafe SDK. The typesafe: model adapter is part of pydantic-ai itself from pydantic-ai-slim 2.45.0, which is newer than the floor the provider’s other extras set, so check that release or later is installed.

  2. Create a connection (Admin > Connections):

    • Connection Id: jev_default

    • Connection Type: Pydantic AI

    • Password: your TypeSafe API key, from your TypeSafe account

    • Extra: {"model": "typesafe:jev-1.13.0"}

Leave Host empty unless you are pointing at a proxy; the provider defaults to TypeSafe’s own endpoint.

Pin the version rather than using jev-latest. A threshold you tuned against one release is not guaranteed to mean the same thing after the next one, and jev-latest moves under you.

When a classifier model is the right choice

All four of these have to hold.

The answer is a label, a number, or a pick, not prose. A bool, a Literal or Enum of strings, a bounded float, a whole-number rubric, or a list of picks. A str field is refused, so anything that writes a summary, a query, a migration, or a reply to a person is out.

You can name the options up front, and there are at most 255. Downstream task ids, error categories, severity levels, environments, a repository list. If the set is open, or is discovered at run time and could grow past the cap, this is the wrong tool. The cap is per question, so tools attached to an agent get their own 255, with the output type as one option in that question.

The decision is on a path where latency or cost is the constraint. One decision per Dag run rarely justifies changing models. One per row, per file, or per retrieved document does, and so does a branch a scheduler is waiting on.

You want a number to branch on, not a sentence to trust. A pick, a rubric, and a yes/no each come back with a confidence, so “act automatically above 0.8, ask a human below it” becomes something you can write down. A bounded float is the exception: there the probability is the answer, so nothing separate is reported and you gate on the value itself. If you would not do anything different with a confidence of 0.55 than with 0.95, that is a sign a general-purpose model is fine here.

And one case where the answer is neither: if a deterministic rule already sorts the input correctly, use the rule. A classifier model is cheap, not free, and a rule you can read is worth more than a probability you have to calibrate.

Where it fits in this provider

Surface

Fits

Why

LLMBranchOperator

Yes, with a caveat

The downstream task ids are already presented to the model as a constrained set of choices, which is exactly the shape a classifier model answers. Setting model_id is the only change, as long as there are two or more downstream tasks – a one-option pick is refused. Describe each branch in branches and set a decision_policy so an unsure pick goes to a person instead of branching; see Branch on an answer: LLMBranchOperator.

LLMOperator / AgentOperator with a typed output_type

Yes

A Literal, Enum, bool or bounded number works. Describe the field, which becomes the question, and describe each option, which is what tells them apart. An option with no description is read from its name alone.

ClassifierRetryPolicy

Yes

The model names one of the policy’s categories and nothing else; retry or fail, the delay and the confidence bar come from each category’s entry in the worker. Set min_confidence and an unsure answer goes to fallback_policy (typically an LLMRetryPolicy on a text model), then fallback_rules, then the task’s own retry behaviour, instead of ending the task on the model’s say-so. LLMRetryPolicy itself asks for free text, which a classifier model refuses. This is the surface where the model’s speed and price matter most: it runs on every task failure.

Agents with toolsets

Partly

Which tool the text calls for is itself a pick, so a classifier model can make it. What it cannot write is a tool’s arguments. A tool taking none it calls itself; one taking arguments raises ToolCallProposed after the request, which is a ModelAPIError rather than a refusal, so FallbackModel hands those requests to a text model behind it and only they cost a full call.

Reading the confidence

The confidence lives in provider_details on the model response. It is reported per output field, and a bare output type is reported under "response". A bounded float field reports none at all, because there the probability is the answer.

Two surfaces act on it for you. LLMBranchOperator and LLMOperator take a decision_policy whose min_confidence sends an unsure answer to a person, or fails the task, before anything downstream runs on it, and record the confidence, the probabilities and the bar in the decision XCom (see Branch on an answer: LLMBranchOperator). ClassifierRetryPolicy takes the same min_confidence and hands an unsure answer to fallback_policy, then its deterministic fallback rules. In the branch operator and the retry policy, a per-option bar lets the choice whose wrong pick costs most demand more certainty than the rest.

Outside those, read it yourself. AgentOperator carries it inside the message_history transcript when that is enabled, and a hook-level call has it on the result:

from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
from airflow.sdk import task
from typing import Literal


@task
def triage(log_line: str) -> dict:
    agent = PydanticAIHook(llm_conn_id="jev_default").create_agent(
        output_type=Literal["transient", "resource", "permanent"],
        instructions="Classify why this Airflow task failed.",
    )
    result = agent.run_sync(log_line)
    details = result.response.provider_details or {}
    confidence = (details.get("confidence") or {}).get("response")
    return {"category": result.output, "confidence": confidence}

Then branch on the returned confidence in a downstream task, so an unsure answer escalates instead of acting. Use a higher bar for acting automatically than for flagging something for review: collapsing both into one number is the easier thing to tune and the wrong shape for the decision.

Worked example

example_classifier_model.py has both halves: a branch whose only classifier-specific line is model_id, and a classification that escalates when the confidence is low.

What it answers badly

Read pydantic-ai’s model page and TypeSafe’s own documentation before you trust a number from one of these models. Two of its failure modes matter more than the rest in a Dag:

  • The text is treated as data, not as hostile. An injected instruction, a misleading framing, or an argument for its own answer can move the result. A guard built on a classifier model belongs alongside deterministic checks, not instead of them.

  • Option order is part of what the model sees. Reordering a Literal’s members can change the answer, so a threshold measured against one ordering is measured against that ordering only.

Was this entry helpful?