This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new d8efc66ea76 Add a Model tab to the task-instance page for `common.ai`
(#74334)
d8efc66ea76 is described below
commit d8efc66ea76efac8a70725bc9eedea2e4001a093
Author: Pierre Jeambrun <[email protected]>
AuthorDate: Wed Oct 14 01:44:33 2026 +0200
Add a Model tab to the task-instance page for `common.ai` (#74334)
---
.../core_api/routes/public/test_plugins.py | 11 +-
.../tests/unit/plugins/test_plugins_manager.py | 2 +-
providers/common/ai/docs/observability.rst | 15 +++
providers/common/ai/provider.yaml | 2 +
providers/common/ai/pyproject.toml | 1 +
.../providers/common/ai/get_provider_info.py | 6 +-
.../providers/common/ai/plugins/_plugin_urls.py | 50 ++++++++
.../providers/common/ai/plugins/hitl_review.py | 36 +-----
.../providers/common/ai/plugins/model_panel.py | 97 ++++++++++++++
.../providers/common/ai/plugins/www/package.json | 3 +-
.../providers/common/ai/plugins/www/pnpm-lock.yaml | 12 ++
.../ai/plugins/www/src/components/ModelPanel.tsx | 142 +++++++++++++++++++++
.../ai/plugins/www/src/components/NoModelInfo.tsx | 58 +++++++++
.../ai/plugins/www/src/components/Tooltip.tsx | 65 ++++++++++
.../ai/plugins/www/src/hooks/useModelInfo.ts | 67 ++++++++++
.../common/ai/plugins/www/src/model-api.ts | 83 ++++++++++++
.../providers/common/ai/plugins/www/src/model.tsx | 73 +++++++++++
.../common/ai/plugins/www/src/types/model.ts | 28 ++++
.../providers/common/ai/plugins/www/vite.config.ts | 23 +++-
.../unit/common/ai/plugins/test__plugin_urls.py | 46 +++++++
.../unit/common/ai/plugins/test_hitl_review.py | 11 --
.../unit/common/ai/plugins/test_model_panel.py | 72 +++++++++++
.../common/ai/plugins/test_model_panel_compat.py | 46 +++++++
scripts/ci/prek/update_providers_dependencies.py | 8 +-
24 files changed, 898 insertions(+), 59 deletions(-)
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py
index ac58f76e334..6775fa3e6ba 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py
@@ -39,11 +39,12 @@ class TestGetPlugins:
# Filters
(
{},
- 19,
+ 20,
[
"InformaticaProviderPlugin",
"MetadataCollectionPlugin",
"OpenLineageProviderPlugin",
+ "ai_model_panel",
"business_day_window_plugin",
"databricks_workflow",
"decreasing_priority_weight_strategy_plugin",
@@ -64,14 +65,14 @@ class TestGetPlugins:
),
(
{"limit": 3, "offset": 3},
- 19,
+ 20,
[
+ "ai_model_panel",
"business_day_window_plugin",
"databricks_workflow",
- "decreasing_priority_weight_strategy_plugin",
],
),
- ({"limit": 1}, 19, ["InformaticaProviderPlugin"]),
+ ({"limit": 1}, 20, ["InformaticaProviderPlugin"]),
],
)
def test_should_respond_200(
@@ -220,7 +221,7 @@ class TestGetPlugins:
assert len(plugins_page) == 8
assert "test_plugin_invalid" not in [p["name"] for p in plugins_page]
- assert body["total_entries"] == 19
+ assert body["total_entries"] == 20
@skip_if_force_lowest_dependencies_marker
diff --git a/airflow-core/tests/unit/plugins/test_plugins_manager.py
b/airflow-core/tests/unit/plugins/test_plugins_manager.py
index e994baa923b..fc52a901849 100644
--- a/airflow-core/tests/unit/plugins/test_plugins_manager.py
+++ b/airflow-core/tests/unit/plugins/test_plugins_manager.py
@@ -978,7 +978,7 @@ class TestPluginsManager:
# Mock/skip loading from plugin dir
with
mock.patch("airflow.plugins_manager._load_plugins_from_plugin_directory",
return_value=([], [])):
plugins = plugins_manager._get_plugins()[0]
- assert len(plugins) == 7
+ assert len(plugins) == 8
class TestWindowPluginRegistration:
diff --git a/providers/common/ai/docs/observability.rst
b/providers/common/ai/docs/observability.rst
index 7e66f1f8197..e4a6fa02060 100644
--- a/providers/common/ai/docs/observability.rst
+++ b/providers/common/ai/docs/observability.rst
@@ -103,6 +103,21 @@ How it works
``@task.llm`` additionally pushes the model name above. The other LLM
operators still emit GenAI spans correlated to the task span by nesting, but
without the identity attributes, the run join key, or the model name.
+* **Model tab.** Separately from tracing, the provider adds an "AI Model" tab
to the
+ task-instance page showing the resolved model name and the ``usage`` XCom
above as a
+ readable table, including an estimated cost. It needs Airflow >= 3.4, since
it relies on
+ the ``applies_to`` operator scoping core added in that version to show up
only for
+ ``LLMOperator`` / ``AgentOperator`` runs (and their ``@task.llm`` /
``@task.agent``
+ decorator forms) -- the match is on the exact operator name, so a subclass
of either does
+ not get the tab automatically. Like the span-level cost above, the tab's
number is a
+ best-effort ``genai-prices`` estimate, not provider-reported billing. With a
custom XCom
+ backend that offloads large values to external storage (e.g. ``common.io``
with
+ ``xcom_objectstorage_threshold`` set low), the tab reads the stored
reference rather than
+ the real value, since the public XCom endpoint it calls never runs the
backend's
+ ``deserialize_value``: the model name badge would show the reference string
(e.g. an
+ ``s3://...`` path) instead of the model, and ``usage`` would come back as
that same
+ string rather than an object -- the tab treats a non-object ``usage`` as
absent and
+ hides the stat boxes instead of rendering ``undefined`` in each one.
* **Content is off by default.** Only token counts, model id, latency, tool
names, and finish reason are recorded. Prompt and completion text is never
emitted unless you opt in (see below).
diff --git a/providers/common/ai/provider.yaml
b/providers/common/ai/provider.yaml
index fb3ccdb768e..974086387c9 100644
--- a/providers/common/ai/provider.yaml
+++ b/providers/common/ai/provider.yaml
@@ -95,6 +95,8 @@ hooks:
plugins:
- name: hitl_review
plugin-class:
airflow.providers.common.ai.plugins.hitl_review.HITLReviewPlugin
+ - name: ai_model_panel
+ plugin-class:
airflow.providers.common.ai.plugins.model_panel.ModelPanelPlugin
config:
diff --git a/providers/common/ai/pyproject.toml
b/providers/common/ai/pyproject.toml
index cecd1f028da..754f0d5e707 100644
--- a/providers/common/ai/pyproject.toml
+++ b/providers/common/ai/pyproject.toml
@@ -222,6 +222,7 @@ provider_info =
"airflow.providers.common.ai.get_provider_info:get_provider_info
[project.entry-points."airflow.plugins"]
hitl_review =
"airflow.providers.common.ai.plugins.hitl_review:HITLReviewPlugin"
+ai_model_panel =
"airflow.providers.common.ai.plugins.model_panel:ModelPanelPlugin"
[tool.hatch.version]
path = "src/airflow/providers/common/ai/__init__.py"
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 f50b3b5256a..b7affc888f1 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
@@ -101,7 +101,11 @@ def get_provider_info():
{
"name": "hitl_review",
"plugin-class":
"airflow.providers.common.ai.plugins.hitl_review.HITLReviewPlugin",
- }
+ },
+ {
+ "name": "ai_model_panel",
+ "plugin-class":
"airflow.providers.common.ai.plugins.model_panel.ModelPanelPlugin",
+ },
],
"config": {
"common.ai": {
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/_plugin_urls.py
b/providers/common/ai/src/airflow/providers/common/ai/plugins/_plugin_urls.py
new file mode 100644
index 00000000000..a1e8517d514
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/_plugin_urls.py
@@ -0,0 +1,50 @@
+# 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.
+"""URL helpers shared by this provider's React plugins (``hitl_review``,
``model_panel``)."""
+
+from __future__ import annotations
+
+from urllib.parse import urlparse
+
+from airflow.providers.common.compat.sdk import conf
+
+
+def get_base_url_path(path: str) -> str:
+ """Construct URL path with webserver base_url prefix for non-root
deployments."""
+ base_url = conf.get("api", "base_url", fallback="/")
+ if base_url.startswith(("http://";, "https://";)):
+ base_path = urlparse(base_url).path
+ else:
+ base_path = base_url
+ base_path = base_path.rstrip("/")
+ return base_path + path
+
+
+def get_bundle_url(plugin_prefix: str, bundle_filename: str) -> str:
+ """
+ Return the bundle URL for a plugin's React bundle.
+
+ Uses an absolute URL when api.base_url is a full URL so the bundle loads
+ correctly in Vite dev mode, where import() resolves relative to the script
+ origin (5173) rather than the document origin (28080).
+ """
+ path = get_base_url_path(f"{plugin_prefix}/static/{bundle_filename}")
+ base_url = conf.get("api", "base_url", fallback="/")
+ if base_url.startswith(("http://";, "https://";)):
+ parsed = urlparse(base_url)
+ return f"{parsed.scheme}://{parsed.netloc}" + path
+ return path
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/hitl_review.py
b/providers/common/ai/src/airflow/providers/common/ai/plugins/hitl_review.py
index 2945e83f9d0..527e72427ab 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/plugins/hitl_review.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/plugins/hitl_review.py
@@ -18,14 +18,11 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Annotated, Any
-from urllib.parse import urlparse
from uuid import UUID
from airflow.plugins_manager import AirflowPlugin
-from airflow.providers.common.compat.sdk import conf
-from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_1_PLUS,
get_base_airflow_version_tuple
-
-AIRFLOW_V_3_4_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0)
+from airflow.providers.common.ai.plugins._plugin_urls import get_bundle_url
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_4_PLUS
if TYPE_CHECKING:
from airflow.plugins_manager import FastAPIAppDict, ReactAppDict
@@ -33,33 +30,6 @@ if TYPE_CHECKING:
_PLUGIN_PREFIX = "/hitl-review"
-def _get_base_url_path(path: str) -> str:
- """Construct URL path with webserver base_url prefix for non-root
deployments."""
- base_url = conf.get("api", "base_url", fallback="/")
- if base_url.startswith(("http://";, "https://";)):
- base_path = urlparse(base_url).path
- else:
- base_path = base_url
- base_path = base_path.rstrip("/")
- return base_path + path
-
-
-def _get_bundle_url() -> str:
- """
- Return bundle URL for the React plugin.
-
- Uses an absolute URL when api.base_url is a full URL so the bundle loads
- correctly in Vite dev mode, where import() resolves relative to the script
- origin (5173) rather than the document origin (28080).
- """
- path = _get_base_url_path(f"{_PLUGIN_PREFIX}/static/main.umd.cjs")
- base_url = conf.get("api", "base_url", fallback="/")
- if base_url.startswith(("http://";, "https://";)):
- parsed = urlparse(base_url)
- return f"{parsed.scheme}://{parsed.netloc}" + path
- return path
-
-
if AIRFLOW_V_3_1_PLUS:
import mimetypes
from pathlib import Path
@@ -658,7 +628,7 @@ class HITLReviewPlugin(AirflowPlugin):
react_apps = [
{
"name": "HITL Review",
- "bundle_url": _get_bundle_url(),
+ "bundle_url": get_bundle_url(_PLUGIN_PREFIX, "main.umd.cjs"),
"destination": "task_instance",
"url_route": "hitl-review",
}
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/model_panel.py
b/providers/common/ai/src/airflow/providers/common/ai/plugins/model_panel.py
new file mode 100644
index 00000000000..966cd3acc4f
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/plugins/model_panel.py
@@ -0,0 +1,97 @@
+# 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 typing import TYPE_CHECKING
+
+from airflow.plugins_manager import AirflowPlugin
+from airflow.providers.common.ai.plugins._plugin_urls import get_bundle_url
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_4_PLUS
+
+if TYPE_CHECKING:
+ from airflow.plugins_manager import FastAPIAppDict, ReactAppDict
+
+_PLUGIN_PREFIX = "/ai-model"
+
+
+if AIRFLOW_V_3_1_PLUS:
+ import mimetypes
+ from pathlib import Path
+
+ from fastapi import FastAPI
+ from fastapi.staticfiles import StaticFiles
+
+ # Ensure proper MIME type for the plugin bundle (FastAPI serves .cjs as
text/plain by default).
+ mimetypes.add_type("application/javascript", ".cjs")
+
+ model_panel_app = FastAPI(
+ title="AI Model Panel",
+ description="Serves the static 'Model' tab bundle for task instances
run with common.ai operators.",
+ )
+
+ _WWW_DIR = Path(__file__).parent / "www"
+ _dist_dir = _WWW_DIR / "dist"
+ if _dist_dir.is_dir():
+ model_panel_app.mount(
+ "/static",
+ StaticFiles(directory=str(_dist_dir.absolute()), html=True),
+ name="model_panel_static",
+ )
+
+
+class ModelPanelPlugin(AirflowPlugin):
+ """Register the 'Model' tab showing the resolved LLM model name and token
usage."""
+
+ name = "ai_model_panel"
+ fastapi_apps: list[FastAPIAppDict] = []
+ react_apps: list[ReactAppDict] = []
+ if AIRFLOW_V_3_1_PLUS:
+ fastapi_apps = [
+ {
+ "name": "ai-model-panel",
+ "app": model_panel_app,
+ "url_prefix": _PLUGIN_PREFIX,
+ }
+ ]
+ if AIRFLOW_V_3_4_PLUS:
+ # `applies_to` (the per-operator scoping below) only came into core
with
+ # apache/airflow#69148 (3.4.0). An older core silently accepts the
extra key and
+ # ignores it, so the tab would show up on every task instance instead
and land on
+ # the empty state for anything that isn't an LLM/Agent run -- the
whole entry is
+ # withheld pre-3.4 rather than shipping that unscoped tab.
+ react_apps = [
+ {
+ "name": "AI Model",
+ "bundle_url": get_bundle_url(_PLUGIN_PREFIX, "model.umd.cjs"),
+ "destination": "task_instance",
+ "url_route": "ai-model",
+ # Only LLMOperator and AgentOperator (and their
@task.llm/@task.agent
+ # TaskFlow forms) push model_name/usage; other operators would
show
+ # a permanently-empty tab. `operator_name` alone covers both
the raw
+ # class name and the decorator's display name, since
TaskInstance falls
+ # back to the class name when there's no custom_operator_name.
+ "applies_to": {
+ "operator_names": [
+ "LLMOperator",
+ "AgentOperator",
+ "@task.llm",
+ "@task.agent",
+ ],
+ },
+ }
+ ]
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/package.json
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/package.json
index 248e84c1c1e..7cb19050067 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/package.json
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/package.json
@@ -21,7 +21,7 @@
],
"scripts": {
"dev": "vite --port 5174 --strictPort",
- "build": "vite build",
+ "build": "vite build && vite build --mode model",
"build:types": "tsc --p tsconfig.lib.json",
"build:lib": "vite build",
"lint": "tsc --p tsconfig.app.json --noEmit",
@@ -33,6 +33,7 @@
"@emotion/react": "^11.14.0",
"react": "^19.2.4",
"react-dom": "^19.2.4",
+ "react-icons": "^5.7.0",
"react-markdown": "^10.1.0",
"remark-gfm": "^4.0.1"
},
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/pnpm-lock.yaml
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/pnpm-lock.yaml
index 55fbf3df6a2..8fde3e887d9 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/pnpm-lock.yaml
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/pnpm-lock.yaml
@@ -34,6 +34,9 @@ importers:
react-dom:
specifier: ^19.2.4
version: 19.2.4([email protected])
+ react-icons:
+ specifier: ^5.7.0
+ version: 5.7.0([email protected])
react-markdown:
specifier: ^10.1.0
version: 10.1.0(@types/[email protected])([email protected])
@@ -1598,6 +1601,11 @@ packages:
peerDependencies:
react: ^19.2.4
+ [email protected]:
+ resolution: {integrity:
sha512-LBLy340Rzqy6+/yVhZKT3B/QpP1BZaesGqasf09HPOBzRarcDIFH0WwXlXQfE7q7ipxK4MSiC5DIBWURCny6fw==}
+ peerDependencies:
+ react: '*'
+
[email protected]:
resolution: {integrity:
sha512-24e6ynE2H+OKt4kqsOvNd8kBpV65zoxbA4BVsEOB3ARVWQki/DHzaUoC5KuON/BiccDaCCTZBuOcfZs70kR8bQ==}
@@ -3837,6 +3845,10 @@ snapshots:
react: 19.2.4
scheduler: 0.27.0
+ [email protected]([email protected]):
+ dependencies:
+ react: 19.2.4
+
[email protected]: {}
[email protected](@types/[email protected])([email protected]):
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/ModelPanel.tsx
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/ModelPanel.tsx
new file mode 100644
index 00000000000..7d729679484
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/ModelPanel.tsx
@@ -0,0 +1,142 @@
+/*!
+ * 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.
+ */
+
+import { Badge, Box, Flex, HStack, Heading, SimpleGrid, Spinner, Text } from
"@chakra-ui/react";
+import type { FC, ReactNode } from "react";
+import { LuInfo } from "react-icons/lu";
+
+import { NoModelInfo } from "src/components/NoModelInfo";
+import { Tooltip } from "src/components/Tooltip";
+import { useModelInfo } from "src/hooks/useModelInfo";
+
+interface ModelPanelProps {
+ dagId: string;
+ runId: string;
+ taskId: string;
+ mapIndex: number;
+}
+
+// Pydantic-ai prices a response via genai-prices' per-token rate lookup (see
+// durable/replay_usage.py::fill_replayed_cost), not from provider-reported
billing --
+// a model or provider the lookup doesn't recognise is left unpriced rather
than guessed.
+const COST_TOOLTIP =
+ "Estimated in USD from a per-token price lookup (genai-prices), not
provider-reported " +
+ "billing -- best-effort, and may be stale or unavailable for some
models/providers.";
+
+// The XCom keeps full precision (see utils/logging.py's
format_usage_for_xcom) so a very
+// cheap run doesn't look free in the logs; round to cents here for a readable
display. A
+// non-zero sub-cent cost is called out explicitly rather than rounded down to
"$0.00",
+// which would read as free.
+function formatCost(cost: string | null | undefined): string {
+ if (cost === null || cost === undefined) return "—";
+ const value = Number(cost);
+ if (value > 0 && value < 0.01) return "< $0.01";
+ return `$${value.toFixed(2)}`;
+}
+
+const StatBox: FC<{ label: string; tooltip?: ReactNode; value: number | string
}> = ({
+ label,
+ tooltip,
+ value,
+}) => (
+ <Box bg="bg.subtle" borderRadius="lg" borderWidth="1px" p={4}>
+ <HStack gap={0.5}>
+ <Text color="fg.muted" fontSize="xs">
+ {label}
+ </Text>
+ {tooltip !== undefined && (
+ <Tooltip content={tooltip} portalled>
+ <Box aria-label={`About ${label}`} as="button" color="fg.muted"
p={0.5}>
+ <LuInfo />
+ </Box>
+ </Tooltip>
+ )}
+ </HStack>
+ <Text fontSize="lg" fontWeight="semibold">
+ {value}
+ </Text>
+ </Box>
+);
+
+export const ModelPanel: FC<ModelPanelProps> = ({ dagId, runId, taskId,
mapIndex }) => {
+ const { modelName, usage, loading, error } = useModelInfo(dagId, runId,
taskId, mapIndex);
+
+ if (loading) {
+ return (
+ <Flex align="center" gap={2} p={2}>
+ <Spinner colorPalette="brand" size="sm" />
+ <Text color="fg.muted" fontSize="sm">
+ Loading model info...
+ </Text>
+ </Flex>
+ );
+ }
+
+ if (!modelName && !usage) {
+ return <NoModelInfo error={error} />;
+ }
+
+ return (
+ <Box p={2}>
+ <Box borderBottomWidth="1px" mb={4} pb={4}>
+ <Heading size="sm">AI Model</Heading>
+ <HStack color="fg.muted" fontSize="sm" gap={3} mt={1}>
+ <Text as="span">
+ <Text as="b">Task:</Text> {taskId}
+ </Text>
+ <Text as="span">
+ <Text as="b">DAG:</Text> {dagId}
+ </Text>
+ </HStack>
+ </Box>
+
+ {error && (
+ <Box
+ bg="red.subtle"
+ borderColor="red.emphasized"
+ borderRadius="lg"
+ borderWidth="1px"
+ color="red.fg"
+ fontSize="sm"
+ mb={4}
+ p={3}
+ >
+ {error}
+ </Box>
+ )}
+
+ {modelName && (
+ <Badge borderRadius="full" colorPalette="brand" fontSize="sm" mb={5}
px={3} py={1}>
+ {modelName}
+ </Badge>
+ )}
+
+ {usage && (
+ <SimpleGrid columns={2} gap={4} maxW="640px">
+ <StatBox label="Requests" value={usage.requests} />
+ <StatBox label="Tool calls" value={usage.tool_calls} />
+ <StatBox label="Input tokens" value={usage.input_tokens} />
+ <StatBox label="Output tokens" value={usage.output_tokens} />
+ <StatBox label="Total tokens" value={usage.total_tokens} />
+ <StatBox label="Cost" tooltip={COST_TOOLTIP}
value={formatCost(usage.cost)} />
+ </SimpleGrid>
+ )}
+ </Box>
+ );
+};
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/NoModelInfo.tsx
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/NoModelInfo.tsx
new file mode 100644
index 00000000000..97c50dc3eb2
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/NoModelInfo.tsx
@@ -0,0 +1,58 @@
+/*!
+ * 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.
+ */
+
+import { Box, Code, Text } from "@chakra-ui/react";
+import type { FC } from "react";
+
+interface NoModelInfoProps {
+ error?: string | null;
+}
+
+export const NoModelInfo: FC<NoModelInfoProps> = ({ error }) => {
+ return (
+ <Box p={5}>
+ <Box
+ bg="bg.subtle"
+ borderRadius="xl"
+ borderWidth="1px"
+ maxW="440px"
+ mx="auto"
+ p={10}
+ textAlign="center"
+ >
+ <Text fontSize="4xl" mb={4}>
+ 🤖
+ </Text>
+ <Text as="h2" fontSize="lg" fontWeight="semibold" mb={2}>
+ No AI Model Info
+ </Text>
+ <Text color="fg.muted" fontSize="sm" lineHeight="tall">
+ {error ?? (
+ <>
+ This task hasn't published a <Code
fontSize="xs">model_name</Code> or{" "}
+ <Code fontSize="xs">usage</Code> XCom yet -- if it's still
running, check back after
+ it finishes. The Model tab shows data for tasks run with{" "}
+ <Code fontSize="xs">LLMOperator</Code> or <Code
fontSize="xs">AgentOperator</Code>.
+ </>
+ )}
+ </Text>
+ </Box>
+ </Box>
+ );
+};
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/Tooltip.tsx
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/Tooltip.tsx
new file mode 100644
index 00000000000..7cbb459fb6e
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/Tooltip.tsx
@@ -0,0 +1,65 @@
+/*!
+ * 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.
+ */
+import { type ReactNode, type RefObject, forwardRef } from "react";
+
+import { Tooltip as ChakraTooltip, Portal } from "@chakra-ui/react";
+
+export type TooltipProps = {
+ readonly content: ReactNode;
+ readonly contentProps?: ChakraTooltip.ContentProps;
+ readonly disabled?: boolean;
+ readonly portalled?: boolean;
+ readonly portalRef?: RefObject<HTMLElement>;
+ readonly showArrow?: boolean;
+} & ChakraTooltip.RootProps;
+
+export const Tooltip = forwardRef<HTMLDivElement, TooltipProps>((props, ref)
=> {
+ const {
+ children,
+ content,
+ contentProps,
+ disabled,
+ portalled,
+ portalRef,
+ showArrow = true,
+ ...rest
+ } = props;
+
+ if (disabled) {
+ return children;
+ }
+
+ return (
+ <ChakraTooltip.Root {...rest}>
+ <ChakraTooltip.Trigger asChild>{children}</ChakraTooltip.Trigger>
+ <Portal container={portalRef} disabled={!portalled}>
+ <ChakraTooltip.Positioner zIndex="tooltip">
+ <ChakraTooltip.Content ref={ref} {...contentProps}>
+ {showArrow ? (
+ <ChakraTooltip.Arrow>
+ <ChakraTooltip.ArrowTip />
+ </ChakraTooltip.Arrow>
+ ) : undefined}
+ {content}
+ </ChakraTooltip.Content>
+ </ChakraTooltip.Positioner>
+ </Portal>
+ </ChakraTooltip.Root>
+ );
+});
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/hooks/useModelInfo.ts
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/hooks/useModelInfo.ts
new file mode 100644
index 00000000000..8b8bc24dbd8
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/hooks/useModelInfo.ts
@@ -0,0 +1,67 @@
+/*!
+ * 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.
+ */
+
+import { useCallback, useEffect, useState } from "react";
+
+import { createModelApi } from "src/model-api";
+import type { UsageInfo } from "src/types/model";
+
+interface UseModelInfoReturn {
+ modelName: string | null;
+ usage: UsageInfo | null;
+ loading: boolean;
+ error: string | null;
+}
+
+function describeFailure(result: PromiseSettledResult<unknown>): string | null
{
+ if (result.status === "fulfilled") return null;
+ return result.reason instanceof Error ? result.reason.message :
String(result.reason);
+}
+
+export function useModelInfo(
+ dagId: string,
+ runId: string,
+ taskId: string,
+ mapIndex: number,
+): UseModelInfoReturn {
+ const [modelName, setModelName] = useState<string | null>(null);
+ const [usage, setUsage] = useState<UsageInfo | null>(null);
+ const [loading, setLoading] = useState(true);
+ const [error, setError] = useState<string | null>(null);
+
+ // Built fresh on every call (not memoized via a ref) so a switch to another
task or map
+ // index -- which changes these props without remounting the panel -- is
reflected here too.
+ const fetchInfo = useCallback(async () => {
+ setLoading(true);
+ const api = createModelApi(dagId, runId, taskId, mapIndex);
+ // allSettled, not all: one XCom missing or erroring must not hide a model
name or usage
+ // the other call did resolve.
+ const [nameResult, usageResult] = await
Promise.allSettled([api.fetchModelName(), api.fetchUsage()]);
+ setModelName(nameResult.status === "fulfilled" ? nameResult.value : null);
+ setUsage(usageResult.status === "fulfilled" ? usageResult.value : null);
+ setError(describeFailure(nameResult) ?? describeFailure(usageResult));
+ setLoading(false);
+ }, [dagId, runId, taskId, mapIndex]);
+
+ useEffect(() => {
+ void fetchInfo();
+ }, [fetchInfo]);
+
+ return { modelName, usage, loading, error };
+}
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model-api.ts
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model-api.ts
new file mode 100644
index 00000000000..66770627c5f
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model-api.ts
@@ -0,0 +1,83 @@
+/*!
+ * 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.
+ */
+
+import type { UsageInfo } from "src/types/model";
+
+function getApiBase(): string {
+ if (typeof document === "undefined") return "/api/v2";
+ const baseHref = document.querySelector("head > base")?.getAttribute("href")
?? "";
+ const baseUrl = new URL(baseHref, globalThis.location.origin);
+ const basePath = baseUrl.pathname.replace(/\/$/, "") || "";
+ return basePath ? `${basePath}/api/v2` : "/api/v2";
+}
+
+const API_BASE = getApiBase();
+
+// Namespaced to match
airflow.providers.common.ai.utils.logging.MODEL_NAME_XCOM_KEY --
+// avoids colliding with a user's own "model_name" XCom.
+const MODEL_NAME_XCOM_KEY = "__AIRFLOW__COMMON_AI_MODEL_NAME__";
+
+/** Read one XCom entry; returns `null` when the task hasn't pushed that key
(404). */
+async function fetchXComValue<T>(
+ dagId: string,
+ runId: string,
+ taskId: string,
+ mapIndex: number,
+ key: string,
+): Promise<T | null> {
+ const path =
+ `${API_BASE}/dags/${encodeURIComponent(dagId)}` +
+ `/dagRuns/${encodeURIComponent(runId)}` +
+ `/taskInstances/${encodeURIComponent(taskId)}` +
+ `/xcomEntries/${encodeURIComponent(key)}`;
+ const res = await fetch(`${path}?map_index=${mapIndex}&deserialize=true`, {
+ credentials: "same-origin",
+ });
+ if (res.status === 404) return null;
+ if (!res.ok) {
+ const body = await res.json().catch(() => ({}));
+ const detail = (body as { detail?: unknown }).detail;
+ // `statusText` is always empty over HTTP/2, and a 422's `detail` is an
array of
+ // validation errors (which would print as "[object Object]") rather than
a string --
+ // fall back to the status code so the error is never a blank or useless
message.
+ throw new Error(typeof detail === "string" && detail.length > 0 ? detail :
`HTTP ${res.status}`);
+ }
+ const data = (await res.json()) as { value: T };
+ return data.value;
+}
+
+function isUsageInfo(value: unknown): value is UsageInfo {
+ return typeof value === "object" && value !== null;
+}
+
+export function createModelApi(dagId: string, runId: string, taskId: string,
mapIndex: number) {
+ return {
+ fetchModelName: () =>
+ fetchXComValue<string>(dagId, runId, taskId, mapIndex,
MODEL_NAME_XCOM_KEY),
+ fetchUsage: async () => {
+ const value = await fetchXComValue<unknown>(dagId, runId, taskId,
mapIndex, "usage");
+ // A reference-storing XCom backend (e.g. common.io with a low
+ // xcom_objectstorage_threshold) returns the stored reference string
here instead of
+ // the real value, since this public endpoint never runs the backend's
+ // deserialize_value. Treat anything that isn't a plain object as absent
rather than
+ // let every stat box below render "undefined".
+ return isUsageInfo(value) ? value : null;
+ },
+ };
+}
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model.tsx
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model.tsx
new file mode 100644
index 00000000000..a5bad684625
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model.tsx
@@ -0,0 +1,73 @@
+/*!
+ * 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.
+ */
+
+import { Box, ChakraProvider } from "@chakra-ui/react";
+import type { FC } from "react";
+
+import { ModelPanel } from "src/components/ModelPanel";
+import { NoModelInfo } from "src/components/NoModelInfo";
+
+import { localSystem } from "./theme";
+
+export interface PluginComponentProps {
+ dagId?: string;
+ mapIndex?: string;
+ runId?: string;
+ taskId?: string;
+}
+
+/**
+ * Main plugin component for the AI Model tab.
+ *
+ * Receives dagId, runId, taskId, mapIndex as props from the Airflow React
plugin
+ * host (from route params). Renders ModelPanel when all params are present,
+ * otherwise shows NoModelInfo fallback.
+ */
+const PluginComponent: FC<PluginComponentProps> = ({
+ dagId = "",
+ runId = "",
+ taskId = "",
+ mapIndex: mapIndexProp = "-1",
+}) => {
+ const mapIndex = /^-?\d+$/.test(String(mapIndexProp)) ?
parseInt(String(mapIndexProp), 10) : -1;
+
+ if (!dagId || !runId || !taskId) {
+ return <NoModelInfo />;
+ }
+
+ return <ModelPanel dagId={dagId} runId={runId} taskId={taskId}
mapIndex={mapIndex} />;
+};
+
+/**
+ * Plugin component wrapped with ChakraProvider for consistent theming with
the host.
+ * Chakra semantic tokens handle light/dark mode automatically.
+ */
+const WrappedPluginComponent: FC<PluginComponentProps> = (props) => {
+ const system = (globalThis as Record<string, unknown>).ChakraUISystem ??
localSystem;
+
+ return (
+ <ChakraProvider value={system}>
+ <Box height="100%" minHeight={0}>
+ <PluginComponent {...props} />
+ </Box>
+ </ChakraProvider>
+ );
+};
+
+export default WrappedPluginComponent;
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/types/model.ts
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/types/model.ts
new file mode 100644
index 00000000000..a15438f8555
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/types/model.ts
@@ -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.
+ */
+
+/** Shape of the `usage` XCom pushed by LLMOperator/AgentOperator (see
format_usage_for_xcom). */
+export interface UsageInfo {
+ requests: number;
+ input_tokens: number;
+ output_tokens: number;
+ total_tokens: number;
+ tool_calls: number;
+ cost: string | null;
+}
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/vite.config.ts
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/vite.config.ts
index f2507cf8f58..efdf271c1db 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/vite.config.ts
+++
b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/vite.config.ts
@@ -22,18 +22,29 @@ import cssInjectedByJsPlugin from
"vite-plugin-css-injected-by-js";
import dts from "vite-plugin-dts";
import { defineConfig } from "vite";
-export default defineConfig(({ command }) => {
+export default defineConfig(({ command, mode }) => {
const isLibraryBuild = command === "build";
+ // Vite's lib mode cannot bundle multiple entries into one UMD/IIFE output
(each needs its
+ // own global name), so each plugin bundle is built via a separate `vite
build --mode <name>`
+ // invocation into the same dist/ directory -- see package.json's "build"
script.
+ const isModelEntry = mode === "model";
+ const entryFile = isModelEntry ? "model.tsx" : "main.tsx";
+ const entryName = isModelEntry ? "model" : "main";
return {
base: "./",
build: isLibraryBuild
? {
chunkSizeWarningLimit: 1600,
+ // Only the first build of a `pnpm build` run should clear stale
dist/ output.
+ emptyOutDir: !isModelEntry,
lib: {
- entry: resolve("src", "main.tsx"),
- fileName: "main",
+ entry: resolve("src", entryFile),
+ fileName: entryName,
formats: ["umd"],
+ // The host (ReactPlugin.tsx) clears `globalThis.AirflowPlugin`
before each dynamic
+ // import and recaptures it under `globalThis[reactApp.name]`
right after -- every
+ // plugin bundle must use this same generic UMD global name, not a
bundle-specific one.
name: "AirflowPlugin",
},
rollupOptions: {
@@ -70,7 +81,7 @@ export default defineConfig(({ command }) => {
...(isLibraryBuild
? [
dts({
- include: ["src/main.tsx"],
+ include: [`src/${entryFile}`],
insertTypesEntry: true,
outDir: "dist",
}),
@@ -81,6 +92,10 @@ export default defineConfig(({ command }) => {
server: {
cors: true,
proxy: {
+ "/api": {
+ changeOrigin: true,
+ target: "http://localhost:28080";,
+ },
"/hitl-review": {
changeOrigin: true,
target: "http://localhost:28080";,
diff --git
a/providers/common/ai/tests/unit/common/ai/plugins/test__plugin_urls.py
b/providers/common/ai/tests/unit/common/ai/plugins/test__plugin_urls.py
new file mode 100644
index 00000000000..a586cdbfe9f
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/plugins/test__plugin_urls.py
@@ -0,0 +1,46 @@
+# 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 airflow.providers.common.ai.plugins._plugin_urls import
get_base_url_path, get_bundle_url
+
+from tests_common.test_utils.config import conf_vars
+
+
+class TestGetBaseUrlPath:
+ def test_default_base_url(self):
+ with conf_vars({("api", "base_url"): "/"}):
+ assert get_base_url_path("/ai-model") == "/ai-model"
+
+ def test_http_base_url_extracts_path(self):
+ with conf_vars({("api", "base_url"): "http://example.com/airflow/"}):
+ assert get_base_url_path("/ai-model") == "/airflow/ai-model"
+
+
+class TestGetBundleUrl:
+ def test_default_base_url_returns_a_relative_path(self):
+ with conf_vars({("api", "base_url"): "/"}):
+ assert get_bundle_url("/ai-model", "model.umd.cjs") ==
"/ai-model/static/model.umd.cjs"
+
+ def test_http_base_url_returns_an_absolute_url(self):
+ """An absolute bundle_url is needed in Vite dev mode, where import()
resolves
+ relative to the script origin (5173), not the document origin
(28080)."""
+ with conf_vars({("api", "base_url"): "http://example.com/airflow/"}):
+ assert (
+ get_bundle_url("/ai-model", "model.umd.cjs")
+ == "http://example.com/airflow/ai-model/static/model.umd.cjs";
+ )
diff --git
a/providers/common/ai/tests/unit/common/ai/plugins/test_hitl_review.py
b/providers/common/ai/tests/unit/common/ai/plugins/test_hitl_review.py
index 783df6cb5fb..40d91d9b402 100644
--- a/providers/common/ai/tests/unit/common/ai/plugins/test_hitl_review.py
+++ b/providers/common/ai/tests/unit/common/ai/plugins/test_hitl_review.py
@@ -38,7 +38,6 @@ from airflow.models.xcom import XComModel
from airflow.providers.common.ai.plugins.hitl_review import (
HITLReviewPlugin,
_build_session_response,
- _get_base_url_path,
_get_map_index,
_is_task_completed,
_read_xcom,
@@ -680,16 +679,6 @@ class TestWriteXcom:
_clear_db()
-class TestGetBaseUrlPath:
- def test_default_base_url(self):
- with conf_vars({("api", "base_url"): "/"}):
- assert _get_base_url_path("/hitl-review") == "/hitl-review"
-
- def test_http_base_url_extracts_path(self):
- with conf_vars({("api", "base_url"): "http://example.com/airflow/"}):
- assert _get_base_url_path("/hitl-review") == "/airflow/hitl-review"
-
-
class TestIsTaskCompleted:
"""Test _is_task_completed."""
diff --git
a/providers/common/ai/tests/unit/common/ai/plugins/test_model_panel.py
b/providers/common/ai/tests/unit/common/ai/plugins/test_model_panel.py
new file mode 100644
index 00000000000..8c27a43cd70
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/plugins/test_model_panel.py
@@ -0,0 +1,72 @@
+# 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
+
+import pytest
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_4_PLUS
+
+if not AIRFLOW_V_3_1_PLUS:
+ pytest.skip("The AI Model panel is only compatible with Airflow >= 3.1.0",
allow_module_level=True)
+
+from airflow.providers.common.ai.decorators.agent import
_AgentDecoratedOperator
+from airflow.providers.common.ai.decorators.llm import _LLMDecoratedOperator
+from airflow.providers.common.ai.operators.agent import AgentOperator
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.plugins.model_panel import ModelPanelPlugin
+
+requires_3_4 = pytest.mark.skipif(
+ not AIRFLOW_V_3_4_PLUS, reason="react_apps' applies_to scoping needs
Airflow >= 3.4.0"
+)
+
+
+class TestModelPanelPlugin:
+ def test_plugin_name(self):
+ assert ModelPanelPlugin.name == "ai_model_panel"
+
+ def test_fastapi_apps_registered(self):
+ assert len(ModelPanelPlugin.fastapi_apps) == 1
+ assert ModelPanelPlugin.fastapi_apps[0]["name"] == "ai-model-panel"
+ assert "url_prefix" in ModelPanelPlugin.fastapi_apps[0]
+
+ @requires_3_4
+ def test_react_apps_registered(self):
+ assert len(ModelPanelPlugin.react_apps) == 1
+ app = ModelPanelPlugin.react_apps[0]
+ assert app["name"] == "AI Model"
+ assert app["url_route"] == "ai-model"
+ assert app["destination"] == "task_instance"
+ assert "model.umd.cjs" in app["bundle_url"]
+
+ @pytest.mark.skipif(
+ AIRFLOW_V_3_4_PLUS, reason="exercises the pre-3.4 gate that withholds
react_apps entirely"
+ )
+ def test_react_apps_withheld_before_3_4(self):
+ assert ModelPanelPlugin.react_apps == []
+
+ @requires_3_4
+ def test_applies_to_scopes_to_operators_that_publish_model_name(self):
+ """Tied to the actual operator/decorator names a task instance would
record,
+ so a rename of either doesn't silently desync the tab's scoping from
reality."""
+ app = ModelPanelPlugin.react_apps[0]
+ operator_names = app["applies_to"]["operator_names"]
+ assert set(operator_names) == {
+ LLMOperator.__name__,
+ AgentOperator.__name__,
+ _LLMDecoratedOperator.custom_operator_name,
+ _AgentDecoratedOperator.custom_operator_name,
+ }
diff --git
a/providers/common/ai/tests/unit/common/ai/plugins/test_model_panel_compat.py
b/providers/common/ai/tests/unit/common/ai/plugins/test_model_panel_compat.py
new file mode 100644
index 00000000000..aacd250f5fc
--- /dev/null
+++
b/providers/common/ai/tests/unit/common/ai/plugins/test_model_panel_compat.py
@@ -0,0 +1,46 @@
+# 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
+
+import pytest
+
+from airflow.providers.common.ai.plugins.model_panel import ModelPanelPlugin
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_4_PLUS
+
+
+def test_model_panel_plugin_registration_matches_airflow_version():
+ # react_apps needs the applies_to scoping that only core >= 3.4
understands (see
+ # model_panel.py); fastapi_apps (the static bundle server) has no such
dependency.
+ if AIRFLOW_V_3_1_PLUS:
+ assert len(ModelPanelPlugin.fastapi_apps) == 1
+ else:
+ assert ModelPanelPlugin.fastapi_apps == []
+ if AIRFLOW_V_3_4_PLUS:
+ assert len(ModelPanelPlugin.react_apps) == 1
+ else:
+ assert ModelPanelPlugin.react_apps == []
+
+
[email protected](not AIRFLOW_V_3_1_PLUS, reason="Requires Airflow 3.1+")
+def test_model_panel_plugin_registers_expected_fastapi_app_name():
+ assert ModelPanelPlugin.fastapi_apps[0]["name"] == "ai-model-panel"
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="react_apps' applies_to
scoping needs Airflow >= 3.4.0")
+def test_model_panel_plugin_registers_expected_react_app_name():
+ assert ModelPanelPlugin.react_apps[0]["name"] == "AI Model"
diff --git a/scripts/ci/prek/update_providers_dependencies.py
b/scripts/ci/prek/update_providers_dependencies.py
index d5c2297182a..72cf23b7ce2 100755
--- a/scripts/ci/prek/update_providers_dependencies.py
+++ b/scripts/ci/prek/update_providers_dependencies.py
@@ -59,7 +59,7 @@ errors: list[str] = []
suspended_paths: list[str] = []
-ALL_DEPENDENCIES: dict[str, dict[str, list[str]]] = defaultdict(lambda:
defaultdict(list))
+ALL_DEPENDENCIES: dict[str, dict[str, list[Any]]] = defaultdict(lambda:
defaultdict(list))
ALL_PROVIDERS: dict[str, dict[str, Any]] = defaultdict(lambda: defaultdict())
ALL_PROVIDER_FILES: list[Path] = []
@@ -204,11 +204,13 @@ if __name__ == "__main__":
for error in errors:
console.print(f"[red] {error}")
console.print(f"[bright_blue]Total: {len(errors)} errors.")
- unique_sorted_dependencies: dict[str, dict[str, list[str] | str]] =
defaultdict(dict)
+ unique_sorted_dependencies: dict[str, dict[str, list[Any] | str]] =
defaultdict(dict)
for key in sorted(ALL_DEPENDENCIES.keys()):
unique_sorted_dependencies[key]["deps"] =
sorted(ALL_DEPENDENCIES[key]["deps"])
unique_sorted_dependencies[key]["devel-deps"] =
sorted(ALL_DEPENDENCIES[key]["devel-deps"])
- unique_sorted_dependencies[key]["plugins"] =
sorted(ALL_DEPENDENCIES[key]["plugins"])
+ unique_sorted_dependencies[key]["plugins"] = sorted(
+ ALL_DEPENDENCIES[key]["plugins"], key=lambda plugin: plugin["name"]
+ )
unique_sorted_dependencies[key]["cross-providers-deps"] = sorted(
set(ALL_DEPENDENCIES[key]["cross-providers-deps"])
)
