This is an automated email from the ASF dual-hosted git repository. Lee-W pushed a commit to branch common-ai-supported-services-page in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 69164c3f242160230b3acaac48e1d8864b99c449 Author: Wei Lee <[email protected]> AuthorDate: Fri Sep 11 17:44:34 2026 +0800 Add a generated supported services page for the common.ai provider Readers had no single place to see what this provider reaches. The answer was spread across the toolsets guide, the index page, seven connection pages and provider.yaml, so the differentiating half — the Airflow hooks, MCP servers, SQL warehouses, DataFusion tables, sandboxes and vendor-managed agents a toolset gives you over a raw SDK call — was effectively unwritten. The model list had also drifted. It was a hand-copied snapshot of nine vendors while the connection actually resolves anything pydantic_ai.providers can construct, so services we already support, Snowflake Cortex and OpenRouter among them, were absent from the page a customer searches. Making provider.yaml the one source both the page and the registry read keeps them from disagreeing, and a test that derives the reachable set from pydantic-ai turns a drift into a red build naming the vendor rather than a quiet omission. That list is derived from a current pydantic-ai, while the provider supports pydantic-ai-slim from 2.23.0 upward and CI exercises exactly that floor. Snowflake Cortex and Crusoe reach back only to 2.27.0 and 2.28.0, so a vendor being absent downstream is a version difference rather than drift. The test therefore asserts on every supported version that an installed vendor is declared, and defers the opposite direction to the release that carries all of them. --- airflow-core/src/airflow/provider.yaml.schema.json | 24 ++ .../23_provider_hook_migration_to_yaml.rst | 27 ++ .../src/sphinx_exts/providers_extensions.py | 96 +++++- .../provider_connection_services.rst.jinja2 | 28 ++ .../templates/provider_toolset_services.rst.jinja2 | 28 ++ .../unit/sphinx_exts/test_providers_extensions.py | 115 +++++++ docs/spelling_wordlist.txt | 4 + providers/common/ai/docs/connections/langchain.rst | 10 +- providers/common/ai/docs/index.rst | 1 + providers/common/ai/docs/supported_services.rst | 53 +++ providers/common/ai/provider.yaml | 61 +++- .../providers/common/ai/get_provider_info.py | 29 +- .../tests/unit/common/ai/test_provider_metadata.py | 357 +++++++++++++++++++++ 13 files changed, 820 insertions(+), 13 deletions(-) diff --git a/airflow-core/src/airflow/provider.yaml.schema.json b/airflow-core/src/airflow/provider.yaml.schema.json index 4c4a199a882..a4215cf640c 100644 --- a/airflow-core/src/airflow/provider.yaml.schema.json +++ b/airflow-core/src/airflow/provider.yaml.schema.json @@ -388,6 +388,30 @@ "items": { "type": "string" } + }, + "external-services": { + "description": "Per-module list of external services each toolset reaches. Representative, non-exhaustive. Each 'module' must also appear in 'python-modules'.", + "type": "array", + "items": { + "type": "object", + "properties": { + "module": { + "type": "string", + "description": "Dotted path of the toolset module, as listed in 'python-modules'." + }, + "services": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "additionalProperties": false, + "required": [ + "module", + "services" + ] + } } }, "additionalProperties": false, diff --git a/contributing-docs/23_provider_hook_migration_to_yaml.rst b/contributing-docs/23_provider_hook_migration_to_yaml.rst index 3b0069dbeb4..b32aa2a5501 100644 --- a/contributing-docs/23_provider_hook_migration_to_yaml.rst +++ b/contributing-docs/23_provider_hook_migration_to_yaml.rst @@ -109,6 +109,33 @@ supported field options, see type: string default: "my-project" +Toolset metadata is defined under the separate top-level ``toolsets`` key: + +toolsets.external-services +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Per-module list of external services each toolset reaches, added as a sibling of +``python-modules`` inside a ``toolsets`` entry. Unlike the connection-type +``external-services`` key above, this one is not surfaced on the registry's +provider version page, and there is no built-in Sphinx rendering for it either -- +a provider that wants to show it must add its own docs page with the +``provider-toolset-services`` directive, the way ``common.ai`` does in +``providers/common/ai/docs/supported_services.rst`` (which also renders the +connection-type list there via ``provider-connection-services``). As with the +connection-type list, the entries are representative, not exhaustive. Every +``module`` value must also appear in that entry's ``python-modules`` list. + +.. code-block:: yaml + + toolsets: + - integration-name: Common AI + python-modules: + - airflow.providers.common.ai.toolsets.hook + external-services: + - module: airflow.providers.common.ai.toolsets.hook + services: + - Any Airflow connection, through its provider hook + Migration Tool -------------- diff --git a/devel-common/src/sphinx_exts/providers_extensions.py b/devel-common/src/sphinx_exts/providers_extensions.py index 5aa40ea9b59..ca45c16981d 100644 --- a/devel-common/src/sphinx_exts/providers_extensions.py +++ b/devel-common/src/sphinx_exts/providers_extensions.py @@ -25,13 +25,12 @@ from pathlib import Path from typing import Any # No stub exists for docutils.parsers.rst.directives. See https://github.com/python/typeshed/issues/5755. -from provider_yaml_utils import load_package_data - from sphinx_exts.operators_and_hooks_ref import ( DEFAULT_HEADER_SEPARATOR, BaseJinjaReferenceDirective, _render_template, ) +from sphinx_exts.provider_yaml_utils import load_package_data def find_class_methods_with_specific_calls( @@ -515,8 +514,101 @@ class OpenLineageSupportedClassesDirective(BaseJinjaReferenceDirective): return _render_openlineage_supported_classes_content() +def _find_provider_package_data(package_name: str) -> dict[str, Any]: + for provider in load_package_data(): + if provider["package-name"] == package_name: + return provider + raise ValueError(f"No provider.yaml found for package '{package_name}'") + + +def _render_connection_services_content(package_name: str) -> str: + provider = _find_provider_package_data(package_name) + rows = [ + { + "hook_name": conn["hook-name"], + "services": conn.get("external-services") or [], + "ref": f"howto/connection:{conn['connection-type']}", + } + for conn in provider.get("connection-types", []) + ] + return _render_template("provider_connection_services.rst.jinja2", rows=rows) + + +# Display name for each toolset module: the class name for modules that map to one +# (`toolsets.rst` documents each under a `` ``ClassName`` `` heading), or a short +# descriptive name for the two that don't (`managed_agent` documents a family of +# provider-specific subclasses under "Managed Agent Toolsets"; `langchain_bridge` is a +# function, documented under "Working with LangChain"). Kept here rather than in +# `provider.yaml` because it is presentation-only, not metadata the registry also needs. +# +# Deliberately no `.get(basename, default)` fallback below: every module in +# `python-modules` must have an entry here, the same way `LABELS` and the anchor check +# in `test_provider_metadata.py` require full coverage. A missing entry for a new +# toolset module fails loudly with a `ValueError` naming the provider and module, not +# a silently rendered basename. +_TOOLSET_DISPLAY_NAMES = { + "hook": "HookToolset", + "sql": "SQLToolset", + "datafusion": "DataFusionToolset", + "logging": "LoggingToolset", + "mcp": "MCPToolset", + "skills": "AgentSkillsToolset", + "sandbox": "SandboxToolset", + "langchain_bridge": "LangChain Bridge", + "managed_agent": "Managed Agent Toolsets", +} + + +def _render_toolset_services_content(package_name: str) -> str: + provider = _find_provider_package_data(package_name) + rows = [] + for toolset in provider.get("toolsets", []): + services_by_module = { + entry["module"]: entry["services"] for entry in toolset.get("external-services") or [] + } + for module in toolset.get("python-modules", []): + basename = module.rsplit(".", 1)[-1] + try: + display_name = _TOOLSET_DISPLAY_NAMES[basename] + except KeyError as e: + raise ValueError( + f"No display name for toolset module '{module}' of provider " + f"'{package_name}'. Add an entry for '{basename}' to `_TOOLSET_DISPLAY_NAMES`." + ) from e + rows.append( + { + "display_name": display_name, + "services": services_by_module.get(module, []), + "ref": f"howto/toolset:{basename}", + } + ) + return _render_template("provider_toolset_services.rst.jinja2", rows=rows) + + +class ProviderConnectionServicesDirective(BaseJinjaReferenceDirective): + """Render a table of a provider's connection types and the external services each reaches.""" + + required_arguments = 1 + optional_arguments = 0 + + def render_content(self, *, tags: set[str] | None, header_separator: str = DEFAULT_HEADER_SEPARATOR): + return _render_connection_services_content(self.arguments[0]) + + +class ProviderToolsetServicesDirective(BaseJinjaReferenceDirective): + """Render a table of a provider's toolset modules and the external services each reaches.""" + + required_arguments = 1 + optional_arguments = 0 + + def render_content(self, *, tags: set[str] | None, header_separator: str = DEFAULT_HEADER_SEPARATOR): + return _render_toolset_services_content(self.arguments[0]) + + def setup(app): """Setup plugin""" app.add_directive("airflow-providers-openlineage-supported-classes", OpenLineageSupportedClassesDirective) + app.add_directive("provider-connection-services", ProviderConnectionServicesDirective) + app.add_directive("provider-toolset-services", ProviderToolsetServicesDirective) return {"parallel_read_safe": True, "parallel_write_safe": True} diff --git a/devel-common/src/sphinx_exts/templates/provider_connection_services.rst.jinja2 b/devel-common/src/sphinx_exts/templates/provider_connection_services.rst.jinja2 new file mode 100644 index 00000000000..5d7bb00b66d --- /dev/null +++ b/devel-common/src/sphinx_exts/templates/provider_connection_services.rst.jinja2 @@ -0,0 +1,28 @@ +{# + Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. +#} +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Connection type + - External services + {% for row in rows -%} + * - :ref:`{{ row.hook_name }} <{{ row.ref }}>` + - {{ row.services | join(', ') }} + {% endfor %} diff --git a/devel-common/src/sphinx_exts/templates/provider_toolset_services.rst.jinja2 b/devel-common/src/sphinx_exts/templates/provider_toolset_services.rst.jinja2 new file mode 100644 index 00000000000..bc9059cc8cf --- /dev/null +++ b/devel-common/src/sphinx_exts/templates/provider_toolset_services.rst.jinja2 @@ -0,0 +1,28 @@ +{# + Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. +#} +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Toolset + - External services + {% for row in rows -%} + * - :ref:`{{ row.display_name }} <{{ row.ref }}>` + - {{ row.services | join(', ') }} + {% endfor %} diff --git a/devel-common/tests/unit/sphinx_exts/test_providers_extensions.py b/devel-common/tests/unit/sphinx_exts/test_providers_extensions.py new file mode 100644 index 00000000000..04ce64269e0 --- /dev/null +++ b/devel-common/tests/unit/sphinx_exts/test_providers_extensions.py @@ -0,0 +1,115 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest import mock + +import pytest + +from sphinx_exts.providers_extensions import ( + _render_connection_services_content, + _render_toolset_services_content, +) + +FAKE_PACKAGE_DATA = [ + { + "package-name": "apache-airflow-providers-example", + "connection-types": [ + { + "hook-class-name": "airflow.providers.example.hooks.example.ExampleHook", + "hook-name": "Example", + "connection-type": "example", + "external-services": ["OpenAI", "Anthropic"], + } + ], + "toolsets": [ + { + "integration-name": "Example", + "python-modules": [ + "airflow.providers.example.toolsets.hook", + "airflow.providers.example.toolsets.mcp", + ], + "external-services": [ + { + "module": "airflow.providers.example.toolsets.hook", + "services": ["Any Airflow connection, through its provider hook"], + }, + { + "module": "airflow.providers.example.toolsets.mcp", + "services": ["Any MCP server (user-supplied endpoint)"], + }, + ], + } + ], + } +] + + [email protected]("sphinx_exts.providers_extensions.load_package_data", autospec=True) +def test_render_connection_services_content_renders_one_row_per_connection_type(mock_load_package_data): + mock_load_package_data.return_value = FAKE_PACKAGE_DATA + + rendered = _render_connection_services_content("apache-airflow-providers-example") + + assert ":ref:`Example <howto/connection:example>`" in rendered + assert "OpenAI, Anthropic" in rendered + + [email protected]("sphinx_exts.providers_extensions.load_package_data", autospec=True) +def test_render_connection_services_content_unknown_package_raises(mock_load_package_data): + mock_load_package_data.return_value = FAKE_PACKAGE_DATA + + with pytest.raises(ValueError, match="No provider.yaml found"): + _render_connection_services_content("apache-airflow-providers-does-not-exist") + + [email protected]("sphinx_exts.providers_extensions.load_package_data", autospec=True) +def test_render_toolset_services_content_renders_one_row_per_module(mock_load_package_data): + mock_load_package_data.return_value = FAKE_PACKAGE_DATA + + rendered = _render_toolset_services_content("apache-airflow-providers-example") + + assert ":ref:`HookToolset <howto/toolset:hook>`" in rendered + assert "Any Airflow connection, through its provider hook" in rendered + assert ":ref:`MCPToolset <howto/toolset:mcp>`" in rendered + assert "Any MCP server (user-supplied endpoint)" in rendered + + +FAKE_PACKAGE_DATA_UNKNOWN_TOOLSET_MODULE = [ + { + "package-name": "apache-airflow-providers-example", + "toolsets": [ + { + "integration-name": "Example", + "python-modules": [ + "airflow.providers.example.toolsets.unknown_toolset", + ], + "external-services": [], + } + ], + } +] + + [email protected]("sphinx_exts.providers_extensions.load_package_data", autospec=True) +def test_render_toolset_services_content_unknown_module_raises(mock_load_package_data): + mock_load_package_data.return_value = FAKE_PACKAGE_DATA_UNKNOWN_TOOLSET_MODULE + + with pytest.raises(ValueError, match="apache-airflow-providers-example") as exc_info: + _render_toolset_services_content("apache-airflow-providers-example") + + assert "airflow.providers.example.toolsets.unknown_toolset" in str(exc_info.value) diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index c31da3d2400..1f6c03977c5 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -232,6 +232,7 @@ ccache cdc celltags centric +Cerebras cfg cgi cgitb @@ -1147,6 +1148,7 @@ NCUs nd ndjson nearText +Nebius neighbours neo Neo4j @@ -1914,6 +1916,7 @@ Vectorizers vendored venv venvs +Vercel Vertica vertica Vespa @@ -1965,6 +1968,7 @@ WTF wtf wtforms www +xAI XCom Xcom xcom diff --git a/providers/common/ai/docs/connections/langchain.rst b/providers/common/ai/docs/connections/langchain.rst index 1949d99dea8..656b991ddc4 100644 --- a/providers/common/ai/docs/connections/langchain.rst +++ b/providers/common/ai/docs/connections/langchain.rst @@ -66,11 +66,11 @@ Supported providers The hook forwards two values to LangChain: the connection's password as ``api_key`` and its host as ``base_url``. Any provider whose LangChain model -class accepts those two keyword arguments works, which includes OpenAI, -Anthropic, Groq, Mistral AI, DeepSeek, Ollama and vLLM. Providers whose -classes expect their own credential shape (AWS Bedrock, Google Vertex AI and -GenAI, Azure OpenAI, Cohere, HuggingFace) reject these kwargs and are not -usable through this connection type. +class accepts those two keyword arguments works. Providers whose classes +expect their own credential shape (AWS Bedrock, Google Vertex AI and GenAI, +Azure OpenAI, Cohere, HuggingFace) reject these kwargs and are not usable +through this connection type. See :doc:`../supported_services` for the +current list of providers this connection type reaches. Model resolution order ----------------------- diff --git a/providers/common/ai/docs/index.rst b/providers/common/ai/docs/index.rst index 0da4eec1bde..dffbe99f3e6 100644 --- a/providers/common/ai/docs/index.rst +++ b/providers/common/ai/docs/index.rst @@ -103,6 +103,7 @@ waits for the result, use that vendor's provider. What you can build <use_cases/index> Models and providers <model_providers> + Supported services <supported_services> Operators <operators/index> Toolsets <toolsets/index> LLM and agent features <features> diff --git a/providers/common/ai/docs/supported_services.rst b/providers/common/ai/docs/supported_services.rst new file mode 100644 index 00000000000..8e51f8f64a5 --- /dev/null +++ b/providers/common/ai/docs/supported_services.rst @@ -0,0 +1,53 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +Supported services +=================== + +``common.ai`` reaches external services on two axes: the model providers its connections +talk to, and the systems its toolsets talk to. For per-model detail beyond the vendor level +(exact model ids, pricing, capabilities), see `pydantic-ai's model list +<https://ai.pydantic.dev/models/>`__. + +Connections +----------- + +.. provider-connection-services:: apache-airflow-providers-common-ai + +Toolsets +-------- + +.. provider-toolset-services:: apache-airflow-providers-common-ai + +Notes +----- + +* The ``langchain`` connection type's row above is limited to OpenAI-compatible credential + surfaces (``api_key`` + optional ``base_url``); see :ref:`Supported providers + <langchain-supported-providers>` for the providers that reject those kwargs and are not + usable through it. +* Azure OpenAI, Google Vertex AI, and AWS Bedrock are also reachable through the generic + ``pydanticai`` connection type, but each has its own dedicated connection type + (:doc:`connections/pydantic_ai_azure`, :doc:`connections/pydantic_ai_vertex`, + :doc:`connections/pydantic_ai_bedrock`) for their non-standard auth. +* Most model providers need an extra installed alongside ``apache-airflow-providers-common-ai`` + — see the "Choosing extras" section of :doc:`index`. +* "Pydantic AI Gateway" in the ``pydanticai`` row is a routing layer, not an upstream vendor + in its own right: set the connection's Model field to ``gateway/<vendor>:<model>`` (for + example ``gateway/anthropic:claude-sonnet-5``) to send the request through it instead of + directly to the vendor. It currently routes to Anthropic, AWS Bedrock, Google Gemini, + Google Vertex AI, Groq, and OpenAI — all already listed above in their own right. diff --git a/providers/common/ai/provider.yaml b/providers/common/ai/provider.yaml index fb4a01d613d..e7490068335 100644 --- a/providers/common/ai/provider.yaml +++ b/providers/common/ai/provider.yaml @@ -158,15 +158,37 @@ connection-types: hook-name: "Pydantic AI" connection-type: pydanticai external-services: - - OpenAI + - Alibaba Cloud Model Studio - Anthropic - - Google - AWS Bedrock + - AWS Bedrock Mantle + - Azure OpenAI + - Cerebras + - Cohere + - Crusoe + - DeepSeek + - Fireworks AI + - Google Gemini + - Google Vertex AI - Groq + - Heroku + - Hugging Face + - LiteLLM - Mistral AI - - DeepSeek + - Moonshot AI + - Nebius - Ollama + - OpenAI + - OpenRouter + - OVHcloud + - Pydantic AI Gateway + - SambaNova + - Snowflake Cortex + - Together AI + - Vercel AI Gateway - vLLM + - xAI + - Z.AI ui-field-behaviour: hidden-fields: - schema @@ -411,6 +433,8 @@ connection-types: - hook-class-name: airflow.providers.common.ai.hooks.mcp.MCPHook hook-name: "MCP Server" connection-type: mcp + external-services: + - Any MCP server (user-supplied endpoint) ui-field-behaviour: hidden-fields: - schema @@ -552,6 +576,37 @@ toolsets: - airflow.providers.common.ai.toolsets.skills - airflow.providers.common.ai.toolsets.langchain_bridge - airflow.providers.common.ai.toolsets.managed_agent + external-services: + - module: airflow.providers.common.ai.toolsets.hook + services: + - Any Airflow connection, through its provider hook + - module: airflow.providers.common.ai.toolsets.sql + services: + - Any common.sql database connection (with table allowlists and AST validation) + - module: airflow.providers.common.ai.toolsets.datafusion + services: + - Object storage (S3, local filesystem, Apache Iceberg) via Apache DataFusion + - module: airflow.providers.common.ai.toolsets.logging + services: + - Airflow task logs + - module: airflow.providers.common.ai.toolsets.mcp + services: + - Any MCP server (user-supplied endpoint) + - module: airflow.providers.common.ai.toolsets.sandbox + services: + - Docker Sandboxes (the shipped sbx backend; other backends can be added via SandboxBackend) + - module: airflow.providers.common.ai.toolsets.skills + services: + - Agent skills from local paths and remote Git repositories + - module: airflow.providers.common.ai.toolsets.langchain_bridge + services: + - LangChain tools + - module: airflow.providers.common.ai.toolsets.managed_agent + services: + - Snowflake Cortex Agents + - Amazon Bedrock AgentCore + - Azure AI Foundry + - Vertex AI Agent Engine retry-policies: - integration-name: Common AI diff --git a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py index b70a7f8382b..dae5fac3a58 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py +++ b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py @@ -132,15 +132,37 @@ def get_provider_info(): "hook-name": "Pydantic AI", "connection-type": "pydanticai", "external-services": [ - "OpenAI", + "Alibaba Cloud Model Studio", "Anthropic", - "Google", "AWS Bedrock", + "AWS Bedrock Mantle", + "Azure OpenAI", + "Cerebras", + "Cohere", + "Crusoe", + "DeepSeek", + "Fireworks AI", + "Google Gemini", + "Google Vertex AI", "Groq", + "Heroku", + "Hugging Face", + "LiteLLM", "Mistral AI", - "DeepSeek", + "Moonshot AI", + "Nebius", "Ollama", + "OpenAI", + "OpenRouter", + "OVHcloud", + "Pydantic AI Gateway", + "SambaNova", + "Snowflake Cortex", + "Together AI", + "Vercel AI Gateway", "vLLM", + "xAI", + "Z.AI", ], "ui-field-behaviour": { "hidden-fields": ["schema", "port", "login"], @@ -323,6 +345,7 @@ def get_provider_info(): "hook-class-name": "airflow.providers.common.ai.hooks.mcp.MCPHook", "hook-name": "MCP Server", "connection-type": "mcp", + "external-services": ["Any MCP server (user-supplied endpoint)"], "ui-field-behaviour": { "hidden-fields": ["schema", "port", "login"], "relabeling": {"password": "Auth Token"}, diff --git a/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py b/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py new file mode 100644 index 00000000000..d684257cc6b --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py @@ -0,0 +1,357 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Drift tripwires between ``provider.yaml``'s declared ``external-services`` and what +pydantic-ai actually installs. + +How ``LABELS``, ``NON_MODULE_SERVICES``, ``UNREACHABLE_MODULES`` and +``DEPRECATED_UPSTREAM_MODULES`` below were derived, including the empirical checks that +back each exclusion, is documented in the comment directly above each constant. +""" + +from __future__ import annotations + +import importlib +import pkgutil +from importlib.metadata import version +from pathlib import Path + +import pydantic_ai.providers as pydantic_ai_providers +import pytest +import yaml +from packaging.version import Version +from pydantic_ai.exceptions import UserError +from pydantic_ai.models import infer_model, known_model_names + +from airflow.providers.common.ai.get_provider_info import get_provider_info + +# `toolsets` never flowed through the runtime `provider_info.schema.json` +# (memory `reference_two_provider_schemas_only_authoring_may_tighten`), so +# `get_provider_info()` doesn't carry it -- read `provider.yaml` directly instead. +PROVIDER_YAML_PATH = Path(__file__).resolve().parents[4] / "provider.yaml" + + +def _extract_toolset_block() -> dict: + provider_yaml = yaml.safe_load(PROVIDER_YAML_PATH.read_text()) + (toolset_block,) = provider_yaml["toolsets"] + return toolset_block + + +# Module basename -> display label used in `provider.yaml`'s `pydanticai.external-services`. +# Sourced from each module's own `name` property / class docstring; kept here (not in +# `src/`) because it is test data, not runtime behaviour. +LABELS = { + "alibaba": "Alibaba Cloud Model Studio", + "anthropic": "Anthropic", + "azure": "Azure OpenAI", + "bedrock": "AWS Bedrock", + "bedrock_mantle": "AWS Bedrock Mantle", + "cerebras": "Cerebras", + "cohere": "Cohere", + "crusoe": "Crusoe", + "deepseek": "DeepSeek", + "fireworks": "Fireworks AI", + "gateway": "Pydantic AI Gateway", + "google": "Google Gemini", + "google_cloud": "Google Vertex AI", + "groq": "Groq", + "heroku": "Heroku", + "huggingface": "Hugging Face", + "litellm": "LiteLLM", + "mistral": "Mistral AI", + "moonshotai": "Moonshot AI", + "nebius": "Nebius", + "ollama": "Ollama", + "openai": "OpenAI", + "openrouter": "OpenRouter", + "ovhcloud": "OVHcloud", + "sambanova": "SambaNova", + "snowflake": "Snowflake Cortex", + "together": "Together AI", + "vercel": "Vercel AI Gateway", + "xai": "xAI", + "zai": "Z.AI", +} + +# Declared in `pydanticai.external-services` but not a pydantic-ai provider module: reached +# by pointing the `openai` provider's `base_url` at a vLLM endpoint instead +# (`hooks/pydantic_ai.py`), so it cannot be derived from `pydantic_ai.providers`. +NON_MODULE_SERVICES = {"vLLM"} + +# Modules `pkgutil` discovers under `pydantic_ai.providers` that `infer_model` never +# routes to -- `infer_model("voyageai:...")` / `infer_model("sentence-transformers:...")` +# raise `UserError: Unknown model`, confirmed by `test_unreachable_modules_stay_unreachable` +# below. Not listed in `LABELS`/`provider.yaml` because there is no `common.ai` code path +# (chat or embedding) that reaches them. +UNREACHABLE_MODULES = {"voyageai", "sentence_transformers"} + +# Modules that ARE reachable through `infer_model` (unlike UNREACHABLE_MODULES) but are +# excluded by decision because pydantic-ai itself marks the upstream backend dead: +# `pydantic_ai.providers.github.GitHubProvider` is decorated `@deprecated` with the message +# "GitHub Models was retired on 2026-07-30 ... this provider is deprecated and will be +# removed in v3." There is no GitHub Copilot provider anywhere in pydantic-ai; the only +# "github" support that ever existed is this now-retired integration. Advertising a vendor +# whose backend the vendor itself shut down is worse than not listing it (Wei's call, +# 2026-09-11 review round). `test_deprecated_upstream_modules_stay_deprecated` below +# guards the reversal: if pydantic-ai un-deprecates this module, or a real GitHub Copilot +# provider appears under a different module name, the corresponding test goes red instead +# of silently staying excluded. +DEPRECATED_UPSTREAM_MODULES = {"github"} + +# `provider.yaml`'s vendor list is derived from a current pydantic-ai, but the provider +# supports `pydantic-ai-slim>=2.33.0` (`pyproject.toml`) and CI's "Low dep tests" job installs +# exactly that floor. Vendors upstream added after the floor are legitimately absent there -- +# bisected empirically: `vllm` landed in 2.38.0 and `github_copilot` in 2.42.0 (`openai_codex` +# also landed after the floor, in 2.41.0, but is excluded from this reasoning -- see +# `OPENAI_CODEX_AVAILABLE_FROM` above) -- so only the "installed but not declared" direction +# holds on every supported version. The opposite direction ("declared but not installed") is a +# real drift signal solely from the release that carries every vendor currently declared, i.e. +# the later of the two, `github_copilot`'s 2.42.0: raise this when a newly declared vendor lands +# in a later release. +ALL_DECLARED_VENDORS_AVAILABLE_FROM = Version("2.42.0") + +requires_all_declared_vendors = pytest.mark.skipif( + Version(Version(version("pydantic-ai-slim")).base_version) < ALL_DECLARED_VENDORS_AVAILABLE_FROM, + reason=( + f"pydantic-ai-slim is older than {ALL_DECLARED_VENDORS_AVAILABLE_FROM}, the first release " + f"carrying every vendor provider.yaml declares" + ), +) + +# `known_model_names()` prefixes that don't match a `pydantic_ai.providers` module basename +# directly (hyphen normalized to underscore) but still resolve through `infer_model`. +PREFIX_ALIASES = { + "openai-chat": "openai", + "openai-responses": "openai", + "azure-responses": "azure", +} + +# The vendors `supported_services.rst`'s Notes section names, in hand-written prose, as the +# ones `gateway/<vendor>:<model>` currently routes to. Deliberately kept as a human-readable +# sentence there rather than a generated list (2026-09-11 review round: this is a customer- +# facing page, not a place to trade readability for testability) -- this constant is the +# reference value the two tests below check it against, in both directions: +# `test_gateway_note_vendors_match_known_model_names` guards against pydantic-ai adding/dropping +# a gateway vendor, and `test_gateway_note_sentence_names_the_expected_vendors` guards against +# the prose and this constant drifting apart from each other. +GATEWAY_NOTE_VENDOR_LABELS = { + "Anthropic", + "AWS Bedrock", + "Google Gemini", + "Google Vertex AI", + "Groq", + "OpenAI", +} + + +def _find_discovered_provider_modules() -> set[str]: + """All non-private modules under `pydantic_ai.providers`, including `gateway`. + + `gateway` is a real, reachable upstream service (`PydanticAIHook.get_conn()` forwards + a `gateway/<vendor>:<model>` model string straight to `infer_model`, which resolves it + to a working `Model` instance), so it is not excluded here. + """ + return { + m.name for m in pkgutil.iter_modules(pydantic_ai_providers.__path__) if not m.name.startswith("_") + } + + +def _find_reachable_provider_modules() -> set[str]: + """The discovered modules `common.ai` actually reaches, i.e. minus the documented exclusions.""" + return _find_discovered_provider_modules() - UNREACHABLE_MODULES - DEPRECATED_UPSTREAM_MODULES + + +def _find_derived_vendor_labels() -> set[str]: + """The vendor labels the *installed* pydantic-ai implies `provider.yaml` should declare.""" + return { + LABELS.get(m, f"<unlabelled:{m}>") for m in _find_reachable_provider_modules() + } | NON_MODULE_SERVICES + + +def _find_pydanticai_external_services() -> list[str]: + connection_types = get_provider_info()["connection-types"] + (pydanticai,) = (c for c in connection_types if c["connection-type"] == "pydanticai") + return pydanticai["external-services"] + + +class TestConnectionTypeExternalServices: + def test_every_connection_type_declares_external_services(self): + for conn in get_provider_info()["connection-types"]: + assert conn.get("external-services"), ( + f"connection-type {conn['connection-type']!r} has no external-services" + ) + + +class TestPydanticAIExternalServicesDrift: + def test_every_installed_vendor_is_declared(self): + """ + The direction that holds on every supported pydantic-ai: a vendor the installed release + can reach must be declared, or the page silently under-advertises what the provider does. + """ + missing = _find_derived_vendor_labels() - set(_find_pydanticai_external_services()) + assert not missing, f"pydantic-ai vendors not declared in provider.yaml: {sorted(missing)}" + + @requires_all_declared_vendors + def test_no_declared_vendor_is_absent_upstream(self): + """ + The opposite direction, gated on ``ALL_DECLARED_VENDORS_AVAILABLE_FROM``: once every + declared vendor exists upstream, one that pydantic-ai does not have is either a typo or a + vendor whose module was dropped, and the page would overclaim what the provider reaches. + """ + extra = set(_find_pydanticai_external_services()) - _find_derived_vendor_labels() + assert not extra, f"provider.yaml declares vendors pydantic-ai doesn't have: {sorted(extra)}" + + def test_every_installed_module_has_a_label(self): + """As `test_every_installed_vendor_is_declared`, for `LABELS` rather than `provider.yaml`.""" + unlabelled = _find_reachable_provider_modules() - set(LABELS) + assert not unlabelled, f"no label for pydantic-ai module(s): {sorted(unlabelled)}" + + @requires_all_declared_vendors + def test_no_label_is_stale(self): + """As `test_no_declared_vendor_is_absent_upstream`, for `LABELS`.""" + stale = set(LABELS) - _find_reachable_provider_modules() + assert not stale, f"label(s) for module(s) pydantic-ai no longer has: {sorted(stale)}" + + def test_unreachable_modules_stay_unreachable(self): + """ + Guard the other direction: if a future pydantic-ai release makes one of + ``UNREACHABLE_MODULES`` reachable through ``infer_model``, this must fail so a human + adds it to ``LABELS`` and ``provider.yaml`` instead of it silently staying missing. + """ + for module in UNREACHABLE_MODULES: + prefix = module.replace("_", "-") + with pytest.raises(UserError, match="Unknown model"): + infer_model(f"{prefix}:placeholder", provider_factory=lambda _p: object()) + + def test_deprecated_upstream_modules_stay_deprecated(self): + """ + Guard the other direction, for every member of ``DEPRECATED_UPSTREAM_MODULES`` (not just + ``github`` -- this set is the one escape hatch in this file that isn't validated against + pydantic-ai by name elsewhere, so adding a module here that isn't genuinely + upstream-deprecated must not pass silently). If pydantic-ai un-deprecates a member + (reactivating its backend, or repurposing the module), this must fail so a human + re-evaluates whether it belongs back in ``LABELS`` and ``provider.yaml``, instead of it + staying excluded forever on a stale rationale. + """ + for module in DEPRECATED_UPSTREAM_MODULES: + try: + mod = importlib.import_module(f"pydantic_ai.providers.{module}") + except ImportError as exc: + # Deliberately not `pytest.importorskip`: an unmet optional dependency here + # would silently turn off this guard, not just this one test run. + pytest.fail( + f"pydantic_ai.providers.{module} (in DEPRECATED_UPSTREAM_MODULES) failed to " + f"import: {exc}. Install the missing dependency rather than skipping this " + f"test." + ) + provider_classes = [ + obj + for obj in vars(mod).values() + if isinstance(obj, type) + and obj.__module__ == mod.__name__ + and obj.__name__.endswith("Provider") + ] + assert provider_classes, f"no Provider class found in pydantic_ai.providers.{module}" + assert all(hasattr(cls, "__deprecated__") for cls in provider_classes), ( + f"pydantic_ai.providers.{module} is no longer marked @deprecated -- re-evaluate " + f"whether '{module}' should be re-added to LABELS and provider.yaml" + ) + + def test_known_model_name_prefixes_are_all_labelled(self): + labelled = set(LABELS) + for name in known_model_names(): + if ":" not in name: + continue # the bare "test" entry, pydantic-ai's built-in keyless TestModel + prefix = name.removeprefix("gateway/").split(":", 1)[0] + module = PREFIX_ALIASES.get(prefix, prefix.replace("-", "_")) + assert module in labelled, f"model prefix {prefix!r} (from {name!r}) has no label" + + def test_gateway_note_vendors_match_known_model_names(self): + """ + Guard `supported_services.rst`'s Notes sentence naming which vendors + ``gateway/<vendor>:<model>`` currently routes to. If pydantic-ai's gateway adds or drops + a vendor, this must fail so a human updates that sentence (and + ``GATEWAY_NOTE_VENDOR_LABELS``) instead of it silently going stale. + """ + gateway_prefixes = { + name.split("/", 1)[1].split(":", 1)[0] + for name in known_model_names() + if name.startswith("gateway/") + } + modules = {PREFIX_ALIASES.get(p, p.replace("-", "_")) for p in gateway_prefixes} + derived = {LABELS.get(m, f"<unlabelled:{m}>") for m in modules} + + assert derived == GATEWAY_NOTE_VENDOR_LABELS, ( + f"gateway now routes to {sorted(derived)} -- update the Notes sentence in " + f"supported_services.rst and GATEWAY_NOTE_VENDOR_LABELS to match" + ) + + def test_gateway_note_sentence_names_the_expected_vendors(self): + """ + Guard the prose itself, in both directions: every vendor in + ``GATEWAY_NOTE_VENDOR_LABELS`` must be named in `supported_services.rst`'s gateway Notes + bullet specifically -- not just somewhere in the file, where several of these vendor + names already appear in other sentences -- and no *other* labelled vendor may be named + in that bullet either. Editing the bullet without updating the constant above (or vice + versa) must be caught either way: underclaiming a vendor is a stale doc, overclaiming one + is a false statement about what this service reaches. + """ + rst_text = (PROVIDER_YAML_PATH.parent / "docs" / "supported_services.rst").read_text() + marker = '* "Pydantic AI Gateway" in the' + start = rst_text.index(marker) + end = rst_text.find("\n* ", start) + gateway_bullet = rst_text[start:] if end == -1 else rst_text[start:end] + gateway_bullet = " ".join(gateway_bullet.split()) + + for vendor in GATEWAY_NOTE_VENDOR_LABELS: + assert vendor in gateway_bullet, ( + f"{vendor!r} missing from supported_services.rst's gateway Notes bullet" + ) + + # `LABELS["gateway"]` ("Pydantic AI Gateway") is excluded: it is itself a value of + # `LABELS` and rightfully appears in this bullet (it's the row's own name), so it must + # not be flagged as an "other" vendor the bullet shouldn't mention. + other_labels = set(LABELS.values()) - GATEWAY_NOTE_VENDOR_LABELS - {LABELS["gateway"]} + unexpected = sorted(vendor for vendor in other_labels if vendor in gateway_bullet) + assert not unexpected, ( + f"gateway Notes bullet names {unexpected}, which are not in GATEWAY_NOTE_VENDOR_LABELS" + ) + + +class TestToolsetExternalServices: + def test_every_toolset_module_has_external_services(self): + toolset_block = _extract_toolset_block() + modules = set(toolset_block["python-modules"]) + services_by_module = { + entry["module"]: entry["services"] for entry in toolset_block["external-services"] + } + + assert set(services_by_module) == modules, ( + f"toolsets.external-services modules {sorted(services_by_module)} != " + f"python-modules {sorted(modules)}" + ) + for module, services in services_by_module.items(): + assert services, f"toolset module {module!r} has an empty services list" + + def test_every_toolset_module_has_a_docs_anchor(self): + toolsets_rst = (PROVIDER_YAML_PATH.parent / "docs" / "toolsets.rst").read_text() + toolset_block = _extract_toolset_block() + + for module in toolset_block["python-modules"]: + basename = module.rsplit(".", 1)[-1] + anchor = f".. _howto/toolset:{basename}:" + assert anchor in toolsets_rst, f"missing docs anchor {anchor!r} for module {module!r}"
