Skip to content

Document using toolsets directly with PydanticAI in any task#64234

Merged
kaxil merged 1 commit intoapache:mainfrom
astronomer:docs/toolsets-direct-pydanticai-usage
Mar 26, 2026
Merged

Document using toolsets directly with PydanticAI in any task#64234
kaxil merged 1 commit intoapache:mainfrom
astronomer:docs/toolsets-direct-pydanticai-usage

Conversation

@kaxil
Copy link
Member

@kaxil kaxil commented Mar 25, 2026

  • Add a "Using Toolsets Directly with PydanticAI" section to toolsets.rst documenting that toolsets work anywhere you can run Python within Airflow (@task, PythonOperator, custom operator execute()), not just via AgentOperator or @task.agent.
  • Add an example DAG showing @task + PydanticAIHook.create_agent() + SQLToolset.
  • Explain the tradeoffs: full agent lifecycle control vs losing durable execution, HITL review, and automatic tool logging.
image
@task
def analyze():
    import pydantic_ai

    agent = pydantic_ai.Agent(
        "anthropic:claude-sonnet-4-6",
        system_prompt=(
            "You are a sales analytics assistant. "
            "Use the available SQL tools to query the database and answer questions. "
            "Always start by listing tables and checking their schema before writing queries."
        ),
        toolsets=[SQLToolset(db_conn_id="sqlite_default")]
    )
    result = agent.run_sync(
        user_prompt=(
            "Which customers have spent the most in total? "
            "Show the top 5 with their total spend and number of orders. "
            "Also tell me which product category generates the most revenue."
        )
    )
    return result

Toolsets like SQLToolset are standard pydantic-ai AbstractToolset
implementations that work independently of AgentOperator. Document this
pattern with an example showing @task + PydanticAIHook + SQLToolset,
and explain the tradeoffs vs AgentOperator (no durable execution, HITL,
or automatic tool logging).

SQLToolset import is inside the task function because common-sql is an
optional dependency of common-ai -- a top-level import would break the
other example DAGs in the file at parse time.
@eladkal
Copy link
Contributor

eladkal commented Mar 26, 2026

@kaxil sent you a question in Slack

@kaxil kaxil merged commit 42d2424 into apache:main Mar 26, 2026
91 of 92 checks passed
@kaxil kaxil deleted the docs/toolsets-direct-pydanticai-usage branch March 26, 2026 01:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants