Add test_connection support to LangChainHook and LlamaIndexHook

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

@@ -171,3 +171,18 @@ def get_embedding_model(self) -> Embeddings:
             kind="embedding",
         )
         return init_embeddings(model_id, **self._connection_kwargs(conn))
+
+    def test_connection(self) -> tuple[bool, str]:
+        """
+        Test connection by resolving the chat model.
+
+        Validates that the model identifier is valid and the provider 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).
+        """
+        try:
+            self.get_chat_model()
+            return True, "Model resolved successfully."
+        except Exception as e:
+            return False, str(e)

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

@@ -187,3 +187,18 @@ def get_llm(self) -> LLM:
             kind="llm",
         )
         return OpenAI(model=model_id, **self._connection_kwargs(conn))
+
+    def test_connection(self) -> tuple[bool, str]:
+        """
+        Test connection by resolving the LLM.
+
+        Validates that the model identifier is valid and the provider 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).
+        """
+        try:
+            self.get_llm()
+            return True, "Model resolved successfully."
+        except Exception as e:
+            return False, str(e)

providers/common/ai/tests/unit/common/ai/hooks/test_langchain.py

@@ -271,6 +271,41 @@ def test_no_credentials_passes_empty_kwargs(self, mock_get_conn, mock_init_embed
         mock_init_embeddings.assert_called_once_with("openai:text-embedding-3-small")
 
 
+class TestConnectionTest:
+    @patch("langchain.chat_models.init_chat_model")
+    @patch.object(LangChainHook, "get_connection")
+    def test_successful_connection(self, mock_get_conn, mock_init_chat_model):
+        mock_get_conn.return_value = _conn(password="sk-test", extra={"model": "openai:gpt-4o"})
+
+        hook = LangChainHook()
+        success, message = hook.test_connection()
+
+        assert success is True
+        assert message == "Model resolved successfully."
+
+    @patch("langchain.chat_models.init_chat_model")
+    @patch.object(LangChainHook, "get_connection")
+    def test_failed_connection(self, mock_get_conn, mock_init_chat_model):
+        mock_get_conn.return_value = _conn(password="sk-test", extra={"model": "openai:gpt-4o"})
+        mock_init_chat_model.side_effect = ValueError("Unknown provider 'badprovider'")
+
+        hook = LangChainHook()
+        success, message = hook.test_connection()
+
+        assert success is False
+        assert "Unknown provider" in message
+
+    @patch.object(LangChainHook, "get_connection")
+    def test_failed_connection_no_model(self, mock_get_conn):
+        mock_get_conn.return_value = _conn()
+
+        hook = LangChainHook()
+        success, message = hook.test_connection()
+
+        assert success is False
+        assert "No chat model identifier set" in message
+
+
 class TestSameHookForBoth:
     """A single hook instance must serve both chat and embedding calls."""
 

providers/common/ai/tests/unit/common/ai/hooks/test_llamaindex.py

@@ -168,3 +168,38 @@ def test_raises_when_no_llm_model(self, mock_get_conn):
 
         with pytest.raises(ValueError, match="No llm model identifier set"):
             hook.get_llm()
+
+
+class TestConnectionTest:
+    @patch("llama_index.llms.openai.OpenAI")
+    @patch.object(LlamaIndexHook, "get_connection")
+    def test_successful_connection(self, mock_get_conn, mock_cls):
+        mock_get_conn.return_value = _conn(password="sk-test", extra={"llm_model": "gpt-4o"})
+
+        hook = LlamaIndexHook()
+        success, message = hook.test_connection()
+
+        assert success is True
+        assert message == "Model resolved successfully."
+
+    @patch("llama_index.llms.openai.OpenAI")
+    @patch.object(LlamaIndexHook, "get_connection")
+    def test_failed_connection(self, mock_get_conn, mock_cls):
+        mock_get_conn.return_value = _conn(password="sk-test", extra={"llm_model": "gpt-4o"})
+        mock_cls.side_effect = ValueError("Invalid API key")
+
+        hook = LlamaIndexHook()
+        success, message = hook.test_connection()
+
+        assert success is False
+        assert "Invalid API key" in message
+
+    @patch.object(LlamaIndexHook, "get_connection")
+    def test_failed_connection_no_model(self, mock_get_conn):
+        mock_get_conn.return_value = _conn()
+
+        hook = LlamaIndexHook()
+        success, message = hook.test_connection()
+
+        assert success is False
+        assert "No llm model identifier set" in message