Lee-W commented on code in PR #73984:
URL: https://github.com/apache/airflow/pull/73984#discussion_r4151545217
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -146,9 +148,49 @@ def get_link(
)
-def _is_concrete_toolset_capability(capability: Any) -> bool:
- """Whether *capability* is a ``Toolset`` holding a toolset, not a callable
factory resolved per run."""
- return isinstance(capability, Toolset) and isinstance(capability.toolset,
AbstractToolset)
+def _capability_toolset(capability: object) -> AbstractToolset[Any] | None:
Review Comment:
```suggestion
def _resolve_capability_toolset(capability: object) -> AbstractToolset[Any]
| None:
```
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -236,6 +278,18 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin):
Dag file is not modified. Derive the connection ID from values the Dag
controls rather than ``params`` or ``dag_run.conf``, which whoever
triggers
the Dag controls.
+ :param capabilities: pydantic-ai capabilities for the agent, e.g.
+ ``[Thinking(effort="high"), WebSearch()]``. A capability bundles tools,
+ instructions, model settings and lifecycle hooks; pydantic-ai wraps
their
+ hooks in list order, first outermost, unless a capability declares its
+ own position. A ``Toolset`` capability holding one of the
+ toolsets above has its connection IDs templated the same way as
+ ``toolsets=``. Capabilities passed here are not stored in the
serialized
+ Dag, except on a mapped task: the worker builds them from the Dag file.
+ Passing ``capabilities``
Review Comment:
```suggestion
``toolsets=``. Capabilities passed here are not stored in the
serialized Dag (the worker builds them from the Dag file), except on a mapped
task, where they are stored as their repr.
Passing ``capabilities``
```
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -719,11 +798,15 @@ def _declared_toolsets(self) ->
list[AbstractToolset[Any]]:
for toolset in (*(self.toolsets or []),
*(self.agent_params.get("toolsets") or []))
if isinstance(toolset, AbstractToolset)
]
- for capability in self.agent_params.get("capabilities") or ():
- if _is_concrete_toolset_capability(capability):
- candidates.append(capability.toolset)
+ for capability in self._declared_capabilities():
+ if (toolset := _capability_toolset(capability)) is not None:
+ candidates.append(toolset)
return candidates
+ def _declared_capabilities(self) -> list[AgentCapability[Any]]:
Review Comment:
```suggestion
@property
def _declared_capabilities(self) -> list[AgentCapability[Any]]:
```
--
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]
