Lee-W commented on code in PR #76877:
URL: https://github.com/apache/airflow/pull/76877#discussion_r4534688770
##########
providers/common/ai/src/airflow/providers/common/ai/durable/capability.py:
##########
@@ -97,6 +101,100 @@
DURABLE_UNIT_TOOLSETS = (FunctionToolset, DynamicToolset, MCPToolset)
"""Leaf toolsets pydantic-ai's durable backend runs as durable units; each
needs a unique ``id``."""
+# Airflow imports a Dag file as a module named ``unusual_prefix_<sha1 of the
file's path>_<file name>``
+# (``get_unique_dag_module_name`` in the shared ``module_loading`` library,
which providers cannot import).
+_DAG_MODULE_PREFIX = re.compile(r"^unusual_prefix_[0-9a-f]{40}_")
+
+
+def callable_toolset_id(function: object) -> str | None:
Review Comment:
```suggestion
def get_callable_toolset_id(function: object) -> str | None:
```
##########
providers/common/ai/src/airflow/providers/common/ai/durable/capability.py:
##########
@@ -97,6 +101,100 @@
DURABLE_UNIT_TOOLSETS = (FunctionToolset, DynamicToolset, MCPToolset)
"""Leaf toolsets pydantic-ai's durable backend runs as durable units; each
needs a unique ``id``."""
+# Airflow imports a Dag file as a module named ``unusual_prefix_<sha1 of the
file's path>_<file name>``
+# (``get_unique_dag_module_name`` in the shared ``module_loading`` library,
which providers cannot import).
+_DAG_MODULE_PREFIX = re.compile(r"^unusual_prefix_[0-9a-f]{40}_")
+
+
+def callable_toolset_id(function: object) -> str | None:
+ """
+ Return the id of a toolset that ``function`` builds per run: its module
and qualified name.
+
+ The hash of the file's path that Airflow puts in a Dag file's module name
is left out,
+ because the path can differ between workers and bundle versions. ``None``
for a callable
+ without a qualified name, such as a ``functools.partial``.
+ """
+ module = getattr(function, "__module__", None)
+ qualname = getattr(function, "__qualname__", None)
+ if not isinstance(module, str) or not isinstance(qualname, str):
+ return None
+ return f"{_DAG_MODULE_PREFIX.sub('', module)}.{qualname}"
+
+
+def chosen_toolset_id(leaf: AbstractToolset[Any]) -> str | None:
+ """
+ Return the id pydantic-ai names a leaf toolset's durable steps after, or
``None`` when it has none.
+
+ An ``MCPToolset`` without an id derives one from its server's hostname,
but durable steps
+ are only named after an id that was chosen for the toolset.
+ """
+ # Before ``_explicit_id`` (pydantic-ai 2.62) no toolset derived its id, so
``id`` is the chosen one.
+ return getattr(leaf, "_explicit_id", leaf.id)
+
+
+# The id ``give_durable_ids`` gives the n-th ``FunctionToolset`` without one.
The same agent gets the
+# same ids on every attempt, so its requests to the model match the recorded
ones. Angle brackets
+# mark an id the framework chose, as in pydantic-ai's own ``<agent>``; there
is no colon, which
+# pydantic-ai rejects in a toolset id because it builds instruction ids from
it.
+_ASSIGNED_ID_PREFIX = "<airflow-function_toolset-"
+
+
+def _assigned_toolset_id(number: int) -> str:
+ return f"{_ASSIGNED_ID_PREFIX}{number}>"
+
+
+# Each leaf ``give_durable_ids`` copied, keyed by ``id()``, with its copy. The
leaf is kept so
+# that its ``id()`` cannot be reused by another object while the map is in use.
+GivenIds: TypeAlias = "dict[int, tuple[AbstractToolset[Any],
AbstractToolset[Any]]]"
+
+
+def give_durable_ids(
+ toolset: AbstractToolset[Any] | ToolsetFunc[Any], given: GivenIds
+) -> AbstractToolset[Any]:
+ """
+ Return ``toolset`` with an id on each leaf pydantic-ai runs as durable
steps, where one can be given.
+
+ pydantic-ai runs a ``FunctionToolset`` or ``DynamicToolset`` as durable
steps only when
+ it has an id. An id-less ``FunctionToolset`` gets one numbered in the
order the toolsets
+ are given, which its step names leave out (a function tool's step is named
after the
+ tool); a function that builds a toolset per run, or an id-less
``DynamicToolset``, gets
+ :func:`callable_toolset_id`. A leaf that already has an id, and an
``MCPToolset``, are
+ left as they are. Leaves are copied, not changed, and ``given`` maps each
leaf to its
+ copy so a toolset listed twice stays one toolset to pydantic-ai.
+ """
+ if not isinstance(toolset, AbstractToolset):
+ return DynamicToolset(toolset, id=callable_toolset_id(toolset))
+ if not any(_lacks_id(leaf) for leaf in iter_toolsets(toolset)):
+ # Rebuilding wrappers via visit_and_replace breaks wrapper subclasses
with their own __init__.
+ return toolset
+
+ def give(leaf: AbstractToolset[Any]) -> AbstractToolset[Any]:
+ if not _lacks_id(leaf):
+ return leaf
+ if (entry := given.get(id(leaf))) is None:
+ copied: AbstractToolset[Any]
+ if isinstance(leaf, DynamicToolset):
+ copied = DynamicToolset(
+ leaf.toolset_func,
+ per_run_step=leaf.per_run_step,
+ id=callable_toolset_id(leaf.toolset_func),
+ )
+ else:
+ numbered = sum(isinstance(other, FunctionToolset) for _, other
in given.values())
+ function_toolset = copy.copy(leaf)
+ # FunctionToolset takes its id only in the constructor and
keeps it here.
+ function_toolset._id = _assigned_toolset_id(numbered)
Review Comment:
Nit: `function_toolset._id = ...` writes a pydantic-ai private attribute on
a `copy.copy` of the toolset. It works now, but a refactor there would break it
without any error here. We should add a test to ensure this is caught when
happens
##########
providers/common/ai/docs/durable_execution.rst:
##########
@@ -65,30 +65,64 @@ The capability needs a name for the agent's steps: the
agent's ``name``, or
``agent_params`` sets one, and the task id otherwise. Outside a running
Airflow task
the capability does nothing, so the same agent runs normally in a test or a
notebook.
-Give every toolset an id
-------------------------
+.. _durable-toolset-ids:
-pydantic-ai records each call of a ``FunctionToolset``, ``MCPToolset`` or
-``DynamicToolset`` under the toolset's ``id``, and refuses to run one without
it. Set
-the id on the toolset itself:
+Toolset ids
+-----------
+
+pydantic-ai names the recorded steps of a ``FunctionToolset``, ``MCPToolset``
or
+``DynamicToolset`` after the toolset's ``id``. ``AgentOperator`` and
``@task.agent`` fill
+in the ids you leave out of a ``FunctionToolset`` and of a function that
builds a toolset
+for each run passed in ``toolsets`` or ``agent_params["toolsets"]``, and of a
+``FunctionToolset`` held directly by a ``Toolset`` capability, so this agent
needs none:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py
+ :language: python
+ :start-after: [START howto_agent_durable_toolsets_without_id]
+ :end-before: [END howto_agent_durable_toolsets_without_id]
+
+A tool call of a ``FunctionToolset`` without an id is recorded under the
tool's name
+alone, ``durable_order_status__function_toolset.call_tool:get_order`` above,
so a retry
+of the same agent finds the calls the failed attempt recorded. A function that
builds a
+toolset, such as ``carrier_tools``, is named by its module and qualified name
+(``example_agent_durable.carrier_tools``). A toolset you give an ``id`` keeps
it in its
+step names. Adding a tool to a toolset, reordering the toolsets, or renaming a
function
+or an id between attempts changes what the model is offered, so the retry runs
again
+from the first model request (:ref:`durable-replay`). A function that builds a
toolset
+per run is the exception for its own tools: its tool list is a recorded step,
so the
+steps a retry replays, and the request that failed, keep the tools the failed
attempt
+saw, and later steps get the new list. A clear that deletes the journal
+(:ref:`durable-clear`) starts over with the new list.
+
+An ``MCPToolset`` from pydantic-ai still needs an id from you. Its calls are
named after
+the toolset alone, and no id is derived from the server's URL or command,
which can carry
+credentials. The Dag fails to load with:
-.. code-block:: python
+.. code-block:: text
- orders = FunctionToolset(tools=[get_order], id="orders")
+ ValueError: durable=True needs a unique id on every MCPToolset, such as
MCPToolset(..., id='weather'): durable replay records each tool call under it,
and none is derived from the server's URL or command, which can carry
credentials. Set it on the toolset itself; the id of a Toolset capability
around a toolset does not reach it.
-``AgentOperator(durable=True)`` checks this when the Dag is parsed and fails
with:
+So does a callable that has no ``__qualname__`` to name it by, such as a
+``functools.partial``; pass it as ``DynamicToolset(function, id=...)``:
.. code-block:: text
- ValueError: durable=True needs a unique id on every FunctionToolset, such
as FunctionToolset(..., id='orders'): durable replay records each tool call
under it. Set it on the toolset itself, and pass a function that builds a
toolset as DynamicToolset(function, id=...); the id of a Toolset capability
around a toolset does not reach it.
+ ValueError: durable=True names the steps of a toolset built per run after
the function that builds it, but functools.partial(<function carrier_tools at
0xffff89c904a0>, carrier='dhl') has no __qualname__. Pass it as
DynamicToolset(function, id=...).
Review Comment:
Nit: the sample error contains a memory address, which will differ on every
run. Add "(address varies)" or elide it so the sample does not read as fixed
output.
##########
providers/common/ai/src/airflow/providers/common/ai/durable/capability.py:
##########
@@ -380,12 +482,36 @@ async def _model_for_fingerprint(
return model
-class _AirflowOperationBackend(JournalCallableOperationBackend[_StepMode]):
+class _AirflowOperationNamer(JournalOperationNamer):
+ """
+ pydantic-ai's journal step names, without an id Airflow gave a function
toolset.
+
+ A function tool's step name ends in the tool's name, so the id Airflow
numbered the
+ toolset by adds nothing to it and would only make the name harder to read
in the task
+ log. An id the Dag author chose stays in the name, so steps recorded under
it before
+ this naming existed replay as they did.
+ """
+
+ def operation_name(self, operation_id: DurableOperationId) -> str:
Review Comment:
```suggestion
def get_operation_name(self, operation_id: DurableOperationId) -> str:
```
Nit: it rewrites both `ToolsetCallToolId` and
`ToolsetValidateToolArgumentsId`, but the new tests assert only the `call_tool`
step name. Assert the `validate_tool_arguments` step name too, or drop that
branch if pydantic-ai does not journal it.
--
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]
