airflow.providers.common.ai.hooks.pydantic_ai

Attributes

OutputT

Classes

PydanticAIHook

Hook for LLM access via pydantic-ai.

Module Contents

airflow.providers.common.ai.hooks.pydantic_ai.OutputT[source]
class airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook(llm_conn_id=default_conn_name, model_id=None, **kwargs)[source]

Bases: airflow.providers.common.compat.sdk.BaseHook

Hook for LLM access via pydantic-ai.

Manages connection credentials and model creation. Uses pydantic-ai’s model inference to support any provider (OpenAI, Anthropic, Google, Bedrock, Ollama, vLLM, etc.).

Connection fields:

  • Model (conn-field): Model in provider:model format (e.g. "anthropic:claude-sonnet-4-20250514")

  • password: API key (OpenAI, Anthropic, Groq, Mistral, etc.)

  • host: Base URL (optional — for custom endpoints like Ollama, vLLM, Azure)

Cloud providers (Bedrock, Vertex) that use native auth chains should leave password empty and configure environment-based auth (AWS_PROFILE, GOOGLE_APPLICATION_CREDENTIALS).

Parameters:

  • llm_conn_id (str) – Airflow connection ID for the LLM provider.

  • model_id (str | None) – Model identifier in provider:model format (e.g. "openai:gpt-5.3"). Overrides the model stored in the connection’s extra field.

conn_name_attr = 'llm_conn_id'[source]
default_conn_name = 'pydanticai_default'[source]
conn_type = 'pydanticai'[source]
hook_name = 'Pydantic AI'[source]
llm_conn_id = 'pydanticai_default'[source]
model_id = None[source]
static get_ui_field_behaviour()[source]

Return custom field behaviour for the Airflow connection form.

get_conn()[source]

Return a configured pydantic-ai Model.

Reads API key from connection password, base_url from connection host, and model from (in priority order):

  1. model_id parameter on the hook

  2. extra["model"] on the connection (set by the “Model” conn-field in the UI)

The result is cached for the lifetime of this hook instance.

create_agent(output_type: type[OutputT], *, instructions: str, **agent_kwargs) → pydantic_ai.Agent[None, OutputT][source]
create_agent(*, instructions: str, **agent_kwargs) → pydantic_ai.Agent[None, str]

Create a pydantic-ai Agent configured with this hook’s model.

Parameters:

  • output_type – The expected output type from the agent (default: str).

  • instructions – System-level instructions for the agent.

  • agent_kwargs – Additional keyword arguments passed to the Agent constructor.

test_connection()[source]

Test connection by resolving the model.

Validates that the model string is valid, the provider package is installed, and the provider class can be instantiated. Does NOT make an LLM API call — that would be expensive, flaky, and fail for reasons unrelated to connectivity (quotas, billing, rate limits).

Was this entry helpful?