ContextQMD
Back to Airflow

PydanticAIHook

providers/common/ai/docs/hooks/pydantic_ai.rst

3.2.12.5 KB
Original Source

.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at

.. http://www.apache.org/licenses/LICENSE-2.0

.. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

.. _howto/hook:pydantic_ai:

PydanticAIHook

Use :class:~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook to interact with LLM providers via pydantic-ai <https://ai.pydantic.dev/>__.

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.

.. seealso:: :ref:Connection configuration <howto/connection:pydanticai>

Basic Usage

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

.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py :language: python :start-after: [START howto_hook_pydantic_ai_basic] :end-before: [END howto_hook_pydantic_ai_basic]

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. (No default — raises an error if neither is set)

.. code-block:: python

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

.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py :language: python :start-after: [START howto_hook_pydantic_ai_structured_output] :end-before: [END howto_hook_pydantic_ai_structured_output]