kaxil commented on code in PR #76541:
URL: https://github.com/apache/airflow/pull/76541#discussion_r4481389329
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/skills.py:
##########
@@ -229,11 +239,12 @@ def _push_skills_xcom(self) -> None:
except RuntimeError:
# Used outside an Airflow task: there is no task instance to
publish on.
return
- if not context["ti"].task.do_xcom_push:
+ if self._history_length 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)
+ # History seeded from another task holds that task's calls, made
against its own clone.
+ counts = _count_skill_calls(self._messages[self._history_length :])
Review Comment:
On a real approval resume this drops the approved call. `AgentOperator`
resumes with the paused transcript as `message_history`, so the approved call's
`ToolCallPart` sits before `_history_length` while its `ToolReturnPart` lands
in the resumed run's own messages, and `calls.get(part.tool_call_id)` misses
it. The paused run had no return yet, so neither run counts it. With
`.approval_required()` on the toolset, a pause and an approve, the XCom ends at
`calls: {}` where main gives `load_skill: 1`. Could the call map be built from
all of `self._messages`, counting only the returns after the boundary? The new
approval-pause test chains two completed runs with no `DeferredToolResults`, so
it passes on main as well and never reaches this path.
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/skills.py:
##########
@@ -166,12 +174,14 @@ async def for_run(self, ctx: RunContext) ->
AbstractToolset:
# Per-run isolation: pydantic-ai shares one toolset instance across
runs,
# but we hold per-run clone/cleanup state on __aenter__/__aexit__. Hand
# each run its own instance so concurrent runs never clobber each
other.
- return AgentSkillsToolset(
+ toolset = AgentSkillsToolset(
self._sources,
exclude_tools=self._exclude_tools,
exclude_resources=self._exclude_resources,
max_retries=self._max_retries,
)
+ toolset._history_length = len(ctx.messages)
Review Comment:
`len(ctx.messages)` can stop matching the run's own messages, because
pydantic-ai rewrites the message list in place with whatever history processors
return. A seeded `message_history` plus a trimming `ProcessHistory` capability
leaves the list shorter than this index, so all of the run's calls get sliced
off (I probed it at the provider's pydantic-ai floor and at the locked version:
1 call made, 0 counted). Could this key on `ctx.run_id` instead? Looking calls
up across the whole transcript would still be needed for the approval resume
case.
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/skills.py:
##########
@@ -243,11 +254,19 @@ def _push_skills_xcom(self) -> None:
}
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.
+ # Another skills toolset on the agent, or an earlier run of this
try (before an approval
+ # pause or a review regeneration), may have published already;
this run's calls add to
+ # theirs. 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
+ 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 []
+ }
+ for name, entry in skills.items():
+ if name in stored:
+ entry["calls"] = dict(Counter(stored[name]["calls"]) +
Counter(entry["calls"]))
Review Comment:
This also adds within a single run, not only across runs.
`_count_skill_calls` matches any `*_load_skill` by suffix, so with two
`prefixed()` skills toolsets that both have a skill named `alpha`, each one
counts the same `one_load_skill(alpha)` call and the second publish turns 1
into 2 (main replaced the entry, so it stayed at 1). Should this only add when
the stored entry came from an earlier run, or should each toolset count only
calls made through its own tools? The two-toolset test uses `alpha` and `beta`,
so it doesn't catch this.
##########
providers/common/ai/tests/unit/common/ai/test_skills.py:
##########
@@ -137,6 +137,16 @@ def test_records_the_cloned_commit(self, mock_repo):
assert locations[0].directory.endswith(os.sep + "skills")
assert locations[1] == _SkillsLocation("/local/skills")
+ @patch("git.Repo")
+ def test_empty_repository_clones_with_no_commit(self, mock_repo):
+ mock_repo.clone_from.return_value.head.is_valid.return_value = False
Review Comment:
`head.commit` on this mock never raises, so the test doesn't reproduce the
empty-repo failure: a `try`/`except ValueError` fix would fail it, and a later
unguarded `head.commit` read would still pass. Cloning a real `git init --bare`
repo from `tmp_path` would test the actual case with no network.
##########
providers/common/ai/docs/toolsets/skills.rst:
##########
@@ -140,10 +140,14 @@ 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 that completed count, not denied, failed or interrupted ones, and a
+``message_history`` seeded from another task does not add that task's calls.
When a run
+resumes from a tool-approval pause, the source shows the commit the resumed
run cloned.
Review Comment:
A review regeneration has the same caveat. `regenerate_with_feedback` builds
a new agent, so the toolset clones the branch again, possibly at a newer
commit. Can this sentence mention it too?
##########
providers/common/ai/tests/unit/common/ai/toolsets/test_skills.py:
##########
@@ -372,36 +420,33 @@ def
test_two_prefixed_toolsets_each_count_their_own_skills(self, mock_get_contex
]
@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()
+ assert ti.pushes == 0
Review Comment:
For `outside_a_task` this `ti` never reaches the toolset, so `ti.pushes ==
0` holds by construction. `mock_get_context.assert_called_once()` would make
that case check something.
##########
providers/common/ai/src/airflow/providers/common/ai/plugins/www/tests/skill-source.test.ts:
##########
@@ -30,17 +30,22 @@ 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"],
Review Comment:
Could this also cover a remote that doesn't parse, like
`/srv/git/skills.git` or `http://github.com/org/skills`? `GitSkills` accepts
both, and removing the `try`/`catch` around `new URL()` still passes all 9
tests, even though that would throw during render instead of falling back to
the raw location.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
