Skip to content

Add support for TypeSafe Jev classifier models - #73363

Merged
kaxil merged 1 commit into
apache:mainfrom
astronomer:classifier-model-support
Sep 20, 2026
Merged

kaxil merged 1 commit into
apache:mainfrom
astronomer:classifier-model-support

Conversation

@kaxil

@kaxil kaxil commented Sep 18, 2026 •

Copy link
Copy Markdown
Member

A classifier model answers typed questions and refuses anything else. You give it text and a question whose answers you named in advance, and it returns one of them plus a confidence; ask it for a string and the request is rejected before it leaves your process. pydantic-ai 2.45.0 added TypeSafeModel for TypeSafe's Jev in pydantic/pydantic-ai#8450, which is the first model of that shape it supports.

It needs no provider code. TypeSafeProvider.__init__ accepts exactly the api_key and base_url that PydanticAIHook._get_provider_kwargs already passes, and infer_model("typesafe:jev-1.13.0") resolves through the same path as every other model, so setting model_id or the connection's Model field is the whole integration. The diff is one [project.optional-dependencies] line, a guide, an example Dag, and the dependency plumbing that lets the extra resolve. Nothing under src/ changes except the new example, which is why there are no new tests.

Over half the diff is the guide. "It works" is not much use to a Dag author facing a model that cannot write text, so the guide covers when a classifier model is the right choice, where it fits among the operators already here, and where it does not.

LLMBranchOperator(
    task_id="route_failure",
    prompt=failure_text,
    llm_conn_id="jev_default",
    model_id="typesafe:jev-1.13.0",
    system_prompt="Pick the remediation that addresses the cause, not the symptom.",
)

Design rationale

Opt-in extra rather than a base dependency. The extra pulls typesafe-sdk, which is pre-1.0. Keeping it opt-in means its churn cannot break a plain provider install, which is the same reasoning already applied to code-mode and shields.

The extra installs the SDK only and does not raise the pydantic-ai floor. The adapter itself ships inside pydantic-ai-slim from 2.45.0, so the obvious extra is pydantic-ai-slim[typesafe]>=2.45.0. That is what the first push carried, and CI showed why it cannot be: from 2.31 onwards pydantic-ai-slim's openai extra requires openai>=3.8, while llama-index-llms-openai (this provider's llamaindex extra) and the Google provider's litellm pin openai<3, and the workspace lock resolves every provider's extras together. Any floor above 2.31 is therefore unsatisfiable until those packages move to openai 3, which is a repo-wide migration rather than something to carry here. The extra is typesafe-sdk>=0.6.0; the guide tells the reader to have pydantic-ai-slim 2.45.0 or later installed, and on an older pydantic-ai the typesafe: prefix fails the way any unknown provider does. One temporary exclude-newer pin remains, for typesafe-sdk, whose first release is inside the rolling four-day window; it follows the existing pydantic-ai-skills and google-cloud-aiplatform entries and is deletable once the window passes 2026-09-18. uv.lock gains typesafe-sdk and nothing else.

The guide documents four limits rather than claiming a clean fit, each verified against the live API rather than inferred:

  • A bounded float field carries no confidence, because there the probability is the answer. This is the trap most likely to bite, since the natural gating pattern (confidence or 0.0) silently reads as zero confidence on a model that was certain.
  • Which tool a call needs is itself a pick, so a classifier model can make it, but it cannot write a tool's arguments. A zero-argument tool it calls itself; an argument-taking one raises ToolCallProposed, a ModelAPIError rather than a refusal, which FallbackModel hands to a text model behind it.
  • LLMBranchOperator needs two or more downstream tasks. The operator permits one, and a one-option pick is refused.
  • LLMRetryPolicy does not fit: three of ErrorClassification's four fields are shapes Jev refuses.

Placed with the model-provider extras, not on its own. The group's stated membership rule is that each extra name mirrors the identically named pydantic-ai-slim group, which typesafe does, and it is selected the same way as the others, through extra["model"] on a generic pydanticai connection. The bullet says it differs in kind and links the guide, which seemed better than inventing a fifth extras category for one member.

Gotchas

Confidence is not reachable from the operators, and that limits what branching on this model can do. It lives in provider_details on the model response, which no operator exposes on XCom by default. AgentOperator carries it inside the message_history transcript when that is enabled, and otherwise it is dropped. A branch is decided inside LLMBranchOperator before any downstream task can look at it, so when two options both fit the evidence the branch turns on a margin nothing surfaces. The guide says this plainly and shows the hook-direct pattern for gating, and the example Dag uses a separable branch for the operator and keeps the confidence gate in a task.

jev-latest moves when TypeSafe ship a release, so the guide and the example pin typesafe:jev-1.13.0. A confidence threshold tuned against one version is not guaranteed to mean the same thing after the next.

Follow-ups

Upstream's own confidence gate is a response handler on FallbackModel:

FallbackModel("typesafe:jev-latest", "openai:…", fallback_on=[ModelAPIError, unsure])

That is currently unreachable from every operator in this provider, because create_agent passes the model positionally, so agent_params={"model": FallbackModel(...)} raises TypeError: got multiple values for argument 'model'. A model= passthrough would close the gap above without this provider inventing any routing policy of its own. Left out of this PR to keep it to the extra and the docs.

@kaxil
kaxil force-pushed the classifier-model-support branch from 2c15cdd to 703bafd Compare September 20, 2026 22:13
A classifier model answers typed questions -- pick one of these options, a
bounded score, a yes/no probability -- and refuses anything that would need
prose. pydantic-ai 2.45.0 added ``TypeSafeModel`` for TypeSafe's Jev, and its
provider takes exactly the ``api_key`` and ``base_url`` that
``PydanticAIHook`` already passes, so no provider code changes: setting
``model_id`` or the connection's Model field is the whole integration.

Adds the opt-in ``typesafe`` extra, a guide covering when a classifier model is
the right choice and where it fits among the existing operators, and an example
Dag with a branch and a confidence-gated classification.

The extra pulls ``typesafe-sdk``, which is pre-1.0, so it stays opt-in rather
than a base dependency. The temporary per-package ``exclude-newer`` pins are
needed because 2.45.0 was published inside the rolling four-day window.
@kaxil
kaxil force-pushed the classifier-model-support branch from 703bafd to 48baa4a Compare September 20, 2026 22:19
@kaxil
kaxil marked this pull request as ready for review September 20, 2026 23:31
@kaxil
kaxil merged commit 134b733 into apache:main Sep 20, 2026
272 of 287 checks passed
@kaxil
kaxil deleted the classifier-model-support branch September 20, 2026 23:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants