From be2829b9175c167cd759238ca75d9eb88c46ead5 Mon Sep 17 00:00:00 2001 From: Pierre Jeambrun Date: Tue, 3 Nov 2026 19:19:51 +0100 Subject: [PATCH] Show which skills an agent was offered and used in the AI Model tab (#76309) * Show which skills an agent was offered and used in the AI Model tab When an agent behaves differently from one run to the next, the skills it had and the ones it loaded are the first thing to check, and a skill pulled from Git can change between runs without the Dag changing. Recording the cloned commit next to each skill turns "it behaved differently yesterday" into "the skill changed between these two commits". * Keep skill calls made before an approval pause or a durable retry A run resumed after a tool approval starts in a new process, and a durable retry replays earlier tool calls without running them, so counting calls as they pass through the toolset showed loaded skills as unused. Both get the earlier calls back as message history, so the counts now come from there. A second skills toolset on the same agent no longer replaces the first one's skills either. --- providers/common/ai/docs/observability.rst | 7 +- providers/common/ai/docs/toolsets/skills.rst | 20 ++ .../plugins/www/src/components/ModelPanel.tsx | 7 +- .../www/src/components/SkillsSection.tsx | 94 ++++++++ .../ai/plugins/www/src/hooks/useModelInfo.ts | 27 ++- .../common/ai/plugins/www/src/model-api.ts | 9 +- .../common/ai/plugins/www/src/skill-source.ts | 54 +++++ .../common/ai/plugins/www/src/types/model.ts | 17 ++ .../ai/plugins/www/tests/model.test.mjs | 44 +++- .../ai/plugins/www/tests/skill-source.test.ts | 69 ++++++ .../src/airflow/providers/common/ai/skills.py | 34 ++- .../providers/common/ai/toolsets/skills.py | 97 +++++++- .../providers/common/ai/utils/logging.py | 4 + .../ai/tests/unit/common/ai/test_skills.py | 15 +- .../unit/common/ai/toolsets/test_skills.py | 213 +++++++++++++++++- 15 files changed, 668 insertions(+), 43 deletions(-) create mode 100644 providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/SkillsSection.tsx create mode 100644 providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/skill-source.ts create mode 100644 providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/skill-source.test.ts diff --git a/providers/common/ai/docs/observability.rst b/providers/common/ai/docs/observability.rst index b1a6f4752db..8a2d537c73b 100644 --- a/providers/common/ai/docs/observability.rst +++ b/providers/common/ai/docs/observability.rst @@ -96,6 +96,10 @@ How it works push it; like ``run_id`` and ``usage``, with ``enable_hitl_review`` it is the initial run's model, not a human-feedback regeneration's, since ``regenerate_with_feedback`` doesn't re-push it. +* **Skills.** An ``AgentSkillsToolset`` publishes the skills it offered the agent, where + each came from (a local path, or a repository, path and commit) and the skill tools the + agent called for each, under the namespaced key ``__AIRFLOW__COMMON_AI_SKILLS__``. See + :ref:`agent-skills-observability`. * **Scope.** The ``run_id`` / ``usage`` XComs and the model name come from ``AgentOperator`` and ``LLMOperator`` and its subclasses (and their decorator forms). Those operators also stamp the ``airflow.*`` identity attributes on their spans, and so @@ -105,7 +109,8 @@ How it works **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 + readable table, including an estimated cost, and the skills XCom above as the list of + skills the agent was offered, with the ones it loaded first. It needs Airflow >= 3.4, since it relies on the ``applies_to`` operator scoping core added in that version to show up only for ``AgentOperator`` runs and runs of ``LLMOperator`` and its SQL / branch / file-analysis / schema-compare subclasses (and their decorator forms) -- the match is on the exact diff --git a/providers/common/ai/docs/toolsets/skills.rst b/providers/common/ai/docs/toolsets/skills.rst index ea7b33c3245..e5261b008d5 100644 --- a/providers/common/ai/docs/toolsets/skills.rst +++ b/providers/common/ai/docs/toolsets/skills.rst @@ -125,6 +125,26 @@ which then read ``reference.md`` and finished the run. ``exclude_resources`` hides files from the resource tools only. A skill's scripts can still read them, which is why the example excludes ``run_skill_script`` as well. +.. _agent-skills-observability: + +Seeing which skills a run used +------------------------------ + +When the toolset runs inside an Airflow task, it publishes on XCom, under the key +``__AIRFLOW__COMMON_AI_SKILLS__``, every skill it offered the agent: its name and +description, where it came from, and how many times the agent called ``load_skill``, +``read_skill_resource`` and ``run_skill_script`` for it. A local skill names its path on +the worker. A ``GitSkills`` skill names the repository, its path inside the repository, +the ``branch`` you set and the commit that was cloned, so a run that behaves differently +from an earlier one can be traced to a change in the skill between two commits. + +The :ref:`AI Model tab ` shows this list, loaded skills +first, with a link to the skill's folder at that commit for a repository on GitHub or +GitLab. The counts cover the whole conversation, so calls made before a tool-approval +pause, a human-review regeneration or a durable retry still count. Several skills toolsets +on one agent, each wrapped with ``prefixed()`` to keep their tool names apart, share the +same list. Nothing is published when the task sets ``do_xcom_push=False``. + Parameters ---------- 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 index 78941cf1b87..c092c3560cb 100644 --- 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 @@ -22,6 +22,7 @@ import type { FC, ReactNode } from "react"; import { LuInfo } from "react-icons/lu"; import { NoModelInfo } from "src/components/NoModelInfo"; +import { SkillsSection } from "src/components/SkillsSection"; import { Tooltip } from "src/components/Tooltip"; import { useModelInfo } from "src/hooks/useModelInfo"; import type { TaskRegion } from "src/region"; @@ -86,13 +87,13 @@ export const ModelLoading: FC = () => ( ); export const ModelPanel: FC = ({ dagId, runId, taskId, mapIndex, region }) => { - const { modelName, usage, loading, error } = useModelInfo(dagId, runId, taskId, mapIndex, region); + const { modelName, usage, skills, loading, error } = useModelInfo(dagId, runId, taskId, mapIndex, region); if (loading) { return ; } - if (!modelName && !usage) { + if (!modelName && !usage && !skills) { return ; } @@ -141,6 +142,8 @@ export const ModelPanel: FC = ({ dagId, runId, taskId, mapIndex )} + + {skills && skills.length > 0 && } ); }; diff --git a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/SkillsSection.tsx b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/SkillsSection.tsx new file mode 100644 index 00000000000..3cda71588c0 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/components/SkillsSection.tsx @@ -0,0 +1,94 @@ +/*! + * 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, HStack, Heading, Link, Stack, Text } from "@chakra-ui/react"; +import type { FC } from "react"; + +import { Tooltip } from "src/components/Tooltip"; +import { describeSkillSource, isLoaded } from "src/skill-source"; +import type { SkillInfo } from "src/types/model"; + +const SkillSourceText: FC<{ skill: SkillInfo }> = ({ skill }) => { + const { location, path, commit, href } = describeSkillSource(skill.source); + const text = [location, path, commit].filter((part) => part !== undefined).join(" · "); + + return href === undefined ? ( + + {text} + + ) : ( + + {text} + + ); +}; + +export const SkillsSection: FC<{ skills: SkillInfo[] }> = ({ skills }) => { + // Loaded skills first: they are what the run actually used. + const ordered = [...skills].sort((a, b) => Number(isLoaded(b.calls)) - Number(isLoaded(a.calls))); + const loadedCount = skills.filter((skill) => isLoaded(skill.calls)).length; + + return ( + + + Skills + + {loadedCount} of {skills.length} loaded + + + + {ordered.map((skill) => { + const loaded = isLoaded(skill.calls); + const calls = Object.entries(skill.calls); + + return ( + + + {skill.description ? {skill.description} : undefined} + + {calls.length === 0 + ? "Available, not used" + : calls.map(([tool, count]) => `${tool} ×${count}`).join(" · ")} + + + } + portalled + > + + {skill.name} + + + + + ); + })} + + + ); +}; 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 index c4b4e4df85a..b3a2fdafa29 100644 --- 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 @@ -21,11 +21,12 @@ import { useEffect, useState } from "react"; import { createModelApi } from "src/model-api"; import { regionOf, type TaskRegion } from "src/region"; -import type { UsageInfo } from "src/types/model"; +import type { SkillInfo, UsageInfo } from "src/types/model"; interface UseModelInfoReturn { modelName: string | null; usage: UsageInfo | null; + skills: SkillInfo[] | null; loading: boolean; error: string | null; } @@ -44,6 +45,7 @@ export function useModelInfo( ): UseModelInfoReturn { const [modelName, setModelName] = useState(null); const [usage, setUsage] = useState(null); + const [skills, setSkills] = useState(null); const [loading, setLoading] = useState(true); const [error, setError] = useState(null); @@ -63,19 +65,22 @@ export function useModelInfo( mapIndex, regionOf({ region_id: regionId, region_index: regionIndex }), ); - // allSettled, not all: one XCom missing or erroring must not hide a model name or usage - // the other call did resolve. - void Promise.allSettled([api.fetchModelName(), api.fetchUsage()]).then(([nameResult, usageResult]) => { - if (!current) return; - setModelName(nameResult.status === "fulfilled" ? nameResult.value : null); - setUsage(usageResult.status === "fulfilled" ? usageResult.value : null); - setError(describeFailure(nameResult) ?? describeFailure(usageResult)); - setLoading(false); - }); + // allSettled, not all: one XCom missing or erroring must not hide what the other calls + // did resolve. + void Promise.allSettled([api.fetchModelName(), api.fetchUsage(), api.fetchSkills()]).then( + ([nameResult, usageResult, skillsResult]) => { + if (!current) return; + setModelName(nameResult.status === "fulfilled" ? nameResult.value : null); + setUsage(usageResult.status === "fulfilled" ? usageResult.value : null); + setSkills(skillsResult.status === "fulfilled" ? skillsResult.value : null); + setError(describeFailure(nameResult) ?? describeFailure(usageResult) ?? describeFailure(skillsResult)); + setLoading(false); + }, + ); return () => { current = false; }; }, [dagId, runId, taskId, mapIndex, regionId, regionIndex]); - return { modelName, usage, loading, error }; + return { modelName, usage, skills, 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 index d1b1af22ba6..d5046717cf1 100644 --- 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 @@ -18,7 +18,7 @@ */ import { regionQuery, type TaskRegion } from "./region.ts"; -import type { UsageInfo } from "src/types/model"; +import type { SkillInfo, UsageInfo } from "src/types/model"; function getApiBase(): string { if (typeof document === "undefined") return "/api/v2"; @@ -33,6 +33,8 @@ 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__"; +// Matches airflow.providers.common.ai.utils.logging.SKILLS_XCOM_KEY. +const SKILLS_XCOM_KEY = "__AIRFLOW__COMMON_AI_SKILLS__"; /** Read one XCom entry; returns `null` when the task hasn't pushed that key (404). */ async function fetchXComValue( @@ -88,5 +90,10 @@ export function createModelApi( // let every stat box below render "undefined". return isUsageInfo(value) ? value : null; }, + fetchSkills: async () => { + const value = await fetchXComValue(dagId, runId, taskId, mapIndex, SKILLS_XCOM_KEY, region); + // Same reference-storing backend caveat as `usage` above. + return Array.isArray(value) ? (value as SkillInfo[]) : null; + }, }; } diff --git a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/skill-source.ts b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/skill-source.ts new file mode 100644 index 00000000000..9ee0d09bae8 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/src/skill-source.ts @@ -0,0 +1,54 @@ +/*! + * 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 { SkillSource } from "src/types/model"; + +/** An HTTPS or SSH Git remote, split into its host and repository path. */ +const GIT_REMOTE = /^(?:https:\/\/|ssh:\/\/git@|git@)([^/:]+)[/:](.+?)(?:\.git)?\/?$/; + +/** + * How a skill's source reads in the tab: the repository or local path, the path inside the + * repository, the short commit, and a link to that commit's folder when the host is GitHub or + * GitLab. The link pins the commit, not the branch, so it shows the files the run actually used. + */ +export function describeSkillSource(source: SkillSource): { + location: string; + path?: string; + commit?: string; + href?: string; +} { + if (source.repo_url === undefined) return { location: source.path }; + const commit = source.commit ?? undefined; + const match = GIT_REMOTE.exec(source.repo_url); + const location = match ? `${match[1]}/${match[2]}` : source.repo_url; + const path = source.path === "." ? undefined : source.path; + if (!match || commit === undefined) return { location, path, commit: commit?.slice(0, 7) }; + const [, host, repo] = match; + const treeSegment = host!.includes("gitlab") ? "-/tree" : host!.includes("github") ? "tree" : undefined; + const href = + treeSegment === undefined + ? undefined + : `https://${host}/${repo}/${treeSegment}/${commit}${path === undefined ? "" : `/${path}`}`; + return { location, path, commit: commit.slice(0, 7), href }; +} + +/** Whether the agent loaded the skill, rather than only seeing it listed. */ +export function isLoaded(calls: Record): boolean { + return (calls.load_skill ?? 0) > 0; +} 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 index a15438f8555..eff59d1b12c 100644 --- 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 @@ -26,3 +26,20 @@ export interface UsageInfo { tool_calls: number; cost: string | null; } + +/** Where a skill was loaded from: a local path, or a path inside a Git repository at a commit. */ +export interface SkillSource { + path: string; + repo_url?: string; + branch?: string | null; + commit?: string | null; +} + +/** One entry of the skills XCom pushed by AgentSkillsToolset (see toolsets/skills.py). */ +export interface SkillInfo { + name: string; + description: string | null; + source: SkillSource; + /** How many times the agent called each skill tool for this skill, e.g. `{ load_skill: 1 }`. */ + calls: Record; +} diff --git a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/model.test.mjs b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/model.test.mjs index fde44e545c6..0913a8d3dd3 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/model.test.mjs +++ b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/model.test.mjs @@ -81,11 +81,11 @@ test("an older host without task instance context reads without a region", async await plugin.render(taskProps); - assert.equal(calls.length, 2); + assert.equal(calls.length, 3); for (const url of calls) assert.equal(url.searchParams.has("region_id"), false); }); -test("the selected loop pass reaches both XCom reads and a pass switch rereads", async (t) => { +test("the selected loop pass reaches every XCom read and a pass switch rereads", async (t) => { const calls = serveXComs(); const plugin = mountPlugin(t); @@ -97,8 +97,10 @@ test("the selected loop pass reaches both XCom reads and a pass switch rereads", [ ["__AIRFLOW__COMMON_AI_MODEL_NAME__", regionId, "2"], ["usage", regionId, "2"], + ["__AIRFLOW__COMMON_AI_SKILLS__", regionId, "2"], ["__AIRFLOW__COMMON_AI_MODEL_NAME__", regionId, "3"], ["usage", regionId, "3"], + ["__AIRFLOW__COMMON_AI_SKILLS__", regionId, "3"], ], ); assert.match(plugin.text(), /model-for-pass-3/); @@ -119,3 +121,41 @@ test("a slow read for the previous loop pass does not replace the current pass", assert.match(plugin.text(), /model-for-pass-3/); assert.doesNotMatch(plugin.text(), /model-for-pass-2/); }); + +test("the skills section lists every available skill with its source, loaded ones first", async (t) => { + const skills = [ + { + name: "unused", + description: "Never loaded.", + source: { path: "/opt/airflow/skills/unused" }, + calls: {}, + }, + { + name: "sql-review", + description: "Reviews SQL.", + source: { + repo_url: "https://github.com/org/skills.git", + path: "skills/sql-review", + branch: "main", + commit: "0123456789abcdef", + }, + calls: { load_skill: 1 }, + }, + ]; + mock.method(globalThis, "fetch", async (input) => { + const key = new URL(input, "http://localhost").pathname.split("/").at(-1); + const value = key === "__AIRFLOW__COMMON_AI_SKILLS__" ? skills : key === "usage" ? { requests: 1 } : "model"; + return new Response(JSON.stringify({ value }), { status: 200 }); + }); + const plugin = mountPlugin(t); + + await plugin.render(taskProps); + + const text = plugin.text(); + assert.match(text, /1 of 2 loaded/); + assert.ok(text.indexOf("sql-review") < text.indexOf("unused")); + assert.match(text, /github\.com\/org\/skills · skills\/sql-review · 0123456/); + assert.match(text, /\/opt\/airflow\/skills\/unused/); + const link = document.querySelector('a[href="https://github.com/org/skills/tree/0123456789abcdef/skills/sql-review"]'); + assert.ok(link); +}); diff --git a/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/skill-source.test.ts b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/skill-source.test.ts new file mode 100644 index 00000000000..9de5421a61c --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/skill-source.test.ts @@ -0,0 +1,69 @@ +/*! + * 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 assert from "node:assert/strict"; +import { test } from "node:test"; + +import { describeSkillSource, isLoaded } from "../src/skill-source.ts"; + +const commit = "0123456789abcdef"; + +test("a local skill reads as its path, with no link", () => { + assert.deepEqual(describeSkillSource({ path: "/opt/airflow/skills/alpha" }), { + location: "/opt/airflow/skills/alpha", + }); +}); + +for (const [repoUrl, href] of [ + ["https://github.com/org/skills.git", `https://github.com/org/skills/tree/${commit}/skills/alpha`], + ["git@github.com:org/skills.git", `https://github.com/org/skills/tree/${commit}/skills/alpha`], + ["ssh://git@gitlab.com/group/sub/skills", `https://gitlab.com/group/sub/skills/-/tree/${commit}/skills/alpha`], +]) { + test(`a Git skill from ${repoUrl} links the folder at the commit that ran`, () => { + assert.deepEqual(describeSkillSource({ repo_url: repoUrl, path: "skills/alpha", branch: "main", commit }), { + location: repoUrl.includes("gitlab") ? "gitlab.com/group/sub/skills" : "github.com/org/skills", + path: "skills/alpha", + commit: "0123456", + href, + }); + }); +} + +test("a skill at the repository root links the commit's root", () => { + const source = { repo_url: "https://github.com/org/skills", path: ".", commit }; + + assert.equal(describeSkillSource(source).href, `https://github.com/org/skills/tree/${commit}`); +}); + +test("a host with no known web UI shows the repository and commit without a link", () => { + const source = { repo_url: "https://git.example.com/org/skills.git", path: "skills/alpha", commit }; + + assert.deepEqual(describeSkillSource(source), { + location: "git.example.com/org/skills", + path: "skills/alpha", + commit: "0123456", + href: undefined, + }); +}); + +test("only a skill the agent loaded counts as loaded", () => { + assert.equal(isLoaded({ load_skill: 1 }), true); + assert.equal(isLoaded({ read_skill_resource: 2 }), false); + assert.equal(isLoaded({}), false); +}); diff --git a/providers/common/ai/src/airflow/providers/common/ai/skills.py b/providers/common/ai/src/airflow/providers/common/ai/skills.py index 13c053313bc..cd71efe9f31 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/skills.py +++ b/providers/common/ai/src/airflow/providers/common/ai/skills.py @@ -97,9 +97,18 @@ class GitSkills: SkillSource = Union[str, "os.PathLike[str]", GitSkills] -def _clone_git(source: GitSkills) -> tuple[str, str]: +@dataclass(frozen=True) +class _SkillsLocation: + """A resolved skills directory, with the Git source and commit it was cloned from.""" + + directory: str + git: GitSkills | None = None + commit: str | None = None + + +def _clone_git(source: GitSkills) -> tuple[str, str, str]: """ - Clone *source* into a fresh temp dir; return (skills_dir, temp_dir_to_remove). + Clone *source* into a fresh temp dir; return (skills_dir, temp_dir_to_remove, commit). When ``conn_id`` is set, credentials come from the Airflow ``git`` connection via ``GitHook`` (HTTPS token or SSH key). The token is stripped from the @@ -162,6 +171,7 @@ def _clone_git(source: GitSkills) -> tuple[str, str]: # Strip any embedded credential from .git/config so a skill script in the # checkout cannot read it back out of the temporary clone. repo.remote("origin").set_url(source.repo_url) + commit = repo.head.commit.hexsha except Exception as exc: shutil.rmtree(temp_dir, ignore_errors=True) # GitPython errors embed the failing command, which may contain the @@ -175,26 +185,26 @@ def _clone_git(source: GitSkills) -> tuple[str, str]: raise skills_dir = os.path.join(temp_dir, source.path) if source.path else temp_dir - return skills_dir, temp_dir + return skills_dir, temp_dir, commit -def _materialize_skills(sources: list[SkillSource]) -> tuple[list[str], Callable[[], None]]: +def _materialize_skills(sources: list[SkillSource]) -> tuple[list[_SkillsLocation], Callable[[], None]]: """ - Resolve *sources* to local directories; return (directories, cleanup). + Resolve *sources* to local directories; return (locations, cleanup). ``cleanup`` removes any directories cloned for ``GitSkills`` sources; local directory sources are returned untouched and are not removed. """ - directories: list[str] = [] + locations: list[_SkillsLocation] = [] temp_dirs: list[str] = [] try: for source in sources: if isinstance(source, GitSkills): - skills_dir, temp_dir = _clone_git(source) - directories.append(skills_dir) + skills_dir, temp_dir, commit = _clone_git(source) + locations.append(_SkillsLocation(skills_dir, git=source, commit=commit)) temp_dirs.append(temp_dir) elif isinstance(source, (str, os.PathLike)): - directories.append(os.fspath(source)) + locations.append(_SkillsLocation(os.fspath(source))) else: raise TypeError( f"Unsupported skill source {type(source).__name__!r}; expected a path or GitSkills." @@ -209,7 +219,7 @@ def _materialize_skills(sources: list[SkillSource]) -> tuple[list[str], Callable for temp_dir in temp_dirs: shutil.rmtree(temp_dir, ignore_errors=True) - return directories, cleanup + return locations, cleanup @contextmanager @@ -222,8 +232,8 @@ def resolve_skills(sources: list[SkillSource]) -> Iterator[list[str]]: ``AgentSkills(skills=...)``). Cloned repositories are removed on exit, so use the returned directories inside the ``with`` block. """ - directories, cleanup = _materialize_skills(sources) + locations, cleanup = _materialize_skills(sources) try: - yield directories + yield [location.directory for location in locations] finally: cleanup() diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/skills.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/skills.py index 325abba642c..caa47f616fa 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/skills.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/skills.py @@ -31,11 +31,19 @@ removed when the toolset context exits. from __future__ import annotations import dataclasses +import logging +import os +from collections import Counter, defaultdict from importlib import metadata +from pathlib import Path, PurePosixPath from typing import TYPE_CHECKING, Any -from airflow.providers.common.ai.skills import SkillSource, _materialize_skills +from pydantic_ai.messages import ToolCallPart, ToolReturnPart + +from airflow.providers.common.ai.skills import SkillSource, _materialize_skills, _SkillsLocation +from airflow.providers.common.ai.utils.logging import SKILLS_XCOM_KEY from airflow.providers.common.ai.utils.toolset_base import validate_max_retries +from airflow.providers.common.compat.sdk import get_current_context # pydantic-ai-skills 2.0 removed SkillsToolset, leaving only SkillsCapability, so the # ``skills`` extra caps it below 2. @@ -50,9 +58,53 @@ if TYPE_CHECKING: from collections.abc import Callable, Sequence from pydantic_ai._run_context import RunContext - from pydantic_ai.messages import InstructionPart + from pydantic_ai.messages import InstructionPart, ModelMessage from pydantic_ai.toolsets.abstract import ToolsetTool +log = logging.getLogger(__name__) + +# The pydantic-ai-skills tools that act on one skill, named by their ``skill_name`` argument. +_PER_SKILL_TOOLS = frozenset({"load_skill", "read_skill_resource", "run_skill_script"}) + + +def _describe_skill_source(skill_uri: str, locations: list[_SkillsLocation]) -> dict[str, Any]: + """Name where a discovered skill came from: its local path, or its repository, path and commit.""" + skill_dir = os.path.realpath(skill_uri) + # Nested sources can both contain the skill; the deepest one is where it came from. + owner, owner_root = None, "" + for location in locations: + root = os.path.realpath(location.directory) + if os.path.commonpath([skill_dir, root]) == root and len(root) > len(owner_root): + owner, owner_root = location, root + if owner is None or owner.git is None: + return {"path": skill_dir} + return { + # GitSkills rejects a repo_url carrying credentials, so it is safe to publish. + "repo_url": owner.git.repo_url, + "path": PurePosixPath(owner.git.path, Path(skill_dir).relative_to(owner_root).as_posix()).as_posix(), + "branch": owner.git.branch, + "commit": owner.commit, + } + + +def _count_skill_calls(messages: Sequence[ModelMessage]) -> defaultdict[str, Counter[str]]: + """Count, per skill, the calls to each per-skill tool that returned a result.""" + calls: dict[str, ToolCallPart] = {} + counts: defaultdict[str, Counter[str]] = defaultdict(Counter) + for message in messages: + for part in message.parts: + if isinstance(part, ToolCallPart): + calls[part.tool_call_id] = part + elif isinstance(part, ToolReturnPart) and (call := calls.get(part.tool_call_id)): + # Two skills toolsets on one agent need ``prefixed()`` to keep their tool names apart. + tool = next( + (t for t in _PER_SKILL_TOOLS if call.tool_name == t or call.tool_name.endswith(f"_{t}")), + None, + ) + if tool is not None and isinstance(skill_name := call.args_as_dict().get("skill_name"), str): + counts[skill_name][tool] += 1 + return counts + class AgentSkillsToolset(AbstractToolset): """ @@ -101,6 +153,10 @@ class AgentSkillsToolset(AbstractToolset): self._max_retries = validate_max_retries(max_retries) self._inner: Any = None self._cleanup: Callable[[], None] | None = None + self._locations: list[_SkillsLocation] = [] + # The run's messages, prior history included. Calls made before an approval pause or + # replayed on a durable retry never reach call_tool, so the skills XCom counts these. + self._messages: Sequence[ModelMessage] = () @property def id(self) -> str | None: @@ -136,10 +192,11 @@ class AgentSkillsToolset(AbstractToolset): ) SkillsToolset = pydantic_ai_skills.SkillsToolset - directories, cleanup = _materialize_skills(self._sources) + locations, cleanup = _materialize_skills(self._sources) self._cleanup = cleanup + self._locations = locations try: - kwargs: dict[str, Any] = {"directories": directories} + kwargs: dict[str, Any] = {"directories": [location.directory for location in locations]} if self._exclude_tools: kwargs["exclude_tools"] = self._exclude_tools if self._exclude_resources: @@ -156,6 +213,7 @@ class AgentSkillsToolset(AbstractToolset): async def __aexit__(self, *args: Any) -> bool | None: try: if self._inner is not None: + self._push_skills_xcom() return await self._inner.__aexit__(*args) return None finally: @@ -164,6 +222,36 @@ class AgentSkillsToolset(AbstractToolset): self._inner = None self._cleanup = None + def _push_skills_xcom(self) -> None: + """Publish the skills offered to the agent and the skill tools it called, for the AI Model tab.""" + try: + context = get_current_context() + except RuntimeError: + # Used outside an Airflow task: there is no task instance to publish on. + return + if not context["ti"].task.do_xcom_push: + return + # Telemetry only: a failure here must not fail the run or skip the clean-up after it. + try: + counts = _count_skill_calls(self._messages) + skills = { + name: { + "name": name, + "description": skill.description, + "source": _describe_skill_source(skill.uri, self._locations), + "calls": dict(counts.get(name, {})), + } + for name, skill in self._inner.skills.items() + } + # Another skills toolset on the agent, or this run before it paused for approval, may + # have published already. XComs from an earlier try are cleared before this one starts. + ti = context["ti"] + stored = ti.xcom_pull(key=SKILLS_XCOM_KEY, task_ids=ti.task_id, map_indexes=ti.map_index) or [] + merged = {entry["name"]: entry for entry in stored} | skills + ti.xcom_push(key=SKILLS_XCOM_KEY, value=[merged[name] for name in sorted(merged)]) + except Exception: + log.warning("Failed to push the skills XCom", exc_info=True) + def _require_inner(self) -> Any: if self._inner is None: raise RuntimeError( @@ -176,6 +264,7 @@ class AgentSkillsToolset(AbstractToolset): # pydantic-ai-skills fixes its tools' budget at one correction, whatever the agent's # retries say; resolve it the way the provider's other toolsets do instead. max_retries = ctx.max_retries if self._max_retries is None else self._max_retries + self._messages = ctx.messages tools = await self._require_inner().get_tools(ctx) return {name: dataclasses.replace(tool, max_retries=max_retries) for name, tool in tools.items()} diff --git a/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py b/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py index 7cd0ecaafff..74cb83e6506 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py +++ b/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py @@ -96,6 +96,10 @@ def _log_cache_and_cost(logger: Logger | logging.Logger, usage: RunUsage) -> Non # name" entry in docs/observability.rst for exactly which operators publish it. MODEL_NAME_XCOM_KEY = "__AIRFLOW__COMMON_AI_MODEL_NAME__" +# Skills an AgentSkillsToolset offered the agent, where each came from, and which skill tools +# the agent called for it. Read by the AI Model tab. See the "Skills" entry in docs/observability.rst. +SKILLS_XCOM_KEY = "__AIRFLOW__COMMON_AI_SKILLS__" + def format_usage_for_xcom(usage: RunUsage) -> dict[str, Any]: """Build the XCom ``usage`` payload -- shared by the success and failure paths.""" diff --git a/providers/common/ai/tests/unit/common/ai/test_skills.py b/providers/common/ai/tests/unit/common/ai/test_skills.py index cbaddb2787e..267bb246463 100644 --- a/providers/common/ai/tests/unit/common/ai/test_skills.py +++ b/providers/common/ai/tests/unit/common/ai/test_skills.py @@ -23,7 +23,7 @@ from unittest.mock import patch import pytest -from airflow.providers.common.ai.skills import GitSkills, resolve_skills +from airflow.providers.common.ai.skills import GitSkills, _materialize_skills, _SkillsLocation, resolve_skills pytest.importorskip("git") pytest.importorskip("airflow.providers.git.hooks.git") @@ -124,6 +124,19 @@ class TestResolveGitWithHook: mock_githook.assert_not_called() assert mock_repo.clone_from.call_args.args[0] == "https://github.com/org/repo" + @patch("git.Repo") + def test_records_the_cloned_commit(self, mock_repo): + mock_repo.clone_from.return_value.head.commit.hexsha = "0123abc" + source = GitSkills(repo_url="https://github.com/org/repo", path="skills") + + locations, cleanup = _materialize_skills([source, "/local/skills"]) + cleanup() + + assert locations[0].git is source + assert locations[0].commit == "0123abc" + assert locations[0].directory.endswith(os.sep + "skills") + assert locations[1] == _SkillsLocation("/local/skills") + @patch("airflow.providers.git.hooks.git.GitHook") @patch("git.Repo") def test_failed_clone_removes_temp_dir(self, mock_repo, mock_githook, tmp_path): diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_skills.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_skills.py index be277d0d936..6f6bd9a7ec5 100644 --- a/providers/common/ai/tests/unit/common/ai/toolsets/test_skills.py +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_skills.py @@ -23,6 +23,7 @@ pytest-asyncio mode. from __future__ import annotations import asyncio +import os import sys import types from unittest.mock import MagicMock, patch @@ -33,8 +34,9 @@ from pydantic_ai.exceptions import UnexpectedModelBehavior from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart from pydantic_ai.models.function import FunctionModel -from airflow.providers.common.ai.skills import GitSkills -from airflow.providers.common.ai.toolsets.skills import AgentSkillsToolset +from airflow.providers.common.ai.skills import GitSkills, _SkillsLocation +from airflow.providers.common.ai.toolsets.skills import AgentSkillsToolset, _describe_skill_source +from airflow.providers.common.ai.utils.logging import SKILLS_XCOM_KEY pytest.importorskip("pydantic_ai_skills") @@ -56,12 +58,13 @@ class _FakeInner: return None -def _write_skill(directory): - skill_dir = directory / "demo-skill" +def _write_skill(directory, name="demo-skill"): + skill_dir = directory / name skill_dir.mkdir(parents=True) (skill_dir / "SKILL.md").write_text( - "---\nname: demo-skill\ndescription: A demo skill for tests.\n---\n\n# Demo\n" + f"---\nname: {name}\ndescription: A {name} skill for tests.\n---\n\n# Demo\n" ) + return skill_dir class TestConstruction: @@ -104,7 +107,7 @@ class TestLifecycle: toolset = AgentSkillsToolset(sources=["/x"], exclude_tools={"run_skill_script"}) with patch( "airflow.providers.common.ai.toolsets.skills._materialize_skills", - return_value=(["/x"], lambda: None), + return_value=([_SkillsLocation("/x")], lambda: None), ): with patch("pydantic_ai_skills.SkillsToolset", fake_skillstoolset): asyncio.run(_enter_exit(toolset)) @@ -123,7 +126,7 @@ class TestLifecycle: with patch( "airflow.providers.common.ai.toolsets.skills._materialize_skills", autospec=True, - return_value=(["/x"], lambda: None), + return_value=([_SkillsLocation("/x")], lambda: None), ): with patch("pydantic_ai_skills.SkillsToolset", fake_skillstoolset): # noqa: spec asyncio.run(_enter_exit(toolset)) @@ -143,7 +146,7 @@ class TestLifecycle: with patch( "airflow.providers.common.ai.toolsets.skills._materialize_skills", autospec=True, - return_value=(["/x"], lambda: None), + return_value=([_SkillsLocation("/x")], lambda: None), ): with patch("pydantic_ai_skills.SkillsToolset", fake_skillstoolset): # noqa: spec asyncio.run(_enter_exit(toolset)) @@ -214,7 +217,7 @@ class TestCleanup: toolset = AgentSkillsToolset(sources=["/x"]) with patch( "airflow.providers.common.ai.toolsets.skills._materialize_skills", - return_value=(["/x"], cleanup), + return_value=([_SkillsLocation("/x")], cleanup), ): with patch("pydantic_ai_skills.SkillsToolset", boom): with pytest.raises(RuntimeError, match="inner build failed"): @@ -250,3 +253,195 @@ class TestMissingExtra: async def _enter_exit(toolset): await toolset.__aenter__() await toolset.__aexit__(None, None, None) + + +def _calling(*tool_calls): + """A model that makes each (tool, args) call in turn, then answers "done".""" + calls = iter(tool_calls) + + def model(messages, info): + if (call := next(calls, None)) is None: + return ModelResponse(parts=[TextPart("done")]) + return ModelResponse(parts=[ToolCallPart(*call)]) + + return FunctionModel(model) + + +_GET_CONTEXT = "airflow.providers.common.ai.toolsets.skills.get_current_context" + + +class _FakeTaskInstance: + """Keeps what is pushed, so a later pull in the same task instance reads it back.""" + + task_id = "agent" + map_index = -1 + + def __init__(self, stored=None): + self.task = types.SimpleNamespace(do_xcom_push=True) + self.xcoms = {} if stored is None else {SKILLS_XCOM_KEY: stored} + + def xcom_pull(self, key, task_ids, map_indexes): + assert (task_ids, map_indexes) == (self.task_id, self.map_index) + return self.xcoms.get(key) + + def xcom_push(self, key, value): + self.xcoms[key] = value + + +def _skill_entry(name, path, calls): + return { + "name": name, + "description": f"A {name} skill for tests.", + "source": {"path": os.path.realpath(path)}, + "calls": calls, + } + + +class TestSkillsXCom: + @patch(_GET_CONTEXT, autospec=True) + def test_pushes_available_skills_and_the_calls_made_for_each(self, mock_get_context, tmp_path): + alpha = _write_skill(tmp_path, "alpha") + (alpha / "notes.md").write_text("notes") + beta = _write_skill(tmp_path, "beta") + ti = _FakeTaskInstance() + mock_get_context.return_value = {"ti": ti} + model = _calling( + ("load_skill", {"skill_name": "alpha"}), + ("read_skill_resource", {"skill_name": "alpha", "resource_name": "notes.md"}), + # Not a skill: the failed lookup is not counted against any skill. + ("load_skill", {"skill_name": "missing"}), + ) + + Agent(model, toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])]).run_sync("go") + + assert ti.xcoms[SKILLS_XCOM_KEY] == [ + _skill_entry("alpha", alpha, {"load_skill": 1, "read_skill_resource": 1}), + _skill_entry("beta", beta, {}), + ] + + @patch(_GET_CONTEXT, autospec=True) + def test_counts_calls_from_the_message_history(self, mock_get_context, tmp_path): + # A run resumed after an approval pause or replayed on a durable retry starts in a new + # process, with its earlier calls given back as history rather than run again. + alpha = _write_skill(tmp_path, "alpha") + ti = _FakeTaskInstance() + mock_get_context.return_value = {"ti": ti} + earlier = Agent( + _calling(("load_skill", {"skill_name": "alpha"})), + toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])], + ).run_sync("go") + + Agent( + _calling(("load_skill", {"skill_name": "alpha"})), + toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])], + ).run_sync("go on", message_history=earlier.all_messages()) + + assert ti.xcoms[SKILLS_XCOM_KEY] == [_skill_entry("alpha", alpha, {"load_skill": 2})] + + @patch(_GET_CONTEXT, autospec=True) + def test_replaces_its_own_stored_skills_and_keeps_the_others(self, mock_get_context, tmp_path): + alpha = _write_skill(tmp_path, "alpha") + other = _skill_entry("other", tmp_path / "elsewhere", {"load_skill": 3}) + ti = _FakeTaskInstance(stored=[_skill_entry("alpha", alpha, {"load_skill": 1}), other]) + mock_get_context.return_value = {"ti": ti} + + Agent(_calling(), toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])]).run_sync("go") + + assert ti.xcoms[SKILLS_XCOM_KEY] == [_skill_entry("alpha", alpha, {}), other] + + @patch(_GET_CONTEXT, autospec=True) + def test_two_prefixed_toolsets_each_count_their_own_skills(self, mock_get_context, tmp_path): + alpha = _write_skill(tmp_path / "one", "alpha") + beta = _write_skill(tmp_path / "two", "beta") + ti = _FakeTaskInstance() + mock_get_context.return_value = {"ti": ti} + toolsets = [ + AgentSkillsToolset(sources=[str(tmp_path / "one")]).prefixed("one"), + AgentSkillsToolset(sources=[str(tmp_path / "two")]).prefixed("two"), + ] + model = _calling( + ("one_load_skill", {"skill_name": "alpha"}), + ("two_load_skill", {"skill_name": "beta"}), + ) + + Agent(model, toolsets=toolsets).run_sync("go") + + assert ti.xcoms[SKILLS_XCOM_KEY] == [ + _skill_entry("alpha", alpha, {"load_skill": 1}), + _skill_entry("beta", beta, {"load_skill": 1}), + ] + + @pytest.mark.parametrize( + "context", + [ + pytest.param(RuntimeError("no task"), id="outside_a_task"), + pytest.param({"ti": MagicMock(**{"task.do_xcom_push": False})}, id="do_xcom_push_off"), + ], + ) + @patch(_GET_CONTEXT, autospec=True) + def test_nothing_pushed_without_a_task_that_pushes_xcom(self, mock_get_context, tmp_path, context): + _write_skill(tmp_path, "alpha") + if isinstance(context, Exception): + mock_get_context.side_effect = context + else: + mock_get_context.return_value = context + model = _calling(("load_skill", {"skill_name": "alpha"})) + + assert ( + Agent(model, toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])]).run_sync("go").output + == "done" + ) + if not isinstance(context, Exception): + context["ti"].xcom_push.assert_not_called() + + @patch(_GET_CONTEXT, autospec=True) + def test_failed_push_does_not_fail_the_run_or_skip_cleanup(self, mock_get_context, tmp_path): + _write_skill(tmp_path, "alpha") + ti = MagicMock() + ti.task.do_xcom_push = True + ti.xcom_pull.return_value = None + ti.xcom_push.side_effect = RuntimeError("xcom backend down") + mock_get_context.return_value = {"ti": ti} + cleanup = MagicMock() + toolset = AgentSkillsToolset(sources=[str(tmp_path)]) + + with patch( + "airflow.providers.common.ai.toolsets.skills._materialize_skills", + return_value=([_SkillsLocation(str(tmp_path))], cleanup), + ): + result = Agent(_calling(), toolsets=[toolset]).run_sync("go") + + assert result.output == "done" + cleanup.assert_called_once() + + +class TestDescribeSkillSource: + @pytest.mark.parametrize( + ("repo_path", "expected_path"), + [ + pytest.param("skills", "skills/alpha", id="sub_path"), + pytest.param("", "alpha", id="repo_root"), + ], + ) + def test_git_skill_names_repository_path_and_commit(self, tmp_path, repo_path, expected_path): + skills_dir = tmp_path / "clone" / repo_path + skill_dir = _write_skill(skills_dir, "alpha") + source = GitSkills(repo_url="https://github.com/org/repo", path=repo_path, branch="main") + locations = [_SkillsLocation(str(skills_dir), git=source, commit="0123abc")] + + assert _describe_skill_source(str(skill_dir), locations) == { + "repo_url": "https://github.com/org/repo", + "path": expected_path, + "branch": "main", + "commit": "0123abc", + } + + def test_nested_source_wins_over_the_one_containing_it(self, tmp_path): + skill_dir = _write_skill(tmp_path / "clone", "alpha") + source = GitSkills(repo_url="https://github.com/org/repo") + locations = [ + _SkillsLocation(str(tmp_path)), + _SkillsLocation(str(tmp_path / "clone"), git=source, commit="0123abc"), + ] + + assert _describe_skill_source(str(skill_dir), locations)["repo_url"] == "https://github.com/org/repo"