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.
This commit is contained in:
Pierre Jeambrun authored and GitHub committed 2026-11-03 13:19:51 -05:00
1 parent eea312fc7c
commit be2829b917
15 files changed
+668 -43

No files matched your search

+6 -1
View File
@@ -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
@@ -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 <howto/observability: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
----------
@@ -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<ModelPanelProps> = ({ 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 <ModelLoading />;
}
if (!modelName && !usage) {
if (!modelName && !usage && !skills) {
return <NoModelInfo error={error} />;
}
@@ -141,6 +142,8 @@ export const ModelPanel: FC<ModelPanelProps> = ({ dagId, runId, taskId, mapIndex
<StatBox label="Cost" tooltip={COST_TOOLTIP} value={formatCost(usage.cost)} />
</SimpleGrid>
)}
{skills && skills.length > 0 && <SkillsSection skills={skills} />}
</Box>
);
};
@@ -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 color="fg.muted" fontFamily="mono" fontSize="xs" wordBreak="break-all">
{text}
</Text>
) : (
<Link color="fg.muted" fontFamily="mono" fontSize="xs" href={href} rel="noreferrer" target="_blank">
{text}
</Link>
);
};
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 (
<Box data-testid="skills-section" maxW="960px" mt={6}>
<HStack gap={2} mb={3}>
<Heading size="sm">Skills</Heading>
<Text color="fg.muted" fontSize="sm">
{loadedCount} of {skills.length} loaded
</Text>
</HStack>
<Stack gap={2}>
{ordered.map((skill) => {
const loaded = isLoaded(skill.calls);
const calls = Object.entries(skill.calls);
return (
<HStack align="center" gap={3} key={skill.name}>
<Tooltip
content={
<Stack gap={1}>
{skill.description ? <Text>{skill.description}</Text> : undefined}
<Text>
{calls.length === 0
? "Available, not used"
: calls.map(([tool, count]) => `${tool} ×${count}`).join(" · ")}
</Text>
</Stack>
}
portalled
>
<Badge
borderRadius="full"
colorPalette={loaded ? "brand" : "gray"}
flexShrink={0}
fontSize="sm"
px={3}
py={1}
variant={loaded ? "solid" : "outline"}
>
{skill.name}
</Badge>
</Tooltip>
<SkillSourceText skill={skill} />
</HStack>
);
})}
</Stack>
</Box>
);
};
@@ -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<string | null>(null);
const [usage, setUsage] = useState<UsageInfo | null>(null);
const [skills, setSkills] = useState<SkillInfo[] | null>(null);
const [loading, setLoading] = useState(true);
const [error, setError] = useState<string | null>(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 };
}
@@ -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<T>(
@@ -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<unknown>(dagId, runId, taskId, mapIndex, SKILLS_XCOM_KEY, region);
// Same reference-storing backend caveat as `usage` above.
return Array.isArray(value) ? (value as SkillInfo[]) : null;
},
};
}
@@ -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<string, number>): boolean {
return (calls.load_skill ?? 0) > 0;
}
@@ -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<string, number>;
}
@@ -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);
});
@@ -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);
});
@@ -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()
@@ -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()}
@@ -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."""
@@ -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):
@@ -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"