kaxil commented on code in PR #76877:
URL: https://github.com/apache/airflow/pull/76877#discussion_r4550718843
##########
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 get_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-"
Review Comment:
Agreed, switched to `__airflow_function_toolset_<n>` in 79212d5d22e. No
brackets and no colon, and the dunder prefix keeps it apart from ids people
pick. It hasn't shipped, so nothing recorded under the old format to worry
about.
##########
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 get_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:
Review Comment:
Renamed to `_get_chosen_toolset_id` in 79212d5d22e.
##########
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 get_callable_toolset_id(function: object) -> str | None:
Review Comment:
Right, only `AgentOperator` and the tests use them. Made all four private in
79212d5d22e (`_get_callable_toolset_id`, `_get_chosen_toolset_id`,
`_give_durable_ids`, `_GivenIds`). `agent.py` imports them across modules like
`tools/strands.py` already does with `_collect_tools`.
##########
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 get_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:
Review Comment:
Private now too, as `_get_chosen_toolset_id`, in 79212d5d22e.
##########
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 get_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]]]"
Review Comment:
Now `_GivenIds`, in 79212d5d22e.
##########
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 get_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(
Review Comment:
Now `_give_durable_ids`, in 79212d5d22e.
--
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]
