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"