GitHub user aman210122 closed a discussion: Proposal: Databricks Unity AI Gateway model class for apache-airflow-providers-common-ai
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://<workspace-host>/serving-endpoints (the existing model serving surface, formerly fronted by Mosaic AI Gateway features). Beta: https://<workspace-host>/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 https://github.com/pydantic/pydantic-ai/issues/2947 and PR https://github.com/pydantic/pydantic-ai/pull/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 https://github.com/orgs/apache/projects/586) 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: <link-to-fork-branch-once-pushed>. 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. GitHub link: https://github.com/apache/airflow/discussions/67581 ---- This is an automatically sent email for [email protected]. To unsubscribe, please send an email to: [email protected]
