added azureopenai support & extandable architecture for pydanticai

providers/common/ai/docs/operators/llm_schema_compare.rst

@@ -91,7 +91,7 @@ on top, concatenate them:
     LLMSchemaCompareOperator(
         task_id="compare_with_custom_rules",
         prompt="Compare schemas and flag breaking changes",
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         db_conn_ids=["postgres_source", "snowflake_target"],
         table_names=["customers"],
         system_prompt=DEFAULT_SYSTEM_PROMPT

providers/common/ai/docs/toolsets.rst

@@ -256,7 +256,7 @@ Using Multiple MCP Servers
     AgentOperator(
         task_id="multi_mcp",
         prompt="Get the weather in London and run a calculation",
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         toolsets=[
             MCPToolset(mcp_conn_id="weather_mcp", tool_prefix="weather"),
             MCPToolset(mcp_conn_id="code_runner_mcp", tool_prefix="code"),
@@ -276,7 +276,7 @@ server instances directly — no Airflow connection needed:
     AgentOperator(
         task_id="direct_mcp",
         prompt="What tools are available?",
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         toolsets=[
             MCPServerStreamableHTTP("http://localhost:3001/mcp"),
             MCPServerStdio("uvx", args=["mcp-run-python"]),

providers/common/ai/provider.yaml

@@ -79,6 +79,176 @@ connection-types:
           type:
             - string
             - 'null'
+  - hook-class-name: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook
+    connection-type: pydanticai-azure
+    ui-field-behaviour:
+      hidden-fields:
+        - schema
+        - port
+        - login
+      relabeling:
+        password: API Key
+        host: Azure Endpoint
+      placeholders:
+        host: "https://<resource>.openai.azure.com"
+    conn-fields:
+      model:
+        label: Model
+        description: "Azure model identifier (e.g. azure:gpt-4o)"
+        schema:
+          type:
+            - string
+            - 'null'
+      api_version:
+        label: API Version
+        description: "Azure OpenAI API version (e.g. 2024-07-01-preview). Falls back to OPENAI_API_VERSION."
+        schema:
+          type:
+            - string
+            - 'null'
+  - hook-class-name: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIBedrockHook
+    connection-type: pydanticai-bedrock
+    ui-field-behaviour:
+      hidden-fields:
+        - schema
+        - port
+        - login
+        - host
+        - password
+      relabeling: {}
+      placeholders: {}
+    conn-fields:
+      model:
+        label: Model
+        description: "Bedrock model identifier (e.g. bedrock:us.anthropic.claude-opus-4-5)"
+        schema:
+          type:
+            - string
+            - 'null'
+      region_name:
+        label: AWS Region
+        description: "AWS region (e.g. us-east-1). Falls back to AWS_DEFAULT_REGION env var."
+        schema:
+          type:
+            - string
+            - 'null'
+      aws_access_key_id:
+        label: AWS Access Key ID
+        description: "IAM access key. Leave empty to use instance role / environment credential chain."
+        schema:
+          type:
+            - string
+            - 'null'
+      aws_secret_access_key:
+        label: AWS Secret Access Key
+        description: "IAM secret key."
+        schema:
+          type:
+            - string
+            - 'null'
+      aws_session_token:
+        label: AWS Session Token
+        description: "Temporary session token (optional)."
+        schema:
+          type:
+            - string
+            - 'null'
+      profile_name:
+        label: AWS Profile Name
+        description: "Named AWS credentials profile (optional)."
+        schema:
+          type:
+            - string
+            - 'null'
+      api_key:
+        label: Bearer Token
+        description: "AWS bearer token (alt. to IAM key/secret). Falls back to AWS_BEARER_TOKEN_BEDROCK."
+        schema:
+          type:
+            - string
+            - 'null'
+      base_url:
+        label: Custom Endpoint URL
+        description: "Override the Bedrock runtime endpoint URL (optional)."
+        schema:
+          type:
+            - string
+            - 'null'
+      aws_read_timeout:
+        label: Read Timeout (s)
+        description: "boto3 read timeout in seconds (float, optional)."
+        schema:
+          type:
+            - number
+            - 'null'
+      aws_connect_timeout:
+        label: Connect Timeout (s)
+        description: "boto3 connect timeout in seconds (float, optional)."
+        schema:
+          type:
+            - number
+            - 'null'
+  - hook-class-name: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIVertexHook
+    connection-type: pydanticai-vertex
+    ui-field-behaviour:
+      hidden-fields:
+        - schema
+        - port
+        - login
+        - host
+        - password
+      relabeling: {}
+      placeholders: {}
+    conn-fields:
+      model:
+        label: Model
+        description: "Google model identifier (e.g. google-vertex:gemini-2.0-flash)"
+        schema:
+          type:
+            - string
+            - 'null'
+      project:
+        label: GCP Project
+        description: "Google Cloud project ID. Falls back to GOOGLE_CLOUD_PROJECT env var."
+        schema:
+          type:
+            - string
+            - 'null'
+      location:
+        label: Location / Region
+        description: "Vertex AI region (e.g. us-central1). Falls back to GOOGLE_CLOUD_LOCATION env var."
+        schema:
+          type:
+            - string
+            - 'null'
+      vertexai:
+        label: Force Vertex AI Mode
+        description: "Force Vertex AI mode. Auto-detected when project/location/credentials are set."
+        schema:
+          type:
+            - boolean
+            - 'null'
+      api_key:
+        label: API Key
+        description: "Google API key for Gen Language API or Vertex AI. Falls back to GOOGLE_API_KEY."
+        schema:
+          type:
+            - string
+            - 'null'
+      service_account_info:
+        label: Service Account Info
+        description: "Service account key as inline dict (JSON with type, project_id, private_key, etc.)."
+        schema:
+          type:
+            - object
+            - 'null'
+      base_url:
+        label: Custom Endpoint URL
+        description: "Override the Google API base URL (optional)."
+        schema:
+          type:
+            - string
+            - 'null'
   - hook-class-name: airflow.providers.common.ai.hooks.mcp.MCPHook
     connection-type: mcp
     ui-field-behaviour:

providers/common/ai/src/airflow/providers/common/ai/decorators/agent.py

@@ -107,7 +107,7 @@ def agent_task(
     Usage::
 
         @task.agent(
-            llm_conn_id="pydantic_ai_default",
+            llm_conn_id="pydanticai_default",
             system_prompt="You are a data analyst.",
             toolsets=[SQLToolset(db_conn_id="postgres_default")],
         )

providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent.py

@@ -36,7 +36,7 @@ def example_agent_operator_sql():
     AgentOperator(
         task_id="analyst",
         prompt="What are the top 5 customers by order count?",
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         system_prompt=(
             "You are a SQL analyst. Use the available tools to explore "
             "the schema and answer the question with data."
@@ -71,7 +71,7 @@ def example_agent_operator_hook():
     AgentOperator(
         task_id="api_explorer",
         prompt="What endpoints are available and what does /status return?",
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         system_prompt="You are an API explorer. Use the tools to discover and call endpoints.",
         toolsets=[
             HookToolset(
@@ -97,7 +97,7 @@ def example_agent_operator_hook():
 @dag
 def example_agent_decorator():
     @task.agent(
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         system_prompt="You are a data analyst. Use tools to answer questions.",
         toolsets=[
             SQLToolset(
@@ -133,7 +133,7 @@ class Analysis(BaseModel):
         row_count: int
 
     @task.agent(
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         system_prompt="You are a data analyst. Return structured results.",
         output_type=Analysis,
         toolsets=[SQLToolset(db_conn_id="postgres_default")],
@@ -158,7 +158,7 @@ def analyze(question: str):
 @dag
 def example_agent_chain():
     @task.agent(
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         system_prompt="You are a SQL analyst.",
         toolsets=[SQLToolset(db_conn_id="postgres_default", allowed_tables=["orders"])],
     )

providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_schema_compare.py

@@ -29,7 +29,7 @@ def example_llm_schema_compare_basic():
     LLMSchemaCompareOperator(
         task_id="detect_schema_drift",
         prompt="Identify schema mismatches that would break data loading between systems",
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         db_conn_ids=["postgres_default", "snowflake_default"],
         table_names=["customers"],
     )
@@ -49,7 +49,7 @@ def example_llm_schema_compare_full_context():
             "Compare schemas and generate a migration plan. "
             "Flag any differences that would break nightly ETL loads."
         ),
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         db_conn_ids=["postgres_source", "snowflake_target"],
         table_names=["customers", "orders"],
         context_strategy="full",
@@ -74,7 +74,7 @@ def example_llm_schema_compare_with_object_storage():
     LLMSchemaCompareOperator(
         task_id="compare_s3_vs_db",
         prompt="Compare S3 Parquet schema against the Postgres table and flag breaking changes",
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         db_conn_ids=["postgres_default"],
         table_names=["customers"],
         data_sources=[s3_source],
@@ -90,7 +90,7 @@ def example_llm_schema_compare_with_object_storage():
 @dag
 def example_llm_schema_compare_decorator():
     @task.llm_schema_compare(
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         db_conn_ids=["postgres_source", "snowflake_target"],
         table_names=["customers"],
     )
@@ -109,7 +109,7 @@ def check_migration_readiness(ds=None):
 @dag
 def example_llm_schema_compare_conditional():
     @task.llm_schema_compare(
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         db_conn_ids=["postgres_source", "snowflake_target"],
         table_names=["customers"],
         context_strategy="full",

providers/common/ai/src/airflow/providers/common/ai/example_dags/example_mcp.py

@@ -34,7 +34,7 @@ def example_mcp_toolset():
     AgentOperator(
         task_id="mcp_agent",
         prompt="What tools are available? Run the hello tool.",
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         system_prompt="You are a helpful assistant with access to MCP tools.",
         toolsets=[
             MCPToolset(mcp_conn_id="my_mcp_server"),
@@ -59,7 +59,7 @@ def example_mcp_multiple_servers():
     AgentOperator(
         task_id="multi_mcp_agent",
         prompt="Get the weather in London and run a Python calculation: 2**10",
-        llm_conn_id="pydantic_ai_default",
+        llm_conn_id="pydanticai_default",
         system_prompt="You have access to weather and code execution tools.",
         toolsets=[
             MCPToolset(mcp_conn_id="weather_mcp", tool_prefix="weather"),
@@ -84,7 +84,7 @@ def example_mcp_multiple_servers():
 #   AgentOperator(
 #       task_id="direct_mcp",
 #       prompt="What tools are available?",
-#       llm_conn_id="pydantic_ai_default",
+#       llm_conn_id="pydanticai_default",
 #       toolsets=[
 #           MCPServerStreamableHTTP("http://localhost:3001/mcp"),
 #           MCPServerStdio("uvx", args=["mcp-run-python"]),