airflow.providers.common.ai.hooks.pydantic_ai¶

Attributes¶

OutputT

Classes¶

`PydanticAIHook`	Hook for LLM access via pydantic-ai.
`PydanticAIAzureHook`	Hook for Azure OpenAI via pydantic-ai.
`PydanticAIBedrockHook`	Hook for AWS Bedrock via pydantic-ai.
`PydanticAIVertexHook`	Hook for Google Vertex AI (or Generative Language API) 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=None, model_id=None, **kwargs)[source]¶

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

Hook for LLM access via pydantic-ai.

Covers providers that use a standard api_key + optional base_url (OpenAI, Anthropic, Groq, Mistral, DeepSeek, Ollama, vLLM, …).

For cloud providers with non-standard auth use the dedicated subclasses: PydanticAIAzureHook, PydanticAIBedrockHook, PydanticAIVertexHook.

Connection fields:

password: API key
host: Base URL (optional, e.g. https://api.openai.com/v1)
extra JSON: {"model": "openai:gpt-5.3"}

Parameters:

llm_conn_id (str | None) – 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[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.

Resolution order:

Explicit credentials — when _get_provider_kwargs() returns a non-empty dict the provider class is instantiated with those kwargs and wrapped in a provider_factory.
Default resolution — delegates to pydantic-ai infer_model which reads standard env vars (OPENAI_API_KEY, AWS_PROFILE, …).

The resolved model 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 and the provider class can be instantiated with the supplied credentials. Does NOT make an LLM API call — that would be expensive and fail for reasons unrelated to connectivity (quotas, billing, rate limits).

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

Bases: PydanticAIHook

Hook for Azure OpenAI via pydantic-ai.

Connection fields:

password: Azure API key
host: Azure endpoint (e.g. https://<resource>.openai.azure.com)

extra JSON:

{"model": "azure:gpt-4o", "api_version": "2024-07-01-preview"}

Parameters:

llm_conn_id (str | None) – Airflow connection ID.
model_id (str | None) – Model identifier, e.g. "azure:gpt-4o".

conn_type = 'pydanticai-azure'[source]¶

default_conn_name = 'pydanticai_azure_default'[source]¶

hook_name = 'Pydantic AI (Azure OpenAI)'[source]¶

static get_ui_field_behaviour()[source]¶

Return custom field behaviour for the Airflow connection form.

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

Bases: PydanticAIHook

Hook for AWS Bedrock via pydantic-ai.

Credentials are resolved in order:

IAM keys from extra (aws_access_key_id + aws_secret_access_key, optionally aws_session_token).
Bearer token in extra (api_key, maps to env AWS_BEARER_TOKEN_BEDROCK).
Environment-variable / instance-role chain (AWS_PROFILE, IAM role, …) when no explicit keys are provided.

Connection fields:

extra JSON:

{
  "model": "bedrock:us.anthropic.claude-opus-4-5",
  "region_name": "us-east-1",
  "aws_access_key_id": "AKIA...",
  "aws_secret_access_key": "...",
  "aws_session_token": "...",
  "profile_name": "my-aws-profile",
  "api_key": "bearer-token",
  "base_url": "https://custom-bedrock-endpoint",
  "aws_read_timeout": 60.0,
  "aws_connect_timeout": 10.0
}

Leave aws_access_key_id / aws_secret_access_key and api_key empty to use the default AWS credential chain.

Parameters:

llm_conn_id (str | None) – Airflow connection ID.
model_id (str | None) – Model identifier, e.g. "bedrock:us.anthropic.claude-opus-4-5".

conn_type = 'pydanticai-bedrock'[source]¶

default_conn_name = 'pydanticai_bedrock_default'[source]¶

hook_name = 'Pydantic AI (AWS Bedrock)'[source]¶

static get_ui_field_behaviour()[source]¶

Return custom field behaviour for the Airflow connection form.

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

Bases: PydanticAIHook

Hook for Google Vertex AI (or Generative Language API) via pydantic-ai.

Credentials are resolved in order:

service_account_info (JSON object) in extra — loaded into a google.auth.credentials.Credentials object and passed as credentials to GoogleProvider.
api_key in extra — for Generative Language API (non-Vertex) or Vertex API-key auth.
Application Default Credentials (GOOGLE_APPLICATION_CREDENTIALS, gcloud auth application-default login, Workload Identity, …) when no explicit credentials are provided.

Connection fields:

extra JSON:

{
    "model": "google-vertex:gemini-2.0-flash",
    "project": "my-gcp-project",
    "location": "us-central1",
    "service_account_info": {...},
    "vertexai": true,
}

Use "service_account_info" to embed the service-account JSON directly (as an object, not a string path).

Set "vertexai": true to force Vertex AI mode when only api_key is provided. Omit vertexai for the Generative Language API (GLA).

Parameters:

llm_conn_id (str | None) – Airflow connection ID.
model_id (str | None) – Model identifier, e.g. "google-vertex:gemini-2.0-flash".

conn_type = 'pydanticai-vertex'[source]¶

default_conn_name = 'pydanticai_vertex_default'[source]¶

hook_name = 'Pydantic AI (Google Vertex AI)'[source]¶

static get_ui_field_behaviour()[source]¶

Return custom field behaviour for the Airflow connection form.