Skip to content

Correct the common-ai toolset list and document the shields extra #71819

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Lee-W wants to merge 1 commit into
base: main
Choose a base branch
from
+14 −5
Open

Correct the common-ai toolset list and document the shields extra #71819

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions providers/common/ai/docs/index.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,8 +102,9 @@ The provider's extras split into a few groups:
``pydantic-ai-slim`` optional dependency group; pydantic-ai supports more model providers
than these four, each under its own extra name, so check the
`pydantic-ai install docs <https://ai.pydantic.dev/install/#slim-install>`__ for the full list.
* **Agent tooling** — ``mcp``, ``skills``, ``code-mode``: MCP servers, Agent Skills, and
code-mode tool execution.
* **Agent tooling** — ``mcp``, ``skills``, ``code-mode``, ``shields``: MCP servers, Agent
Skills, code-mode tool execution, and shield capabilities (input/output guards, tool
guards, cost tracking).
* **Document loading** — ``pdf``, ``docx``, ``avro``, ``parquet``: file formats for
document pipelines.
* **Retrieval / SQL** — ``sql``, ``common.sql``, ``langchain``, ``llamaindex``: RAG and
Expand Down
14 changes: 11 additions & 3 deletions providers/common/ai/docs/toolsets.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,8 @@ Airflow's 350+ provider hooks already have typed methods, rich docstrings,
and managed credentials. Toolsets expose them as pydantic-ai tools so that
LLM agents can call them during multi-turn reasoning.

Four toolsets are included:
Four toolsets are exported directly from the ``airflow.providers.common.ai.toolsets``
package root:

- :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset` — generic
adapter for any Airflow Hook.
Expand All @@ -37,7 +38,14 @@ Four toolsets are included:
the agent a shell and a filesystem inside an isolated sandbox, off the
Airflow worker. See :ref:`which boundary that is <sandbox-boundaries>`.

All four implement pydantic-ai's
Three more toolsets are documented later on this page. They are not re-exported from the
package root, so import each of them from its own submodule::
Comment thread
Lee-W marked this conversation as resolved.

from airflow.providers.common.ai.toolsets.datafusion import DataFusionToolset
from airflow.providers.common.ai.toolsets.logging import LoggingToolset
from airflow.providers.common.ai.toolsets.skills import AgentSkillsToolset

All of the toolsets on this page implement pydantic-ai's
`AbstractToolset <https://ai.pydantic.dev/toolsets/>`__ interface and can be
passed to any pydantic-ai ``Agent``, including via
:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`.
Expand Down Expand Up @@ -842,7 +850,7 @@ Four rules for an implementation:
when it might. The first fails the task for Airflow to retry; the second
becomes a bounded prompt back to the model.
- If you cannot enforce something the ``SandboxSpec`` asks for, **raise**. Never
provision a weaker sandbox than the DAG author asked for.
provision a weaker sandbox than the Dag author asked for.

Parameters
^^^^^^^^^^
Expand Down