kaxil commented on code in PR #72786:
URL: https://github.com/apache/airflow/pull/72786#discussion_r3973648603
##########
providers/common/ai/docs/toolsets.rst:
##########
@@ -1420,6 +1474,13 @@ The standby-served fraction is a ratio over
``managed_agent.served`` alone, so
grep. Both are tagged by platform rather than agent name to keep cardinality
bounded.
+``managed_agent.invoked`` (see `Metrics`_ above) is the denominator for
Review Comment:
`managed_agent.failover` fires once per transition, not once per call, so a
three-member group where the first two fail adds two to the numerator and one
to the denominator. Worth describing this as failovers per call rather than a
rate, since it is not bounded by 1. Also, the `list-table` just above still
says "Two counters" and has no `managed_agent.invoked` row, which is where the
warning in the Metrics section sends the reader ("see the metrics table below").
##########
providers/common/ai/docs/toolsets.rst:
##########
@@ -1358,6 +1394,24 @@ Groups nest.
members=[bedrock_claims_agent, foundry_claims_agent], # same image,
two clouds
)
+.. note::
+
+ Each member still needs its own ``tool_name`` -- the constructor rejects
+ an empty one -- but nothing reads it once the member is inside a group.
+ The group calls only ``member.invoke()``, never ``member.get_tools()``,
+ so a member's ``tool_name``, ``description``, and ``max_retries`` are
+ inert; only the group's own values reach the calling model.
+
+.. note::
+
+ The same is true in reverse for ``timeout``:
``FailoverManagedAgentToolset``
+ does not override the ``timeout`` property, so it just mirrors whatever
+ the group's own constructor received, and the group's ``invoke()`` above
+ calls only ``member.invoke(prompt)`` -- it never reads ``self.timeout``.
+ Passing ``timeout=`` to a failover group is therefore a silent no-op; set
Review Comment:
If `timeout=` on a group is always a no-op, could
`FailoverManagedAgentToolset.__init__` reject it instead? It already raises for
fewer than two members, and the base class raises for an empty `tool_name` or a
negative `max_retries`, so a `ValueError` pointing the author at the members
would fit the existing validation and would not depend on finding this note.
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py:
##########
@@ -232,9 +273,17 @@ async def call_tool(
ctx: RunContext[Any],
tool: ToolsetTool[Any],
) -> Any:
- ref = self.agent_ref
+ ref = _safe_agent_ref(self)
log.info("Consulting managed agent %s on %s", ref.get("name"),
ref.get("platform"))
result = await self.invoke(tool_args["prompt"])
+ # Emitted here, not in FailoverManagedAgentToolset.invoke(), so a lone
+ # toolset gets it too -- it is the only per-tool signal a bare toolset
+ # has, and it is the ratio that turns managed_agent.failover from a raw
+ # count into a rate (failovers / invocations for the same tool).
+ Stats.incr(
Review Comment:
This increments only after `invoke()` returns, so it counts answers rather
than calls (`test_call_tool_does_not_emit_invoked_on_failure` pins that). It
makes an awkward denominator for `managed_agent.failover`: a total outage still
increments `failover` while `invoked` stays flat, so the documented ratio
climbs fastest exactly when the group has stopped answering, and there is
nothing to divide by when nothing succeeds at all. Moving the `Stats.incr`
above the `await` would count attempts instead, which matches the name and
makes `served`/`invoked` a success rate for a group (that test would then flip
to asserting it *is* emitted). Keeping it success-only is fine too, but then
the docs line and the registry description should stop calling it the
denominator.
##########
providers/common/ai/tests/unit/common/ai/toolsets/test_managed_agent.py:
##########
@@ -435,3 +456,84 @@ async def test_model_retry_is_not_a_failover(self,
mock_stats):
with pytest.raises(ModelRetry):
await group.invoke("q")
mock_stats.incr.assert_not_called()
+
+
+class TestSafeAgentRef:
+ """agent_ref only labels a call; resolving it must never block the call
itself."""
+
+ async def _call(self, toolset, prompt="what is the number?"):
+ tools = await toolset.get_tools(ctx=None)
+ tool = tools[toolset._tool_name]
+ return await toolset.call_tool(toolset._tool_name, {"prompt": prompt},
None, tool)
+
+ @staticmethod
+ def _group(*members, **kwargs):
+ kwargs.setdefault("tool_name", "ask_resilient")
+ kwargs.setdefault("description", "Answers questions, on whichever
cloud is up.")
+ return FailoverManagedAgentToolset(members=list(members), **kwargs)
+
+ @pytest.mark.asyncio
+ async def
test_call_tool_survives_a_broken_agent_ref_on_a_lone_toolset(self):
+ toolset = BrokenAgentRefManagedAgentToolset(result="ok")
+ assert await self._call(toolset) == "ok"
+
+ @pytest.mark.asyncio
+ async def test_call_tool_survives_a_broken_standby_agent_ref(self):
+ # Kaxil's original repro: a group whose agent_ref join fails must still
+ # try its healthy primary through call_tool.
+ primary = FakeManagedAgentToolset(result="from primary")
+ standby = BrokenAgentRefManagedAgentToolset(result="from standby")
+ group = self._group(primary, standby)
+ assert await self._call(group) == "from primary"
+
+ @pytest.mark.asyncio
+ async def test_failover_survives_a_broken_primary_agent_ref(self):
+ primary = BrokenAgentRefManagedAgentToolset(result="from primary")
+ standby = FakeManagedAgentToolset(result="from standby")
+ group = self._group(primary, standby)
+ assert await group.invoke("q") == "from primary"
+
+ @pytest.mark.asyncio
+ async def test_failover_survives_a_broken_standby_agent_ref(self):
+ primary =
FakeManagedAgentToolset(raises=ManagedAgentInvocationError("down"))
+ standby = BrokenAgentRefManagedAgentToolset(result="from standby")
+ group = self._group(primary, standby)
+ assert await group.invoke("q") == "from standby"
+
+
+class TestInvokedMetric:
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.common.ai.toolsets.managed_agent.Stats")
+ async def test_call_tool_emits_invoked_on_success(self, mock_stats):
+ toolset = FakeManagedAgentToolset()
+ tools = await toolset.get_tools(ctx=None)
+ await toolset.call_tool("ask_specialist", {"prompt": "q"}, None,
tools["ask_specialist"])
+ mock_stats.incr.assert_called_once_with(
+ "managed_agent.invoked",
+ tags={"tool": "ask_specialist", "platform": "fake.cloud"},
+ )
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.common.ai.toolsets.managed_agent.Stats")
+ async def test_call_tool_does_not_emit_invoked_on_failure(self,
mock_stats):
+ toolset = FakeManagedAgentToolset(raises=RuntimeError("503"))
+ tools = await toolset.get_tools(ctx=None)
+ with pytest.raises(RuntimeError):
+ await toolset.call_tool("ask_specialist", {"prompt": "q"}, None,
tools["ask_specialist"])
+ mock_stats.incr.assert_not_called()
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.common.ai.toolsets.managed_agent.Stats")
+ async def
test_group_call_tool_emits_invoked_once_regardless_of_failover(self,
mock_stats):
+ primary =
FakeManagedAgentToolset(raises=ManagedAgentInvocationError("down"))
+ standby = FakeManagedAgentToolset(result="from standby")
+ group = FailoverManagedAgentToolset(
+ members=[primary, standby],
+ tool_name="ask_resilient",
+ description="Answers questions, on whichever cloud is up.",
+ )
+ tools = await group.get_tools(ctx=None)
+ await group.call_tool("ask_resilient", {"prompt": "q"}, None,
tools["ask_resilient"])
+
+ kinds = [c.args[0] for c in mock_stats.incr.call_args_list]
+ assert kinds == ["managed_agent.failover", "managed_agent.served",
"managed_agent.invoked"]
Review Comment:
Could this assert the tags too? A group emits `invoked` with
`platform="failover"` rather than the cloud that answered, which is the
surprising half of the group case and nothing currently pins it. It also means
summing `managed_agent.invoked` by platform mixes real platforms with
`failover`, which is probably worth a sentence in the Metrics docs section.
--
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]
