Resolve ci failures

Author: gopidesupavan

docs/spelling_wordlist.txt

@@ -817,6 +817,7 @@ init
 initdb
 initialisation
 initialiser
+initialises
 initialising
 initializer
 inout

providers/common/ai/docs/changelog.rst

@@ -25,15 +25,6 @@
 Changelog
 ---------
 
-Next release
-............
-
-Features
-^^^^^^^^
-
-* Add ``BaseAIHook`` contract with framework-agnostic ``create_agent`` / ``run_agent`` /
-  ``get_model`` interface so ``AgentOperator`` selects the agent backend via connection type.
-
 0.3.0
 .....
 

providers/common/ai/provider.yaml

@@ -61,6 +61,9 @@ integrations:
     tags: [ai]
 
 hooks:
+  - integration-name: Base Hook
+    python-modules:
+      - airflow.providers.common.ai.hooks.base_ai
   - integration-name: Pydantic AI
     python-modules:
       - airflow.providers.common.ai.hooks.pydantic_ai

providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py

@@ -67,6 +67,10 @@ def get_provider_info():
             },
         ],
         "hooks": [
+            {
+                "integration-name": "Base Hook",
+                "python-modules": ["airflow.providers.common.ai.hooks.base_ai"],
+            },
             {
                 "integration-name": "Pydantic AI",
                 "python-modules": ["airflow.providers.common.ai.hooks.pydantic_ai"],

providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py

@@ -23,7 +23,7 @@
 import json
 import time
 from abc import ABCMeta, abstractmethod
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
 from typing import Any, ClassVar
 
@@ -123,7 +123,8 @@ class AgentRunRequest:
     framework-neutral structure, so that :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`
     has zero framework-specific imports.
 
-    :param prompt: User prompt for this invocation.
+    :param prompt: User prompt for this invocation (plain ``str`` or a multimodal
+        ``Sequence`` accepted by the backend agent's run API).
     :param output_type: Expected structured output type (default: ``str``).
     :param instructions: System-level instructions for the agent.
     :param toolsets: List of :class:`BaseToolset` instances the agent may call.
@@ -135,7 +136,7 @@ class AgentRunRequest:
         Use this escape hatch for framework-specific options.
     """
 
-    prompt: str
+    prompt: str | Sequence[Any]
     output_type: type[Any] = str
     instructions: str = ""
     toolsets: list[Any] | None = None

providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py

@@ -141,7 +141,7 @@ def get_model(self) -> Model:
         if self._model is not None:
             return self._model
 
-        conn = self.get_connection(self.llm_conn_id)
+        conn = self.get_connection(self.llm_conn_id or self.default_conn_name)
 
         extra: dict[str, Any] = conn.extra_dejson
         model_name: str | KnownModelName = self.model_id or extra.get("model", "")
@@ -271,7 +271,10 @@ def run_agent(self, agent: Agent[None, Any], request: AgentRunRequest) -> AgentR
             if storage is not None and counter is not None:
                 from airflow.providers.common.ai.durable.caching_model import CachingModel
 
-                resolved_model = infer_model(agent.model)
+                if agent.model is None:
+                    raise ValueError("Agent model must be set when durable=True")
+                model = agent.model
+                resolved_model = infer_model(model) if isinstance(model, str) else model
                 caching_model = CachingModel(
                     resolved_model,
                     storage=storage,

providers/common/ai/tests/unit/common/ai/decorators/test_agent.py

@@ -77,7 +77,6 @@ def test_execute_raises_on_invalid_prompt(self, return_value):
 
     def test_execute_accepts_sequence_prompt(self):
         """A non-empty Sequence[UserContent] return value is forwarded as-is."""
-        from pydantic_ai.messages import ImageUrl
 
         image = ImageUrl(url="https://example.com/x.png")
         prompt = ["Describe this:", image]
@@ -101,8 +100,6 @@ def test_sequence_prompt_with_hitl_review_raises(self):
         if not AIRFLOW_V_3_1_PLUS:
             pytest.skip("enable_hitl_review requires Airflow >= 3.1.0")
 
-        from pydantic_ai.messages import ImageUrl
-
         op = _AgentDecoratedOperator(
             task_id="test",
             python_callable=lambda: ["x", ImageUrl(url="https://example.com/x.png")],