kaxil commented on code in PR #74334: URL: https://github.com/apache/airflow/pull/74334#discussion_r4201931277
########## providers/common/ai/src/airflow/providers/common/ai/plugins/model_panel.py: ########## @@ -0,0 +1,119 @@ +# 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 urllib.parse import urlparse + +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 + +if TYPE_CHECKING: + from airflow.plugins_manager import FastAPIAppDict, ReactAppDict + +_PLUGIN_PREFIX = "/ai-model" + + +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/model.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 + + 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, + } + ] + react_apps = [ + { + "name": "Model", + "bundle_url": _get_bundle_url(), + "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": { Review Comment: `applies_to` only came into core with #69148 (3.4.0), and no released core reads it. On 3.1.x through 3.3.x `ReactAppResponse` accepts the extra key and the UI ignores it, so the Model tab shows up on every task instance and lands on the empty state after two 404s. That's the opposite of what this comment and the PR description promise, on every core the provider runs on today. Could we either register `react_apps` only on 3.4+ (needs an `AIRFLOW_V_3_4_PLUS` in common.compat), or say in the docs that scoping needs 3.4 and make the panel tolerate an unrelated task's `usage` XCom (see the `formatCost` comment)? ########## providers/common/ai/src/airflow/providers/common/ai/plugins/model_panel.py: ########## @@ -0,0 +1,119 @@ +# 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 urllib.parse import urlparse + +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 + +if TYPE_CHECKING: + from airflow.plugins_manager import FastAPIAppDict, ReactAppDict + +_PLUGIN_PREFIX = "/ai-model" + + +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/model.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 + + 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): Review Comment: Nothing in `providers/common/ai/docs` mentions the tab yet. Could we add a few lines to `observability.rst`, next to the `usage` XCom description? Users will want to know that only `AgentOperator` fills in tokens and cost, and that the cost is a genai-prices estimate. The tab is also scoped through `operator_names`, which is an exact match, so a user's own `AgentOperator` subclass won't get it. ########## providers/common/ai/src/airflow/providers/common/ai/plugins/model_panel.py: ########## @@ -0,0 +1,119 @@ +# 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 urllib.parse import urlparse + +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 + +if TYPE_CHECKING: + from airflow.plugins_manager import FastAPIAppDict, ReactAppDict + +_PLUGIN_PREFIX = "/ai-model" + + +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/model.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 + + 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, + } + ] + react_apps = [ + { + "name": "Model", Review Comment: The host caches each plugin component on `globalThis[reactApp.name]` (`loadPlugin` in `ReactPlugin.tsx`), and `plugins_manager` only dedupes on `url_route`. Another plugin that also calls its tab "Model" would get whichever bundle loaded first. A more specific name like "AI Model" (matching the panel heading) avoids that. ########## providers/common/ai/src/airflow/providers/common/ai/plugins/model_panel.py: ########## @@ -0,0 +1,119 @@ +# 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 urllib.parse import urlparse + +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 + +if TYPE_CHECKING: + from airflow.plugins_manager import FastAPIAppDict, ReactAppDict + +_PLUGIN_PREFIX = "/ai-model" + + +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/model.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 + + 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, + } + ] + react_apps = [ + { + "name": "Model", + "bundle_url": _get_bundle_url(), + "destination": "task_instance", + "url_route": "ai-model", + # Only LLMOperator and AgentOperator (and their @task.llm/@task.agent Review Comment: `LLMOperator` (and so `@task.llm`) only pushes the model-name key. The only `usage` pushes are in `agent.py` (`_emit_run_metadata` and `_report_failed_run`), and `docs/observability.rst` already says `usage` comes only from `AgentOperator`. So on an LLM task the tab shows a model badge with no tokens or cost. Either push `format_usage_for_xcom(...)` from `LLMOperator.execute` as well, or correct this comment and the `UsageInfo` docstring in `types/model.ts` and render an explicit "usage not available" for that case. ########## providers/common/ai/tests/unit/common/ai/plugins/test_model_panel.py: ########## @@ -0,0 +1,61 @@ +# 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 + +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.plugins.model_panel import ModelPanelPlugin, _get_base_url_path + +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 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] + + def test_react_apps_registered(self): + assert len(ModelPanelPlugin.react_apps) == 1 + app = ModelPanelPlugin.react_apps[0] + assert app["name"] == "Model" + assert app["url_route"] == "ai-model" + assert app["destination"] == "task_instance" + assert "model.umd.cjs" in app["bundle_url"] + + def test_applies_to_scopes_to_operators_that_publish_model_name(self): + app = ModelPanelPlugin.react_apps[0] + operator_names = app["applies_to"]["operator_names"] + assert set(operator_names) == {"LLMOperator", "AgentOperator", "@task.llm", "@task.agent"} Review Comment: This compares the literal set to itself, so a rename of a decorator's `custom_operator_name` wouldn't fail it. Building the expected set from `LLMOperator.__name__`, `AgentOperator.__name__` and the two decorator classes' `custom_operator_name` would tie it to what the TI actually records. The absolute-URL branch of `_get_bundle_url` has no test either. ########## providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/model-api.ts: ########## @@ -0,0 +1,68 @@ +/*! + * 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?: string }).detail; + throw new Error(detail ?? res.statusText); Review Comment: This can throw `new Error("")`: over HTTP/2 `statusText` is always empty, so a 5xx with a non-JSON body (a proxy error page) has neither `detail` nor a status text. `NoModelInfo` then renders `error ?? (...)` as a blank card, since `""` isn't nullish. A 422 also returns `detail` as an array, which prints as `[object Object]`. Falling back to `` `HTTP ${res.status}` `` when `detail` isn't a string covers both. ########## 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, useRef, 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; + refetch: () => Promise<void>; +} + +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); + const apiRef = useRef(createModelApi(dagId, runId, taskId, mapIndex)); Review Comment: `apiRef` captures the props from the first render and `fetchInfo` has empty deps, so the hook never refetches. On a running agent task neither XCom exists yet, so the tab says "This task did not publish a model_name or usage XCom" and keeps saying it after the task finishes, until the user navigates away. On 3.1.0 to 3.1.6 the host's `ExternalView` also doesn't remount on a path change (`key={pathname}` arrived in 3.1.7), so switching task or map index keeps the previous task's numbers on screen. Keying the fetch on `[dagId, runId, taskId, mapIndex]` fixes the second. For the first, the panel can't tell "not finished yet" from "never will" without the TI state, so either fetch it and poll while it's non-terminal, or soften the empty-state wording. ########## providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/ModelPanel.tsx: ########## @@ -0,0 +1,137 @@ +/*! + * 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. +function formatCost(cost: string | null): string { + return cost === null ? "—" : `$${Number(cost).toFixed(2)}`; Review Comment: `toFixed(2)` rounds a sub-cent run down to zero dollars, which is the "looks free" case the comment above is guarding against. Showing something like "under a cent" below that threshold would keep the rounding without the misleading zero. Separately, `cost === null` doesn't catch `undefined`, so a `usage` XCom with a different shape (any task pushing its own `usage` key, which matters while the tab isn't scoped on released cores) renders `$NaN` next to blank stat boxes. ########## providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/ModelPanel.tsx: ########## @@ -0,0 +1,137 @@ +/*! + * 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. +function formatCost(cost: string | null): string { + return cost === null ? "—" : `$${Number(cost).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 && ( Review Comment: Can this branch render? `useModelInfo` only sets `error` when `Promise.all` rejects, and then `modelName` and `usage` are still their initial `null`, so line 86 returns `NoModelInfo` first. That also means a failed `usage` read throws away a successful model-name read. `Promise.allSettled` would make partial data plus an error reachable; otherwise this box, `refetch` in the hook and the `<Toaster />` in `model.tsx` (nothing in this bundle uses it) can go. ########## providers/common/ai/src/airflow/providers/common/ai/plugins/model_panel.py: ########## @@ -0,0 +1,119 @@ +# 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 urllib.parse import urlparse + +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 + +if TYPE_CHECKING: + from airflow.plugins_manager import FastAPIAppDict, ReactAppDict + +_PLUGIN_PREFIX = "/ai-model" + + +def _get_base_url_path(path: str) -> str: Review Comment: `_get_base_url_path` and `_get_bundle_url` are copies of the ones in `hitl_review.py` (only the bundle filename differs), and `getApiBase` / `model.tsx` mirror `api.ts` / `main.tsx` on the TS side. Could the Python pair move to a small shared private module under `plugins/` that takes the bundle name? `TestGetBaseUrlPath` here would then stop duplicating the one in `test_hitl_review.py`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
