Skip to content

Commit 194e0ea

Browse files
authored
Add a generated supported services page for the common.ai provider (#72939)
1 parent fac3157 commit 194e0ea

12 files changed

Lines changed: 749 additions & 14 deletions

File tree

‎devel-common/src/sphinx_exts/providers_extensions.py‎

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,12 @@
2525
from typing import Any
2626

2727
# No stub exists for docutils.parsers.rst.directives. See https://github.com/python/typeshed/issues/5755.
28-
from provider_yaml_utils import load_package_data
29-
3028
from sphinx_exts.operators_and_hooks_ref import (
3129
DEFAULT_HEADER_SEPARATOR,
3230
BaseJinjaReferenceDirective,
3331
_render_template,
3432
)
33+
from sphinx_exts.provider_yaml_utils import load_package_data
3534

3635

3736
def find_class_methods_with_specific_calls(
@@ -515,8 +514,39 @@ def render_content(self, *, tags: set[str] | None, header_separator: str = DEFAU
515514
return _render_openlineage_supported_classes_content()
516515

517516

517+
def _find_provider_package_data(package_name: str) -> dict[str, Any]:
518+
for provider in load_package_data():
519+
if provider["package-name"] == package_name:
520+
return provider
521+
raise ValueError(f"No provider.yaml found for package '{package_name}'")
522+
523+
524+
def _render_connection_services_content(package_name: str) -> str:
525+
provider = _find_provider_package_data(package_name)
526+
rows = [
527+
{
528+
"hook_name": conn["hook-name"],
529+
"services": conn.get("external-services") or [],
530+
"ref": f"howto/connection:{conn['connection-type']}",
531+
}
532+
for conn in provider.get("connection-types", [])
533+
]
534+
return _render_template("provider_connection_services.rst.jinja2", rows=rows)
535+
536+
537+
class ProviderConnectionServicesDirective(BaseJinjaReferenceDirective):
538+
"""Render a table of a provider's connection types and the external services each reaches."""
539+
540+
required_arguments = 1
541+
optional_arguments = 0
542+
543+
def render_content(self, *, tags: set[str] | None, header_separator: str = DEFAULT_HEADER_SEPARATOR):
544+
return _render_connection_services_content(self.arguments[0])
545+
546+
518547
def setup(app):
519548
"""Setup plugin"""
520549
app.add_directive("airflow-providers-openlineage-supported-classes", OpenLineageSupportedClassesDirective)
550+
app.add_directive("provider-connection-services", ProviderConnectionServicesDirective)
521551

522552
return {"parallel_read_safe": True, "parallel_write_safe": True}
Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
{#
2+
Licensed to the Apache Software Foundation (ASF) under one
3+
or more contributor license agreements. See the NOTICE file
4+
distributed with this work for additional information
5+
regarding copyright ownership. The ASF licenses this file
6+
to you under the Apache License, Version 2.0 (the
7+
"License"); you may not use this file except in compliance
8+
with the License. You may obtain a copy of the License at
9+
10+
http://www.apache.org/licenses/LICENSE-2.0
11+
12+
Unless required by applicable law or agreed to in writing,
13+
software distributed under the License is distributed on an
14+
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
KIND, either express or implied. See the License for the
16+
specific language governing permissions and limitations
17+
under the License.
18+
#}
19+
.. list-table::
20+
:header-rows: 1
21+
:widths: 30 70
22+
23+
* - Connection type
24+
- External services
25+
{% for row in rows -%}
26+
* - :ref:`{{ row.hook_name }} <{{ row.ref }}>`
27+
- {{ row.services | join(', ') }}
28+
{% endfor %}
Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
from __future__ import annotations
18+
19+
from unittest import mock
20+
21+
import pytest
22+
23+
from sphinx_exts.providers_extensions import (
24+
_render_connection_services_content,
25+
)
26+
27+
FAKE_PACKAGE_DATA = [
28+
{
29+
"package-name": "apache-airflow-providers-example",
30+
"connection-types": [
31+
{
32+
"hook-class-name": "airflow.providers.example.hooks.example.ExampleHook",
33+
"hook-name": "Example",
34+
"connection-type": "example",
35+
"external-services": ["OpenAI", "Anthropic"],
36+
}
37+
],
38+
}
39+
]
40+
41+
42+
@mock.patch("sphinx_exts.providers_extensions.load_package_data", autospec=True)
43+
def test_render_connection_services_content_renders_one_row_per_connection_type(mock_load_package_data):
44+
mock_load_package_data.return_value = FAKE_PACKAGE_DATA
45+
46+
rendered = _render_connection_services_content("apache-airflow-providers-example")
47+
48+
assert ":ref:`Example <howto/connection:example>`" in rendered
49+
assert "OpenAI, Anthropic" in rendered
50+
51+
52+
@mock.patch("sphinx_exts.providers_extensions.load_package_data", autospec=True)
53+
def test_render_connection_services_content_unknown_package_raises(mock_load_package_data):
54+
mock_load_package_data.return_value = FAKE_PACKAGE_DATA
55+
56+
with pytest.raises(ValueError, match="No provider.yaml found"):
57+
_render_connection_services_content("apache-airflow-providers-does-not-exist")

‎docs/spelling_wordlist.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,7 @@ ccache
233233
cdc
234234
celltags
235235
centric
236+
Cerebras
236237
cfg
237238
cgi
238239
cgitb
@@ -327,6 +328,7 @@ contextmgr
327328
contrib
328329
convertToIri
329330
cooldown
331+
Copilot
330332
copyable
331333
coreutils
332334
CoreV
@@ -1157,6 +1159,7 @@ NCUs
11571159
nd
11581160
ndjson
11591161
nearText
1162+
Nebius
11601163
neighbours
11611164
neo
11621165
Neo4j
@@ -1929,6 +1932,7 @@ Vectorizers
19291932
vendored
19301933
venv
19311934
venvs
1935+
Vercel
19321936
Vertica
19331937
vertica
19341938
Vespa
@@ -1981,6 +1985,7 @@ WTF
19811985
wtf
19821986
wtforms
19831987
www
1988+
xAI
19841989
XCom
19851990
Xcom
19861991
xcom

‎providers/common/ai/docs/connections/langchain.rst‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -71,11 +71,11 @@ Supported providers
7171

7272
The hook forwards two values to LangChain: the connection's password as
7373
``api_key`` and its host as ``base_url``. Any provider whose LangChain model
74-
class accepts those two keyword arguments works, which includes OpenAI,
75-
Anthropic, Groq, Mistral AI, DeepSeek, Ollama and vLLM. Providers whose
76-
classes expect their own credential shape (AWS Bedrock, Google Vertex AI and
77-
GenAI, Azure OpenAI, Cohere, HuggingFace) reject these kwargs and are not
78-
usable through this connection type.
74+
class accepts those two keyword arguments works. Providers whose classes
75+
expect their own credential shape (AWS Bedrock, Google Vertex AI and GenAI,
76+
Azure OpenAI, Cohere, HuggingFace) reject these kwargs and are not usable
77+
through this connection type. See :doc:`../supported_services` for the
78+
current list of providers this connection type reaches.
7979

8080
Model resolution order
8181
-----------------------

‎providers/common/ai/docs/index.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ waits for the result, use that vendor's provider.
105105
What you can build <use_cases/index>
106106
Connections <connections/index>
107107
Models and providers <model_providers>
108+
Supported services <supported_services>
108109
Operators and decorators <operators/index>
109110
Toolsets <toolsets/index>
110111
Agent frameworks <frameworks/index>

‎providers/common/ai/docs/model_providers.rst‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@ name picks the vendor. Switching vendors is a connection
2626
change, not a Dag change. Find your vendor below, install the extra, create the connection
2727
type shown, and set the model name with that prefix.
2828

29+
For the exhaustive, generated-and-hand-written list of every service this provider
30+
reaches (connections and toolsets alike), see :doc:`supported_services`.
31+
2932
.. list-table::
3033
:header-rows: 1
3134
:widths: 18 16 18 20 28
Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
.. Licensed to the Apache Software Foundation (ASF) under one
2+
or more contributor license agreements. See the NOTICE file
3+
distributed with this work for additional information
4+
regarding copyright ownership. The ASF licenses this file
5+
to you under the Apache License, Version 2.0 (the
6+
"License"); you may not use this file except in compliance
7+
with the License. You may obtain a copy of the License at
8+
9+
.. http://www.apache.org/licenses/LICENSE-2.0
10+
11+
.. Unless required by applicable law or agreed to in writing,
12+
software distributed under the License is distributed on an
13+
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
KIND, either express or implied. See the License for the
15+
specific language governing permissions and limitations
16+
under the License.
17+
18+
Supported services
19+
===================
20+
21+
``common.ai`` reaches external services on two axes: the model providers its connections
22+
talk to, and the systems its toolsets talk to. For per-model detail beyond the vendor level
23+
(exact model ids, pricing, capabilities), see `pydantic-ai's model list
24+
<https://ai.pydantic.dev/models/>`__.
25+
26+
This page is the complete list of what this provider reaches; for a curated, install- and
27+
credential-oriented walkthrough of each model vendor, see :doc:`model_providers`.
28+
29+
Connections
30+
-----------
31+
32+
.. provider-connection-services:: apache-airflow-providers-common-ai
33+
34+
Toolsets
35+
--------
36+
37+
.. list-table::
38+
:header-rows: 1
39+
:widths: 30 70
40+
41+
* - Toolset
42+
- External services
43+
* - :doc:`HookToolset <toolsets/hook>`
44+
- Any Airflow connection, through its provider hook
45+
* - :doc:`SQLToolset <toolsets/sql>`
46+
- Any common.sql database connection (with table allowlists and AST validation)
47+
* - :doc:`DataFusionToolset <toolsets/datafusion>`
48+
- Object storage (S3, local filesystem, Apache Iceberg) via Apache DataFusion
49+
* - :doc:`LoggingToolset <toolsets/logging>`
50+
- Airflow task logs
51+
* - :doc:`MCPToolset <toolsets/mcp>`
52+
- Any MCP server (user-supplied endpoint)
53+
* - :doc:`ObjectStorageToolset <toolsets/object_storage>`
54+
- Object storage (S3, GCS, Azure Blob Storage, or any store
55+
:class:`~airflow.sdk.ObjectStoragePath` can open), read-only
56+
* - :doc:`SandboxToolset <sandbox/index>`
57+
- Docker Sandboxes (the shipped sbx backend; other backends can be added via SandboxBackend)
58+
* - :doc:`AgentSkillsToolset <toolsets/skills>`
59+
- Agent skills from local paths and remote Git repositories
60+
* - :doc:`LangChain Bridge <toolsets/langchain>`
61+
- LangChain tools
62+
* - :doc:`Managed Agent Toolsets <toolsets/managed_agent>`
63+
- Snowflake Cortex Agents; Amazon Bedrock AgentCore; Azure AI Foundry; Vertex AI Agent Engine
64+
65+
Notes
66+
-----
67+
68+
* The ``langchain`` connection type's row above is limited to OpenAI-compatible credential
69+
surfaces (``api_key`` + optional ``base_url``); see :ref:`Supported providers
70+
<langchain-supported-providers>` for the providers that reject those kwargs and are not
71+
usable through it.
72+
* Azure OpenAI, Google Vertex AI, and AWS Bedrock are also reachable through the generic
73+
``pydanticai`` connection type, but each has its own dedicated connection type
74+
(:doc:`connections/pydantic_ai_azure`, :doc:`connections/pydantic_ai_vertex`,
75+
:doc:`connections/pydantic_ai_bedrock`) for their non-standard auth.
76+
* Most model providers need an extra installed alongside ``apache-airflow-providers-common-ai``
77+
— see the "Choosing extras" section of :doc:`installation`.
78+
* "Pydantic AI Gateway" in the ``pydanticai`` row is a routing layer, not an upstream vendor
79+
in its own right: set the connection's Model field to ``gateway/<vendor>:<model>`` (for
80+
example ``gateway/anthropic:claude-sonnet-5``) to send the request through it instead of
81+
directly to the vendor. It currently routes to Anthropic, AWS Bedrock, Google Gemini,
82+
Google Vertex AI, Groq, and OpenAI — all already listed above in their own right.

‎providers/common/ai/provider.yaml‎

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -162,15 +162,39 @@ connection-types:
162162
hook-name: "Pydantic AI"
163163
connection-type: pydanticai
164164
external-services:
165-
- OpenAI
165+
- Alibaba Cloud Model Studio
166166
- Anthropic
167-
- Google
168167
- AWS Bedrock
168+
- AWS Bedrock Mantle
169+
- Azure OpenAI
170+
- Cerebras
171+
- Cohere
172+
- Crusoe
173+
- DeepSeek
174+
- Fireworks AI
175+
- GitHub Copilot
176+
- Google Gemini
177+
- Google Vertex AI
169178
- Groq
179+
- Heroku
180+
- Hugging Face
181+
- LiteLLM
170182
- Mistral AI
171-
- DeepSeek
183+
- Moonshot AI
184+
- Nebius
172185
- Ollama
186+
- OpenAI
187+
- OpenRouter
188+
- OVHcloud
189+
- Pydantic AI Gateway
190+
- SambaNova
191+
- Snowflake Cortex
192+
- Together AI
193+
- TypeSafe
194+
- Vercel AI Gateway
173195
- vLLM
196+
- xAI
197+
- Z.AI
174198
ui-field-behaviour:
175199
hidden-fields:
176200
- schema
@@ -415,6 +439,8 @@ connection-types:
415439
- hook-class-name: airflow.providers.common.ai.hooks.mcp.MCPHook
416440
hook-name: "MCP Server"
417441
connection-type: mcp
442+
external-services:
443+
- Any MCP server (user-supplied endpoint)
418444
ui-field-behaviour:
419445
hidden-fields:
420446
- schema

0 commit comments

Comments
 (0)