This is an automated email from the ASF dual-hosted git repository.
gopidesupavan pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 7991cd81c77 Add test_connection support to LangChainHook and
LlamaIndexHook (#71841)
7991cd81c77 is described below
commit 7991cd81c7713c405d2d69248fa4eb5d078c649b
Author: Jyun-An Chen <[email protected]>
AuthorDate: Sun Aug 23 18:52:58 2026 +0800
Add test_connection support to LangChainHook and LlamaIndexHook (#71841)
Clicking Test on a LangChain or LlamaIndex connection in the UI always
reported "doesn't implement or inherit test_connection method" since
neither hook implemented it, unlike PydanticAIHook and MCPHook which
already validate their connections this way.
---
.../airflow/providers/common/ai/hooks/langchain.py | 15 ++++++++++
.../providers/common/ai/hooks/llamaindex.py | 15 ++++++++++
.../tests/unit/common/ai/hooks/test_langchain.py | 35 ++++++++++++++++++++++
.../tests/unit/common/ai/hooks/test_llamaindex.py | 35 ++++++++++++++++++++++
4 files changed, 100 insertions(+)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py
b/providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py
index 81834a83764..b2f4e1a8c22 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py
@@ -171,3 +171,18 @@ class LangChainHook(BaseHook):
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)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/hooks/llamaindex.py
b/providers/common/ai/src/airflow/providers/common/ai/hooks/llamaindex.py
index 05e002d8642..354e12afe45 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/hooks/llamaindex.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/llamaindex.py
@@ -187,3 +187,18 @@ class LlamaIndexHook(BaseHook):
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)
diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_langchain.py
b/providers/common/ai/tests/unit/common/ai/hooks/test_langchain.py
index 646f72aa107..e1af49eef96 100644
--- a/providers/common/ai/tests/unit/common/ai/hooks/test_langchain.py
+++ b/providers/common/ai/tests/unit/common/ai/hooks/test_langchain.py
@@ -271,6 +271,41 @@ class TestGetEmbeddingModel:
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."""
diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_llamaindex.py
b/providers/common/ai/tests/unit/common/ai/hooks/test_llamaindex.py
index 9d6e71790b3..c91866822ce 100644
--- a/providers/common/ai/tests/unit/common/ai/hooks/test_llamaindex.py
+++ b/providers/common/ai/tests/unit/common/ai/hooks/test_llamaindex.py
@@ -168,3 +168,38 @@ class TestGetLlm:
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
