Fix common.ai example DAGs failing to parse without the sql extra

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

@@ -24,9 +24,13 @@
 
 from airflow.providers.common.ai.operators.agent import AgentOperator
 from airflow.providers.common.ai.toolsets.hook import HookToolset
-from airflow.providers.common.ai.toolsets.sql import SQLToolset
 from airflow.providers.common.compat.sdk import dag, task
 
+try:
+    from airflow.providers.common.ai.toolsets.sql import SQLToolset
+except Exception:
+    SQLToolset = None  # type: ignore[assignment,misc]
+
 
 # [START howto_decorator_agent_structured_output_class]
 # Pydantic output classes must be defined at module scope so downstream
@@ -48,29 +52,30 @@ class Analysis(BaseModel):
 
 
 # [START howto_operator_agent_sql]
-@dag(tags=["example"])
-def example_agent_operator_sql():
-    AgentOperator(
-        task_id="analyst",
-        prompt="What are the top 5 customers by order count?",
-        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."
-        ),
-        toolsets=[
-            SQLToolset(
-                db_conn_id="postgres_default",
-                allowed_tables=["customers", "orders"],
-                max_rows=20,
-            )
-        ],
-    )
-
-
-# [END howto_operator_agent_sql]
-
-example_agent_operator_sql()
+if SQLToolset is not None:
+
+    @dag(tags=["example"])
+    def example_agent_operator_sql():
+        AgentOperator(
+            task_id="analyst",
+            prompt="What are the top 5 customers by order count?",
+            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."
+            ),
+            toolsets=[
+                SQLToolset(
+                    db_conn_id="postgres_default",
+                    allowed_tables=["customers", "orders"],
+                    max_rows=20,
+                )
+            ],
+        )
+
+    # [END howto_operator_agent_sql]
+
+    example_agent_operator_sql()
 
 
 # ---------------------------------------------------------------------------
@@ -111,27 +116,28 @@ def example_agent_operator_hook():
 
 
 # [START howto_decorator_agent]
-@dag(tags=["example"])
-def example_agent_decorator():
-    @task.agent(
-        llm_conn_id="pydanticai_default",
-        system_prompt="You are a data analyst. Use tools to answer questions.",
-        toolsets=[
-            SQLToolset(
-                db_conn_id="postgres_default",
-                allowed_tables=["orders"],
-            )
-        ],
-    )
-    def analyze(question: str):
-        return f"Answer this question about our orders data: {question}"
+if SQLToolset is not None:
 
-    analyze("What was our total revenue last month?")
+    @dag(tags=["example"])
+    def example_agent_decorator():
+        @task.agent(
+            llm_conn_id="pydanticai_default",
+            system_prompt="You are a data analyst. Use tools to answer questions.",
+            toolsets=[
+                SQLToolset(
+                    db_conn_id="postgres_default",
+                    allowed_tables=["orders"],
+                )
+            ],
+        )
+        def analyze(question: str):
+            return f"Answer this question about our orders data: {question}"
 
+        analyze("What was our total revenue last month?")
 
-# [END howto_decorator_agent]
+    # [END howto_decorator_agent]
 
-example_agent_decorator()
+    example_agent_decorator()
 
 
 # ---------------------------------------------------------------------------
@@ -140,23 +146,24 @@ def analyze(question: str):
 
 
 # [START howto_decorator_agent_structured]
-@dag(tags=["example"])
-def example_agent_structured_output():
-    @task.agent(
-        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")],
-    )
-    def analyze(question: str):
-        return f"Analyze: {question}"
+if SQLToolset is not None:
 
-    analyze("What are the trending products this week?")
+    @dag(tags=["example"])
+    def example_agent_structured_output():
+        @task.agent(
+            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")],
+        )
+        def analyze(question: str):
+            return f"Analyze: {question}"
 
+        analyze("What are the trending products this week?")
 
-# [END howto_decorator_agent_structured]
+    # [END howto_decorator_agent_structured]
 
-example_agent_structured_output()
+    example_agent_structured_output()
 
 
 # ---------------------------------------------------------------------------
@@ -165,29 +172,30 @@ def analyze(question: str):
 
 
 # [START howto_agent_chain]
-@dag(tags=["example"])
-def example_agent_chain():
-    @task.agent(
-        llm_conn_id="pydanticai_default",
-        system_prompt="You are a SQL analyst.",
-        toolsets=[SQLToolset(db_conn_id="postgres_default", allowed_tables=["orders"])],
-    )
-    def investigate(question: str):
-        return f"Investigate: {question}"
-
-    @task
-    def send_report(analysis: str):
-        """Send the agent's analysis to a downstream system."""
-        print(f"Report: {analysis}")
-        return analysis
-
-    result = investigate("Summarize order trends for last quarter")
-    send_report(result)
-
-
-# [END howto_agent_chain]
-
-example_agent_chain()
+if SQLToolset is not None:
+
+    @dag(tags=["example"])
+    def example_agent_chain():
+        @task.agent(
+            llm_conn_id="pydanticai_default",
+            system_prompt="You are a SQL analyst.",
+            toolsets=[SQLToolset(db_conn_id="postgres_default", allowed_tables=["orders"])],
+        )
+        def investigate(question: str):
+            return f"Investigate: {question}"
+
+        @task
+        def send_report(analysis: str):
+            """Send the agent's analysis to a downstream system."""
+            print(f"Report: {analysis}")
+            return analysis
+
+        result = investigate("Summarize order trends for last quarter")
+        send_report(result)
+
+    # [END howto_agent_chain]
+
+    example_agent_chain()
 
 
 # ---------------------------------------------------------------------------

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

@@ -28,9 +28,13 @@
 from pydantic_ai.capabilities import Thinking, WebSearch
 
 from airflow.providers.common.ai.operators.agent import AgentOperator
-from airflow.providers.common.ai.toolsets.sql import SQLToolset
 from airflow.providers.common.compat.sdk import dag
 
+try:
+    from airflow.providers.common.ai.toolsets.sql import SQLToolset
+except Exception:
+    SQLToolset = None  # type: ignore[assignment,misc]
+
 # ---------------------------------------------------------------------------
 # 1. Thinking capability: enable model reasoning at a configurable effort level
 # ---------------------------------------------------------------------------
@@ -86,29 +90,30 @@ def example_agent_capabilities_web_search():
 
 
 # [START howto_operator_agent_capabilities_composed]
-@dag(tags=["example"])
-def example_agent_capabilities_composed():
-    AgentOperator(
-        task_id="analyst",
-        prompt="Cross-reference our top customers with their recent public news. Think first.",
-        llm_conn_id="pydanticai_default",
-        system_prompt=(
-            "You are a sales analyst. Query the database for customers, then search the web "
-            "for recent news. Reason carefully about which leads to surface."
-        ),
-        toolsets=[
-            SQLToolset(
-                db_conn_id="postgres_default",
-                allowed_tables=["customers", "orders"],
-                max_rows=20,
+if SQLToolset is not None:
+
+    @dag(tags=["example"])
+    def example_agent_capabilities_composed():
+        AgentOperator(
+            task_id="analyst",
+            prompt="Cross-reference our top customers with their recent public news. Think first.",
+            llm_conn_id="pydanticai_default",
+            system_prompt=(
+                "You are a sales analyst. Query the database for customers, then search the web "
+                "for recent news. Reason carefully about which leads to surface."
             ),
-        ],
-        agent_params={
-            "capabilities": [Thinking(effort="medium"), WebSearch()],
-        },
-    )
-
-
-# [END howto_operator_agent_capabilities_composed]
-
-example_agent_capabilities_composed()
+            toolsets=[
+                SQLToolset(
+                    db_conn_id="postgres_default",
+                    allowed_tables=["customers", "orders"],
+                    max_rows=20,
+                ),
+            ],
+            agent_params={
+                "capabilities": [Thinking(effort="medium"), WebSearch()],
+            },
+        )
+
+    # [END howto_operator_agent_capabilities_composed]
+
+    example_agent_capabilities_composed()

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

@@ -21,40 +21,45 @@
 from datetime import timedelta
 
 from airflow.providers.common.ai.operators.agent import AgentOperator
-from airflow.providers.common.ai.toolsets.sql import SQLToolset
 from airflow.providers.common.compat.sdk import dag, task
 
+try:
+    from airflow.providers.common.ai.toolsets.sql import SQLToolset
+except Exception:
+    SQLToolset = None  # type: ignore[assignment,misc]
+
 # ---------------------------------------------------------------------------
 # 1. Durable AgentOperator: resumes from last model call on retry
 # ---------------------------------------------------------------------------
 
 
 # [START howto_operator_agent_durable]
-@dag(default_args={"retries": 3, "retry_delay": timedelta(seconds=30)}, tags=["example"])
-def example_agent_durable_operator():
-    """Agent with durable execution -- resumes from the last model call on retry."""
-    AgentOperator(
-        task_id="durable_analyst",
-        prompt="What are the top 5 customers by order count?",
-        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."
-        ),
-        durable=True,
-        toolsets=[
-            SQLToolset(
-                db_conn_id="postgres_default",
-                allowed_tables=["customers", "orders"],
-                max_rows=20,
-            )
-        ],
-    )
-
-
-# [END howto_operator_agent_durable]
-
-example_agent_durable_operator()
+if SQLToolset is not None:
+
+    @dag(default_args={"retries": 3, "retry_delay": timedelta(seconds=30)}, tags=["example"])
+    def example_agent_durable_operator():
+        """Agent with durable execution -- resumes from the last model call on retry."""
+        AgentOperator(
+            task_id="durable_analyst",
+            prompt="What are the top 5 customers by order count?",
+            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."
+            ),
+            durable=True,
+            toolsets=[
+                SQLToolset(
+                    db_conn_id="postgres_default",
+                    allowed_tables=["customers", "orders"],
+                    max_rows=20,
+                )
+            ],
+        )
+
+    # [END howto_operator_agent_durable]
+
+    example_agent_durable_operator()
 
 
 # ---------------------------------------------------------------------------
@@ -63,25 +68,26 @@ def example_agent_durable_operator():
 
 
 # [START howto_decorator_agent_durable]
-@dag(default_args={"retries": 3, "retry_delay": timedelta(seconds=30)}, tags=["example"])
-def example_agent_durable_decorator():
-    @task.agent(
-        llm_conn_id="pydanticai_default",
-        system_prompt="You are a data analyst. Use tools to answer questions.",
-        durable=True,
-        toolsets=[
-            SQLToolset(
-                db_conn_id="postgres_default",
-                allowed_tables=["orders"],
-            )
-        ],
-    )
-    def analyze(question: str):
-        return f"Answer this question about our orders data: {question}"
-
-    analyze("What was our total revenue last month?")
-
-
-# [END howto_decorator_agent_durable]
-
-example_agent_durable_decorator()
+if SQLToolset is not None:
+
+    @dag(default_args={"retries": 3, "retry_delay": timedelta(seconds=30)}, tags=["example"])
+    def example_agent_durable_decorator():
+        @task.agent(
+            llm_conn_id="pydanticai_default",
+            system_prompt="You are a data analyst. Use tools to answer questions.",
+            durable=True,
+            toolsets=[
+                SQLToolset(
+                    db_conn_id="postgres_default",
+                    allowed_tables=["orders"],
+                )
+            ],
+        )
+        def analyze(question: str):
+            return f"Answer this question about our orders data: {question}"
+
+        analyze("What was our total revenue last month?")
+
+    # [END howto_decorator_agent_durable]
+
+    example_agent_durable_decorator()