Using the hook directly: PydanticAIHook

Use PydanticAIHook to interact with LLM providers via pydantic-ai.

The hook manages API credentials from an Airflow connection and creates pydantic-ai Model and Agent objects. It supports any provider that pydantic-ai supports.

See also

Connection configuration

Basic usage

Use the hook in a @task function to call an LLM:

airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py[source]

@dag(schedule=None, tags=["example"])
def example_pydantic_ai_hook():
    @task
    def generate_summary(text: str) -> str:
        hook = PydanticAIHook(llm_conn_id="pydanticai_default")
        agent = hook.create_agent(output_type=str, instructions="Summarize concisely.")
        result = agent.run_sync(text)
        return result.output

    generate_summary("Apache Airflow is a platform for programmatically authoring...")

Overriding the model

The model can be specified at three levels (highest priority first):

  1. model_id parameter on the hook

  2. model key in the connection’s extra JSON

  3. With create_agent(spec_file=...), the spec file’s model. Otherwise there is no default and the hook raises.

# Use model from the connection's extra JSON
hook = PydanticAIHook(llm_conn_id="my_llm")

# Override with a specific model
hook = PydanticAIHook(llm_conn_id="my_llm", model_id="anthropic:claude-sonnet-5")

Structured output

Pydantic-ai’s structured output works naturally through the hook. Define a Pydantic model for the expected output shape, then pass it as output_type:

airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py[source]

@dag(schedule=None, tags=["example"])
def example_pydantic_ai_structured_output():
    @task
    def generate_sql(prompt: str) -> dict:
        class SQLResult(BaseModel):
            query: str
            explanation: str

        hook = PydanticAIHook(llm_conn_id="pydanticai_default")
        agent = hook.create_agent(
            output_type=SQLResult,
            instructions="Generate a SQL query and explain it.",
        )
        result = agent.run_sync(prompt)
        return result.output.model_dump()

    generate_sql("Find the top 10 customers by revenue")

Loading agent config from a spec file

Instead of hard-coding model name, instructions, and settings in Python, you can store them in a YAML or JSON AgentSpec file and pass its path via spec_file. This keeps prompt engineering separate from Dag logic and lets you version-control agent configs independently.

agent_spec.yaml
model: openai:gpt-5-mini
instructions: >
  You are a concise summarizer. Given any text, respond with a single
  paragraph that captures the key points.
model_settings:
  temperature: 0.3
retries: 2

airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py[source]

@dag(schedule=None, tags=["example"])
def example_pydantic_ai_spec_file():
    """Load agent settings from a YAML spec file instead of inline code.

    The spec file (``example_agent_spec.yaml``) declares model, instructions,
    model_settings, retries, etc. If ``model_id`` or the connection's ``model``
    extra is set, that hook model takes precedence over the file's model.
    """

    @task
    def summarize_from_spec(text: str) -> str:
        spec_path = Path(__file__).parent / "example_agent_spec.yaml"
        hook = PydanticAIHook(llm_conn_id="pydanticai_default")
        # Model, instructions, temperature, and retries all come from the YAML file.
        agent = hook.create_agent(spec_file=spec_path)
        result = agent.run_sync(text)
        return result.output

    @task
    def summarize_with_additional_instructions(text: str) -> str:
        """Add call-time instructions alongside the spec file instructions."""
        spec_path = Path(__file__).parent / "example_agent_spec.yaml"
        hook = PydanticAIHook(llm_conn_id="pydanticai_default")
        agent = hook.create_agent(
            spec_file=spec_path,
            instructions="Summarize in exactly one sentence.",
        )
        result = agent.run_sync(text)
        return result.output

    body = "Apache Airflow is an open-source platform for authoring, scheduling..."
    summarize_from_spec(body)
    summarize_with_additional_instructions(body)

The spec file’s model is the last resort: model_id and the connection’s model extra both take precedence over it, in the order the connection page lists. Passing instructions to create_agent when a spec_file is also given appends additional instructions to the file value.

Using the hook in an async task

In an async def task on Airflow 3.2.1 and later, use the async counterparts aget_hook, aget_conn and acreate_agent, and run the agent with await agent.run(...). They take the same parameters as get_hook, get_conn and create_agent and return the same objects, fallback chain included. Connection lookups make synchronous calls to the Task SDK, which block the event loop while another async call is in flight (from 3.3.1 they raise DeadlockImminentError instead), so the async methods run them on a worker thread of their own.

@task
async def summarize(text: str) -> str:
    hook = await PydanticAIHook.aget_hook("my_llm")
    agent = await hook.acreate_agent(instructions="Summarize the text in one sentence.")
    result = await agent.run(text)
    return result.output

Was this entry helpful?