amoghrajesh commented on code in PR #76877:
URL: https://github.com/apache/airflow/pull/76877#discussion_r4548718017
##########
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:
Rename to `_get_chosen_toolset_id`?
##########
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:
Should these be public? Only agent operator uses it.
##########
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:
We are using angled braces here but pydantic-ai's `AbstractToolset.id`:
https://github.com/pydantic/pydantic-ai/blob/main/pydantic_ai_slim/pydantic_ai/toolsets/abstract.py#L106-L108
reserves those for its own roles (`<agent>`, `<output>`). Nothing checks this
today, but a later `pydantic-ai` release that treats bracketed ids as framework
owned would break replay silently. We can avoid it.
##########
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:
Same as above
##########
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:
Same as above
##########
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:
Same here
--
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]
