Proposal: Databricks Unity AI Gateway model class for apache-airflow-providers-common-ai

[aman210122]

Hello Airflow maintainers and community. I am an enterprise AI/ML architect and a daily user of Databricks Unity AI Gateway (formerly Mosaic AI Gateway) in production. I am opening this as my first contribution thread to apache/airflow to propose adding a Databricks model class to apache-airflow-providers-common-ai, and to ask the maintainers to choose between two implementation paths before any PR lands. I would like to align on direction before writing code so that the work fits cleanly into the AIP-99 "Common Data Access Pattern + AI" scope and does not collide with the upstream pydantic-ai effort.
Background
apache-airflow-providers-common-ai (0.1.0 released April 13, 2026, currently 0.3.0 on PyPI) ships LLMOperator, AgentOperator, LLMBranchOperator, LLMSQLQueryOperator, LLMSchemaCompareOperator, and LLMFileAnalysisOperator, plus matching @task.llm, @task.agent, @task.llm_branch, @task.llm_sql, @task.llm_schema_compare, and @task.llm_file_analysis decorators. The provider exposes connection-typed hooks for OpenAI, Azure OpenAI, AWS Bedrock, Google Vertex AI, and MCP through pydantic-ai Models (PydanticAIHook, PydanticAIAzureHook, PydanticAIBedrockHook, PydanticAIVertexHook, MCPHook). The provider source lives at https://github.com/apache/airflow/tree/main/providers/common/ai and the docs at https://airflow.apache.org/docs/apache-airflow-providers-common-ai/stable/index.html.
Databricks now exposes two OpenAI-compatible model surfaces:

GA: https:///serving-endpoints (the existing model serving surface, formerly fronted by Mosaic AI Gateway features).
Beta: https:///ai-gateway/mlflow/v1 (the new Unity AI Gateway for agents and LLMs, announced by David Nasi on April 15, 2026 at https://www.databricks.com/blog/ai-gateway-governance-layer-agentic-ai, which brings AI Gateway inside Unity Catalog as the governance layer for agents, LLMs, and MCP servers).

Both surfaces speak the OpenAI Chat Completions schema, except that message.content is typed str | list[ContentItem] rather than str | None. That is exactly the divergence tracked by pydantic-ai issue pydantic/pydantic-ai#2947 and PR pydantic/pydantic-ai#4036 by @aayushchou-db, which is currently in draft (the author wrote, verbatim, "Converting this to draft, just so I can iron out the auth a bit more!").
Today, common.ai users who want to target Databricks foundation models, external models, or AI Gateway endpoints have no first-class path. They either reach for OpenAIChatModel with a Databricks base URL and hit the content-typing mismatch on streaming, or they fall outside the provider altogether. That is the gap.
Proposed scope (v1)
A PydanticAIDatabricksHook plus a pydanticai-databricks connection type that returns a pydantic-ai Model usable directly with every existing common.ai operator and decorator. v1 should:

Support both base URLs (GA /serving-endpoints and Beta /ai-gateway/mlflow/v1) selected by a simple surface connection field.
Support PAT and OAuth M2M (service principal client credentials) authentication only.
Return a Model that handles the str | list[ContentItem] content quirk locally (or delegates to upstream once available).
Live under airflow.providers.common.ai.hooks.pydantic_ai next to the existing hooks, with a databricks pip extra in pyproject.toml to pull in cross-provider dependencies.
Work without modifying the existing operators or decorators.

The fork in the road
This is the decision I need maintainer input on before I open a PR.
Path A: Wait for upstream pydantic-ai DatabricksModel. Ship a thin PydanticAIDatabricksHook in common.ai that simply imports pydantic_ai.models.databricks.DatabricksModel once PR #4036 merges, wires up the Airflow connection, and returns the Model. Pro: zero duplication, smallest surface area, future-proof. Con: blocked on upstream timing; PR #4036 has been in draft since the author wrote "Converting this to draft, just so I can iron out the auth a bit more!" and there is no visible ETA.
Path B: Ship a local DatabricksModel in common.ai now. Subclass OpenAIChatModel directly inside airflow/providers/common/ai/models/databricks.py, override _validate_completion, _get_stream_options, and _map_part_delta to handle list[ContentItem] content and the gpt-oss streaming option restriction, and delete the local class once upstream lands. Pro: unblocks users now, exercises the exact same subclassing pattern the upstream PR proposes, gives common.ai a Databricks connection type immediately. Con: temporary duplication with upstream, a small migration when upstream merges.
I am willing to implement either path. I lean toward Path B because the upstream PR's subclassing pattern is small enough that maintaining it inside common.ai for a few releases is cheap, and because the Unity AI Gateway is shipping aggressively right now. But this should be the maintainers' call.
Out of scope for v1

AI Gateway feature toggles on the client side (rate limits, guardrails, payload logging configuration via PUT /api/2.0/serving-endpoints/{name}/ai-gateway). Defer to a follow-up.
Azure SPN federation, managed identity, and Kubernetes OIDC token federation. The existing apache-airflow-providers-databricks Connection supports these and we should mirror that surface, but in a v2 PR.
MCP governance on the Unity AI Gateway side. Defer; common.ai already has an MCPHook and the agent-governance story from the April 15, 2026 announcement is large enough to deserve its own thread.
Embeddings endpoint. Defer until the chat path is stable.
Agent-specific Unity AI Gateway features (OBO execution, agent service policies). Likely a separate AIP if pursued.

Alignment with AIP-99
AIP-99 "Common Data Access Pattern + AI" (tracked on the apache GitHub Project board at AIP-99 Common Data Access Pattern + AI) already covers pydantic-ai-based model providers as the integration pattern for common.ai. Adding Databricks as another supported connection type is a scope-extension of AIP-99 rather than a new AIP. I do not propose changing the operator surface, decorator API, or toolset API. If maintainers disagree and would prefer this be its own AIP, I will write it.
Open questions for maintainers

Path A or Path B? This is the blocker.
Home of the code: common.ai (consistent with Bedrock and Vertex being inside common.ai despite the existence of apache-airflow-providers-amazon and apache-airflow-providers-google) or apache-airflow-providers-databricks (consistent with all other Databricks operators)? My read of precedent is common.ai, but I want confirmation.
Connection type naming: a databricks connection type would conflict with the existing Databricks connection in the standalone provider. pydanticai-databricks as the connection type, matching the existing pydanticai-azure, pydanticai-bedrock, pydanticai-vertex connection type pattern, plus a databricks pip extra to pull in cross-provider dependencies. Confirm.
CI/test infrastructure: common.ai mock-tests its hooks against pydantic-ai's test doubles rather than live providers. I plan to follow the same pattern (no live Databricks workspace in CI). Is that acceptable, or do you want a system test that runs against a real workspace behind a secret?
Surface selection ergonomics: should the GA-vs-Beta surface be a connection extra field, an environment variable, or two separate hook subclasses?

Closing
I am willing to submit the PR end-to-end once a path is chosen, with unit tests against pydantic-ai doubles, docs, a connection-type stub, and an example DAG. Prototype branch placeholder: . If the path choice settles outside this Discussion, I am happy to convert it into a formal issue or a draft AIP. Pinging @kaxil and @gopidesupavan since common.ai is your area, and @vikramkoka for the AIP-99 scope question.
Thanks for reading.

[kaxil]

Thanks for the detailed writeup, @aman210122.

Path A, gated on a pydantic-ai release (not just the merge).

Quick correction on the upstream status: pydantic/pydantic-ai#4036 is no longer in draft. It's open, ~3,400 lines, 42 commits, 63 review comments, assigned to @DouweM, last activity Apr 6. So the "blocked on upstream timing" framing for Path B is weaker than the proposal suggests. The work is happening, the bottleneck is review on that PR.

Why Path A:

1. common.ai's other model hooks (Azure, Bedrock, Vertex) all wrap upstream pydantic-ai Models rather than fork them. A local subclass of OpenAIChatModel here would duplicate a 3.4k-line PR that's still under active review, and we'd inherit any course corrections that review surfaces.
2. The provider is still 0.x. Adding a class we plan to delete is a deprecation chore we don't need.

Why "wait for the release," not just the merge:

The databricks extra has to pin pydantic-ai >= X.Y.Z against a real PyPI version. Until upstream cuts a release containing DatabricksModel, the hook's from pydantic_ai.models.databricks import DatabricksModel has nothing to resolve against, and pinning a git ref in a provider extra isn't something we do. So the highest-leverage place to push this forward is review on pydantic/pydantic-ai#4036 itself, not a fork PR here. Once that lands and ships in a pydantic-ai release, the common.ai wrapper is roughly a 30-line hook plus a connection-type stub plus docs.

On the other questions:

• Home: common.ai, consistent with Bedrock and Vertex precedent.
• Connection name: pydanticai-databricks plus a databricks pip extra, matching the existing pattern.
• CI: mock against pydantic-ai test doubles, same as today. No live workspace.
• Surface selection (GA vs Beta): connection extra field. Two hook subclasses is too much surface for a flag.

The AI Gateway feature toggles, embeddings, agent-specific bits are correctly out of v1 scope.

Concretely: nothing to open on apache/airflow right now. Let's watch pydantic/pydantic-ai#4036 land and get into a pydantic-ai release, and then the common.ai PR is small and mechanical. Happy to review it when that point comes.

[aman210122]

Thanks @kaxil, agreed on Path A and gating on a pydantic-ai release. Confirmed on home (common.ai), connection type (pydanticai-databricks), pip extra (databricks), CI approach (pydantic-ai test doubles), and surface selection (connection extra field). I will track pydantic/pydantic-ai#4036 and engage on review there where I can be useful on the Databricks workspace auth side. Will open the common.ai PR matching your shape once a release ships with DatabricksModel. Appreciate the guidance.

[aayushchou-db]

Hey @aman210122 ! currently that pr is blocked because databricks sdk doesn't have asyncio support just yet: databricks/databricks-sdk-py#246

Seems that issue has been triaged though, once it comes through I'll revisit the pr, hopefully that would unblock this!