pierrejeambrun commented on code in PR #76541:
URL: https://github.com/apache/airflow/pull/76541#discussion_r4482074171
##########
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:
You're right, thanks for catching this. My test chained two completed runs,
so it never hit the deferred path. A call now counts when its return is in a
message stamped with the current `run_id`, and calls are matched across the
whole transcript, so the approved call counts once in the resumed run.
`test_counts_a_skill_call_approved_after_a_pause` uses `.approval_required()`
with real `DeferredToolResults`: 0 on the previous commit, 1 now.
##########
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:
Switched to `ctx.run_id`. A trimming `ProcessHistory` no longer moves the
boundary; covered by
`test_counts_calls_after_a_history_processor_trims_the_transcript`.
##########
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:
Counts are now kept per run under `calls_by_run`, and `calls` totals them. A
second toolset in the same run replaces that run's counts instead of adding to
them, and earlier runs keep theirs. The two-toolset test now also covers two
toolsets with the same skill name.
##########
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:
Done: the sentence now covers both the approval resume and the review
regeneration.
##########
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:
Added `mock_get_context.assert_called_once()`.
##########
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:
Added both remotes; the `/srv/git` case fails without the `try`/`catch`.
--
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]
