This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 92fa9b36b09 Fix skill call counts and source links in the AI Model tab
(#76541)
92fa9b36b09 is described below
commit 92fa9b36b09a4af36c96133e1db47f07d89eba11
Author: Pierre Jeambrun <[email protected]>
AuthorDate: Fri Nov 6 11:06:49 2026 +0100
Fix skill call counts and source links in the AI Model tab (#76541)
* Fix skill call counts and source links in the AI Model tab
Follow-up to the review of #76309. The skills section could show calls that
never
ran (denied or interrupted ones, or ones made by the task that seeded the
message
history), lose every count when the agent was entered with `async with
agent:`,
fail the task for a repository with no commits, break links for remotes on a
custom port, hide itself when the sources held no skills, and keep its
details
out of reach of the keyboard.
* Check the skills XCom for null explicitly and write the apostrophe as is
* Go back to a truthiness check on the skills XCom
* Count skill calls by agent run instead of by history position
The history-length boundary lost a skill call approved after a pause, since
its call
sits in the resumed run's history while its return is in the run itself; it
also
shifted when a history processor trimmed the transcript. Two prefixed
toolsets with
the same skill name each counted the same calls, so adding their publishes
doubled
them.
* Publish skills from a toolset built per run by a factory
pydantic-ai's DynamicToolset, which AgentOperator also wraps a toolset
factory in,
never calls for_run on the toolset it builds, so the run id set there left
such a
toolset unable to publish; get_tools sees every run. A run that made no
call to a
skill no longer adds an empty entry, so the XCom does not grow with every
resume.
* Note that a toolset factory rebuilds the skills toolset on every step
A toolset function is wrapped in DynamicToolset, which calls it again on
each model step by default, so the repository is cloned and the skills XCom
pushed once per step rather than once per run.
---
providers/common/ai/docs/observability.rst | 3 +-
providers/common/ai/docs/toolsets/skills.rst | 19 +-
.../ai/plugins/www/src/components/ModelPanel.tsx | 2 +-
.../plugins/www/src/components/SkillsSection.tsx | 13 +-
.../common/ai/plugins/www/src/skill-source.ts | 43 +++-
.../common/ai/plugins/www/tests/model.test.mjs | 36 ++++
.../ai/plugins/www/tests/skill-source.test.ts | 27 ++-
.../ai/src/airflow/providers/common/ai/skills.py | 5 +-
.../airflow/providers/common/ai/toolsets/skills.py | 61 ++++--
.../common/ai/tests/unit/common/ai/test_skills.py | 11 +
.../tests/unit/common/ai/toolsets/test_skills.py | 232 +++++++++++++++++----
11 files changed, 367 insertions(+), 85 deletions(-)
diff --git a/providers/common/ai/docs/observability.rst
b/providers/common/ai/docs/observability.rst
index 8a2d537c73b..29e2206198d 100644
--- a/providers/common/ai/docs/observability.rst
+++ b/providers/common/ai/docs/observability.rst
@@ -122,7 +122,8 @@ How it works
``deserialize_value``: the model name badge would show the reference string
(e.g. an
``s3://...`` path) instead of the model, and ``usage`` would come back as
that same
string rather than an object -- the tab treats a non-object ``usage`` as
absent and
- hides the stat boxes instead of rendering ``undefined`` in each one.
+ hides the stat boxes instead of rendering ``undefined`` in each one. The
skills XCom
+ comes back as that string too, so the skills section is hidden.
* **Content is off by default.** Only token counts, model id, latency, tool
names, and finish reason are recorded. Prompt and completion text is never
emitted unless you opt in (see below).
diff --git a/providers/common/ai/docs/toolsets/skills.rst
b/providers/common/ai/docs/toolsets/skills.rst
index e5261b008d5..4ef5088c8a7 100644
--- a/providers/common/ai/docs/toolsets/skills.rst
+++ b/providers/common/ai/docs/toolsets/skills.rst
@@ -140,10 +140,21 @@ 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``.
+GitLab, and says so when the sources held no skills. Calls made before a
tool-approval
+pause or a human-review regeneration still count, and so do calls a durable
retry replays.
+Only calls the tool answered count: a denied, retried, failed or interrupted
call does not,
+while a script that exits with an error still does, since ``run_skill_script``
returns its
+output. A ``message_history`` seeded from another task does not add that
task's calls, and
+calls that a history processor or compaction removes from the transcript are
not counted.
+``calls`` totals ``calls_by_run``, which keeps the calls of each agent run
that made any
+under its pydantic-ai run id.
+A run resumed from a tool-approval pause and a human-review regeneration each
clone the
+repository again, so the source shows the commit the latest run cloned. A
toolset built by
+a function, which the agent wraps in pydantic-ai's ``DynamicToolset``, is
built and cloned
+again on every model step unless you pass ``DynamicToolset(fn,
per_run_step=False)``.
+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 c092c3560cb..8e195278224 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
@@ -143,7 +143,7 @@ export const ModelPanel: FC<ModelPanelProps> = ({ dagId,
runId, taskId, mapIndex
</SimpleGrid>
)}
- {skills && skills.length > 0 && <SkillsSection skills={skills} />}
+ {skills && <SkillsSection skills={skills} />}
</Box>
);
};
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
index 3cda71588c0..5fd1ae203f0 100644
---
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
@@ -48,10 +48,17 @@ export const SkillsSection: FC<{ skills: SkillInfo[] }> =
({ skills }) => {
<Box data-testid="skills-section" maxW="960px" mt={6}>
<HStack gap={2} mb={3}>
<Heading size="sm">Skills</Heading>
+ {skills.length > 0 && (
+ <Text color="fg.muted" fontSize="sm">
+ {loadedCount} of {skills.length} loaded
+ </Text>
+ )}
+ </HStack>
+ {skills.length === 0 && (
<Text color="fg.muted" fontSize="sm">
- {loadedCount} of {skills.length} loaded
+ No skills found in the toolset's sources.
</Text>
- </HStack>
+ )}
<Stack gap={2}>
{ordered.map((skill) => {
const loaded = isLoaded(skill.calls);
@@ -73,6 +80,8 @@ export const SkillsSection: FC<{ skills: SkillInfo[] }> = ({
skills }) => {
portalled
>
<Badge
+ aria-label={`About the ${skill.name} skill`}
+ as="button"
borderRadius="full"
colorPalette={loaded ? "brand" : "gray"}
flexShrink={0}
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
index 9ee0d09bae8..28b2254be13 100644
---
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
@@ -19,8 +19,37 @@
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)?\/?$/;
+/** The scp-like `git@host:path` SSH remote, which `new URL()` cannot parse. */
+const SCP_REMOTE = /^git@([^/:]+):(.+)$/;
+
+/** Split an HTTPS or SSH Git remote into the host its web UI is served on and
the repository path. */
+function parseGitRemote(repoUrl: string): { host: string; repo: string } |
undefined {
+ let host: string;
+ let path: string;
+ const scp = SCP_REMOTE.exec(repoUrl);
+ if (scp) {
+ [, host = "", path = ""] = scp;
+ } else {
+ let url: URL;
+ try {
+ url = new URL(repoUrl);
+ } catch {
+ return undefined;
+ }
+ if (url.protocol === "https:") {
+ // Keeps a port, as in https://gitlab.corp.com:8443/team/skills.
+ host = url.host;
+ } else if (url.protocol === "ssh:") {
+ // An SSH port is not the web UI's.
+ host = url.hostname;
+ } else {
+ return undefined;
+ }
+ path = url.pathname.slice(1);
+ }
+ const repo = path.replace(/\/$/, "").replace(/\.git$/, "");
+ return host && repo ? { host, repo } : undefined;
+}
/**
* How a skill's source reads in the tab: the repository or local path, the
path inside the
@@ -35,12 +64,12 @@ export function describeSkillSource(source: SkillSource): {
} {
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 remote = parseGitRemote(source.repo_url);
+ const location = remote ? `${remote.host}/${remote.repo}` : 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;
+ if (!remote || commit === undefined) return { location, path, commit:
commit?.slice(0, 7) };
+ const { host, repo } = remote;
+ const treeSegment = host.includes("gitlab") ? "-/tree" :
host.includes("github") ? "tree" : undefined;
const href =
treeSegment === undefined
? undefined
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 0913a8d3dd3..29702020ed0 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
@@ -159,3 +159,39 @@ test("the skills section lists every available skill with
its source, loaded one
const link =
document.querySelector('a[href="https://github.com/org/skills/tree/0123456789abcdef/skills/sql-review";]');
assert.ok(link);
});
+
+for (const [label, skills, shown] of [
+ ["a toolset whose sources held no skills says so", [], true],
+ ["a task without a skills toolset shows no skills section", null, false],
+]) {
+ test(label, async (t) => {
+ mock.method(globalThis, "fetch", async (input) => {
+ const key = new URL(input,
"http://localhost";).pathname.split("/").at(-1);
+ if (key === "__AIRFLOW__COMMON_AI_SKILLS__" && skills === null) return
new Response("{}", { status: 404 });
+ 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);
+
+ assert.equal(document.querySelector('[data-testid="skills-section"]') !==
null, shown);
+ assert.equal(/No skills found in the toolset's
sources\./.test(plugin.text()), shown);
+ assert.doesNotMatch(plugin.text(), /0 of 0 loaded/);
+ });
+}
+
+test("a skill badge is a button, so its details are reachable from the
keyboard", async (t) => {
+ const skills = [{ name: "sql-review", description: "Reviews SQL.", source: {
path: "/skills/sql-review" }, calls: {} }];
+ 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 badge = document.querySelector('button[aria-label="About the
sql-review skill"]');
+ assert.equal(badge?.textContent, "sql-review");
+});
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
index 9de5421a61c..b55483f1191 100644
---
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
@@ -30,17 +30,32 @@ test("a local skill reads as its path, with no link", () =>
{
});
});
-for (const [repoUrl, href] of [
- ["https://github.com/org/skills.git";,
`https://github.com/org/skills/tree/${commit}/skills/alpha`],
- ["[email protected]:org/skills.git",
`https://github.com/org/skills/tree/${commit}/skills/alpha`],
- ["ssh://[email protected]/group/sub/skills",
`https://gitlab.com/group/sub/skills/-/tree/${commit}/skills/alpha`],
+for (const [repoUrl, location] of [
+ ["https://github.com/org/skills.git";, "github.com/org/skills"],
+ ["[email protected]:org/skills.git", "github.com/org/skills"],
+ ["ssh://gitlab.com/group/sub/skills", "gitlab.com/group/sub/skills"],
+ // The web UI is on the HTTPS port, not on the SSH one.
+ ["ssh://gitlab.com:2222/group/sub/skills.git",
"gitlab.com/group/sub/skills"],
+ ["https://gitlab.corp.com:8443/team/skills.git/";,
"gitlab.corp.com:8443/team/skills"],
]) {
test(`a Git skill from ${repoUrl} links the folder at the commit that ran`,
() => {
+ const tree = location.includes("gitlab") ? "-/tree" : "tree";
+
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",
+ location,
+ path: "skills/alpha",
+ commit: "0123456",
+ href: `https://${location}/${tree}/${commit}/skills/alpha`,
+ });
+ });
+}
+
+for (const repoUrl of ["/srv/git/skills.git", "http://github.com/org/skills";])
{
+ test(`a remote the tab cannot read, ${repoUrl}, shows as given with no
link`, () => {
+ assert.deepEqual(describeSkillSource({ repo_url: repoUrl, path:
"skills/alpha", commit }), {
+ location: repoUrl,
path: "skills/alpha",
commit: "0123456",
- href,
});
});
}
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 cd71efe9f31..2631246f66e 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/skills.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/skills.py
@@ -106,7 +106,7 @@ class _SkillsLocation:
commit: str | None = None
-def _clone_git(source: GitSkills) -> tuple[str, str, str]:
+def _clone_git(source: GitSkills) -> tuple[str, str, str | None]:
"""
Clone *source* into a fresh temp dir; return (skills_dir,
temp_dir_to_remove, commit).
@@ -171,7 +171,8 @@ def _clone_git(source: GitSkills) -> tuple[str, 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
+ # An empty repository has no commit to record, and its clone still
worked.
+ commit = repo.head.commit.hexsha if repo.head.is_valid() else None
except Exception as exc:
shutil.rmtree(temp_dir, ignore_errors=True)
# GitPython errors embed the failing command, which may contain the
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 caa47f616fa..f3a9607ae39 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
@@ -87,15 +87,23 @@ def _describe_skill_source(skill_uri: str, locations:
list[_SkillsLocation]) ->
}
-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."""
+def _count_skill_calls(messages: Sequence[ModelMessage], run_id: str) ->
defaultdict[str, Counter[str]]:
+ """Count, per skill, the calls to each per-skill tool that succeeded in
the agent run ``run_id``."""
calls: dict[str, ToolCallPart] = {}
counts: defaultdict[str, Counter[str]] = defaultdict(Counter)
for message in messages:
for part in message.parts:
+ # A call approved after a pause was made in the paused run, so
calls are matched across
+ # the whole transcript while only this run's returns are counted.
if isinstance(part, ToolCallPart):
calls[part.tool_call_id] = part
- elif isinstance(part, ToolReturnPart) and (call :=
calls.get(part.tool_call_id)):
+ elif (
+ isinstance(part, ToolReturnPart)
+ and message.run_id == run_id
+ # A denied, failed or interrupted call never ran the skill
tool to completion.
+ and part.outcome == "success"
+ 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}")),
@@ -154,9 +162,14 @@ class AgentSkillsToolset(AbstractToolset):
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.
+ # The run's messages, its history included. Calls replayed on a
durable retry never reach
+ # call_tool, so the skills XCom counts these.
self._messages: Sequence[ModelMessage] = ()
+ # The agent run this instance serves, set by get_tools. ``async with
agent:`` also enters
+ # the instance passed to the agent, which never serves a run and so
never publishes. Not
+ # set in for_run: a DynamicToolset builds this toolset on every model
step, unless given
+ # per_run_step=False, without calling its for_run.
+ self._run_id: str | None = None
@property
def id(self) -> str | None:
@@ -229,25 +242,38 @@ class AgentSkillsToolset(AbstractToolset):
except RuntimeError:
# Used outside an Airflow task: there is no task instance to
publish on.
return
- if not context["ti"].task.do_xcom_push:
+ run_id = self._run_id
+ if run_id is None or 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: {
+ # Only this run's calls: history seeded from another task holds
that task's calls.
+ counts = _count_skill_calls(self._messages, run_id)
+ # An earlier run of this try (before an approval pause or a review
regeneration), or
+ # another skills toolset on this run's agent, may have published
already. Calls are kept
+ # per run: another toolset of this run counts the same calls and
replaces them, while an
+ # earlier run's stay. XComs from an earlier try are cleared before
this one starts.
+ ti = context["ti"]
+ stored = {
+ entry["name"]: entry
+ for entry in ti.xcom_pull(key=SKILLS_XCOM_KEY,
task_ids=ti.task_id, map_indexes=ti.map_index)
+ or []
+ }
+ skills = {}
+ for name, skill in self._inner.skills.items():
+ calls_by_run = dict(stored.get(name, {}).get("calls_by_run",
{}))
+ # A run that made no call to the skill is left out, so the
XCom grows with calls,
+ # not with every resume of a paused run.
+ if calls := counts.get(name):
+ calls_by_run[run_id] = dict(calls)
+ skills[name] = {
"name": name,
"description": skill.description,
"source": _describe_skill_source(skill.uri,
self._locations),
- "calls": dict(counts.get(name, {})),
+ "calls": dict(sum((Counter(calls) for calls in
calls_by_run.values()), Counter())),
+ "calls_by_run": calls_by_run,
}
- 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
+ merged = 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)
@@ -265,6 +291,7 @@ class AgentSkillsToolset(AbstractToolset):
# 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
+ self._run_id = ctx.run_id
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/tests/unit/common/ai/test_skills.py
b/providers/common/ai/tests/unit/common/ai/test_skills.py
index 267bb246463..2556cc478f2 100644
--- a/providers/common/ai/tests/unit/common/ai/test_skills.py
+++ b/providers/common/ai/tests/unit/common/ai/test_skills.py
@@ -19,6 +19,7 @@
from __future__ import annotations
import os
+import subprocess
from unittest.mock import patch
import pytest
@@ -137,6 +138,16 @@ class TestResolveGitWithHook:
assert locations[0].directory.endswith(os.sep + "skills")
assert locations[1] == _SkillsLocation("/local/skills")
+ def test_empty_repository_clones_with_no_commit(self, tmp_path):
+ empty = tmp_path / "empty.git"
+ subprocess.run(["git", "init", "--bare", "-q", str(empty)], check=True)
+ source = GitSkills(repo_url=empty.as_uri())
+
+ locations, cleanup = _materialize_skills([source])
+ cleanup()
+
+ assert locations[0].commit is None
+
@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 6f6bd9a7ec5..64d7c4fb985 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
@@ -26,16 +26,23 @@ import asyncio
import os
import sys
import types
+from collections import Counter
from unittest.mock import MagicMock, patch
import pytest
-from pydantic_ai import Agent
+from pydantic_ai import Agent, DeferredToolRequests, DeferredToolResults
+from pydantic_ai.capabilities import ProcessHistory
from pydantic_ai.exceptions import UnexpectedModelBehavior
-from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart
+from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart,
ToolCallPart, ToolReturnPart
from pydantic_ai.models.function import FunctionModel
+from pydantic_ai.toolsets import DynamicToolset
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.toolsets.skills import (
+ AgentSkillsToolset,
+ _count_skill_calls,
+ _describe_skill_source,
+)
from airflow.providers.common.ai.utils.logging import SKILLS_XCOM_KEY
pytest.importorskip("pydantic_ai_skills")
@@ -279,6 +286,7 @@ class _FakeTaskInstance:
def __init__(self, stored=None):
self.task = types.SimpleNamespace(do_xcom_push=True)
self.xcoms = {} if stored is None else {SKILLS_XCOM_KEY: stored}
+ self.pushes = 0
def xcom_pull(self, key, task_ids, map_indexes):
assert (task_ids, map_indexes) == (self.task_id, self.map_index)
@@ -286,14 +294,17 @@ class _FakeTaskInstance:
def xcom_push(self, key, value):
self.xcoms[key] = value
+ self.pushes += 1
-def _skill_entry(name, path, calls):
+def _skill_entry(name, path, calls_by_run):
+ calls = sum((Counter(calls) for calls in calls_by_run.values()), Counter())
return {
"name": name,
"description": f"A {name} skill for tests.",
"source": {"path": os.path.realpath(path)},
- "calls": calls,
+ "calls": dict(calls),
+ "calls_by_run": calls_by_run,
}
@@ -312,47 +323,145 @@ class TestSkillsXCom:
("load_skill", {"skill_name": "missing"}),
)
- Agent(model,
toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])]).run_sync("go")
+ Agent(model,
toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])]).run_sync("go",
run_id="run")
assert ti.xcoms[SKILLS_XCOM_KEY] == [
- _skill_entry("alpha", alpha, {"load_skill": 1,
"read_skill_resource": 1}),
+ _skill_entry("alpha", alpha, {"run": {"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.
+ def test_counts_a_skill_call_approved_after_a_pause(self,
mock_get_context, tmp_path):
+ # The resumed run gets the paused run's transcript as history: the
approved call was made
+ # there, and only its return is in the resumed run's own messages.
alpha = _write_skill(tmp_path, "alpha")
ti = _FakeTaskInstance()
mock_get_context.return_value = {"ti": ti}
+ agent = Agent(
+ _calling(("load_skill", {"skill_name": "alpha"})),
+ toolsets=[
+ AgentSkillsToolset(sources=[str(tmp_path)]).approval_required(
+ lambda ctx, tool_def, args: tool_def.name == "load_skill"
+ )
+ ],
+ output_type=[str, DeferredToolRequests],
+ )
+ paused = agent.run_sync("go", run_id="paused")
+ approvals = {call.tool_call_id: True for call in
paused.output.approvals}
+
+ agent.run_sync(
+ message_history=paused.all_messages(),
+ deferred_tool_results=DeferredToolResults(approvals=approvals),
+ run_id="resumed",
+ )
+
+ assert ti.xcoms[SKILLS_XCOM_KEY] == [_skill_entry("alpha", alpha,
{"resumed": {"load_skill": 1}})]
+
+ @patch(_GET_CONTEXT, autospec=True)
+ def test_history_seeded_from_another_task_is_not_counted(self,
mock_get_context, tmp_path):
+ alpha = _write_skill(tmp_path, "alpha")
+ mock_get_context.return_value = {"ti": _FakeTaskInstance()}
earlier = Agent(
_calling(("load_skill", {"skill_name": "alpha"})),
toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])],
- ).run_sync("go")
+ ).run_sync("go", run_id="other-task")
+ ti = _FakeTaskInstance()
+ mock_get_context.return_value = {"ti": ti}
+
+ Agent(_calling(),
toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])]).run_sync(
+ "go on", message_history=earlier.all_messages(), run_id="run"
+ )
- Agent(
+ assert ti.xcoms[SKILLS_XCOM_KEY] == [_skill_entry("alpha", alpha, {})]
+
+ @patch(_GET_CONTEXT, autospec=True)
+ def test_counts_calls_after_a_history_processor_trims_the_transcript(self,
mock_get_context, tmp_path):
+ alpha = _write_skill(tmp_path, "alpha")
+ mock_get_context.return_value = {"ti": _FakeTaskInstance()}
+ seeded = Agent(_calling(),
toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])]).run_sync(
+ "go", run_id="other-task"
+ )
+ ti = _FakeTaskInstance()
+ mock_get_context.return_value = {"ti": ti}
+ agent = Agent(
_calling(("load_skill", {"skill_name": "alpha"})),
toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])],
- ).run_sync("go on", message_history=earlier.all_messages())
+ capabilities=[ProcessHistory(lambda messages: messages[-3:])],
+ )
+
+ agent.run_sync("go on", message_history=seeded.all_messages() * 3,
run_id="run")
+
+ assert ti.xcoms[SKILLS_XCOM_KEY] == [_skill_entry("alpha", alpha,
{"run": {"load_skill": 1}})]
+
+ @patch(_GET_CONTEXT, autospec=True)
+ def test_adds_this_run_to_the_stored_runs_and_keeps_the_other_skills(self,
mock_get_context, tmp_path):
+ alpha = _write_skill(tmp_path, "alpha")
+ other = _skill_entry("other", tmp_path / "elsewhere", {"earlier":
{"load_skill": 3}})
+ stored_alpha = _skill_entry("alpha", alpha, {"earlier": {"load_skill":
1, "read_skill_resource": 1}})
+ ti = _FakeTaskInstance(stored=[stored_alpha, other])
+ mock_get_context.return_value = {"ti": ti}
+ model = _calling(("load_skill", {"skill_name": "alpha"}))
+
+ Agent(model,
toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])]).run_sync("go",
run_id="run")
+
+ assert ti.xcoms[SKILLS_XCOM_KEY] == [
+ _skill_entry(
+ "alpha",
+ alpha,
+ {"earlier": {"load_skill": 1, "read_skill_resource": 1},
"run": {"load_skill": 1}},
+ ),
+ other,
+ ]
+
+ @patch(_GET_CONTEXT, autospec=True)
+ def test_toolset_built_per_run_by_a_factory_publishes(self,
mock_get_context, tmp_path):
+ # DynamicToolset, which AgentOperator also wraps a toolset factory in,
never calls the
+ # built toolset's for_run.
+ alpha = _write_skill(tmp_path, "alpha")
+ ti = _FakeTaskInstance()
+ mock_get_context.return_value = {"ti": ti}
+ toolset = DynamicToolset(lambda ctx:
AgentSkillsToolset(sources=[str(tmp_path)]))
+
+ Agent(_calling(("load_skill", {"skill_name": "alpha"})),
toolsets=[toolset]).run_sync(
+ "go", run_id="run"
+ )
- assert ti.xcoms[SKILLS_XCOM_KEY] == [_skill_entry("alpha", alpha,
{"load_skill": 2})]
+ assert ti.xcoms[SKILLS_XCOM_KEY] == [_skill_entry("alpha", alpha,
{"run": {"load_skill": 1}})]
@patch(_GET_CONTEXT, autospec=True)
- def test_replaces_its_own_stored_skills_and_keeps_the_others(self,
mock_get_context, tmp_path):
+ def test_entering_the_agent_publishes_only_from_the_run(self,
mock_get_context, tmp_path):
+ # ``async with agent:`` enters the toolset itself, besides the copy
for_run hands the run.
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])
+ ti = _FakeTaskInstance()
mock_get_context.return_value = {"ti": ti}
+ agent = Agent(
+ _calling(("load_skill", {"skill_name": "alpha"})),
+ toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])],
+ )
+
+ async def run():
+ async with agent:
+ await agent.run("go", run_id="run")
- Agent(_calling(),
toolsets=[AgentSkillsToolset(sources=[str(tmp_path)])]).run_sync("go")
+ asyncio.run(run())
- assert ti.xcoms[SKILLS_XCOM_KEY] == [_skill_entry("alpha", alpha, {}),
other]
+ assert ti.xcoms[SKILLS_XCOM_KEY] == [_skill_entry("alpha", alpha,
{"run": {"load_skill": 1}})]
+ assert ti.pushes == 1
+ @pytest.mark.parametrize(
+ ("second_skill", "expected_second_calls"),
+ [
+ pytest.param("beta", {"load_skill": 1}, id="different_skills"),
+ # Each toolset counts both calls to "alpha"; publishing twice must
not count them twice.
+ pytest.param("alpha", None, id="same_skill_name"),
+ ],
+ )
@patch(_GET_CONTEXT, autospec=True)
- def test_two_prefixed_toolsets_each_count_their_own_skills(self,
mock_get_context, tmp_path):
+ def test_two_prefixed_toolsets_each_count_their_own_skills(
+ self, mock_get_context, tmp_path, second_skill, expected_second_calls
+ ):
alpha = _write_skill(tmp_path / "one", "alpha")
- beta = _write_skill(tmp_path / "two", "beta")
+ second = _write_skill(tmp_path / "two", second_skill)
ti = _FakeTaskInstance()
mock_get_context.return_value = {"ti": ti}
toolsets = [
@@ -361,47 +470,49 @@ class TestSkillsXCom:
]
model = _calling(
("one_load_skill", {"skill_name": "alpha"}),
- ("two_load_skill", {"skill_name": "beta"}),
+ ("two_load_skill", {"skill_name": second_skill}),
)
- Agent(model, toolsets=toolsets).run_sync("go")
+ Agent(model, toolsets=toolsets).run_sync("go", run_id="run")
- assert ti.xcoms[SKILLS_XCOM_KEY] == [
- _skill_entry("alpha", alpha, {"load_skill": 1}),
- _skill_entry("beta", beta, {"load_skill": 1}),
- ]
+ if expected_second_calls is None:
+ # Toolsets exit in reverse order, so the first one's source is
published last.
+ assert ti.xcoms[SKILLS_XCOM_KEY] == [_skill_entry("alpha", alpha,
{"run": {"load_skill": 2}})]
+ else:
+ assert ti.xcoms[SKILLS_XCOM_KEY] == [
+ _skill_entry("alpha", alpha, {"run": {"load_skill": 1}}),
+ _skill_entry(second_skill, second, {"run":
expected_second_calls}),
+ ]
@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"),
- ],
+ "in_task",
+ [pytest.param(False, id="outside_a_task"), pytest.param(True,
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):
+ def test_nothing_pushed_without_a_task_that_pushes_xcom(self,
mock_get_context, tmp_path, in_task):
_write_skill(tmp_path, "alpha")
- if isinstance(context, Exception):
- mock_get_context.side_effect = context
+ ti = _FakeTaskInstance()
+ ti.task.do_xcom_push = False
+ if in_task:
+ mock_get_context.return_value = {"ti": ti}
else:
- mock_get_context.return_value = context
+ mock_get_context.side_effect = RuntimeError("no task")
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()
+ mock_get_context.assert_called_once()
+ assert ti.pushes == 0
+ @patch.object(
+ _FakeTaskInstance, "xcom_push", autospec=True,
side_effect=RuntimeError("xcom backend down")
+ )
@patch(_GET_CONTEXT, autospec=True)
- def test_failed_push_does_not_fail_the_run_or_skip_cleanup(self,
mock_get_context, tmp_path):
+ def test_failed_push_does_not_fail_the_run_or_skip_cleanup(self,
mock_get_context, mock_push, 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}
+ mock_get_context.return_value = {"ti": _FakeTaskInstance()}
cleanup = MagicMock()
toolset = AgentSkillsToolset(sources=[str(tmp_path)])
@@ -412,9 +523,37 @@ class TestSkillsXCom:
result = Agent(_calling(), toolsets=[toolset]).run_sync("go")
assert result.output == "done"
+ mock_push.assert_called_once()
cleanup.assert_called_once()
+class TestCountSkillCalls:
+ @pytest.mark.parametrize(
+ ("outcome", "expected"),
+ [
+ pytest.param("success", {"alpha": {"run_skill_script": 1}},
id="success"),
+ pytest.param("denied", {}, id="denied"),
+ pytest.param("failed", {}, id="failed"),
+ pytest.param("interrupted", {}, id="interrupted"),
+ # A missing resource or script: pydantic-ai records the ModelRetry
as a "retried" return.
+ pytest.param("retried", {}, id="retried"),
+ ],
+ )
+ def test_counts_only_calls_that_succeeded(self, outcome, expected):
+ args = {"skill_name": "alpha", "script_name": "run.py"}
+ messages = [
+ ModelResponse(
+ parts=[ToolCallPart("run_skill_script", args,
tool_call_id="call-1")], run_id="run"
+ ),
+ ModelRequest(
+ parts=[ToolReturnPart("run_skill_script", "output",
tool_call_id="call-1", outcome=outcome)],
+ run_id="run",
+ ),
+ ]
+
+ assert _count_skill_calls(messages, "run") == expected
+
+
class TestDescribeSkillSource:
@pytest.mark.parametrize(
("repo_path", "expected_path"),
@@ -436,12 +575,15 @@ class TestDescribeSkillSource:
"commit": "0123abc",
}
- def test_nested_source_wins_over_the_one_containing_it(self, tmp_path):
+ @pytest.mark.parametrize("nested_first", [False, True],
ids=["containing_first", "nested_first"])
+ def test_nested_source_wins_over_the_one_containing_it(self, tmp_path,
nested_first):
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"),
]
+ if nested_first:
+ locations.reverse()
assert _describe_skill_source(str(skill_dir), locations)["repo_url"]
== "https://github.com/org/repo";
