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.

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)
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):

model_id parameter on the hook
model key in the connection’s extra JSON
(No default — raises an error if neither is set)

# 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-opus-4-6")

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)
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")