From 1a56af2fd6290f8315ca83db0460a3e3fd35704a Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Wed, 19 Aug 2026 13:51:08 +0000 Subject: [PATCH] Add test_connection support to LangChainHook and LlamaIndexHook 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. --- .../providers/common/ai/hooks/langchain.py | 15 ++++++++ .../providers/common/ai/hooks/llamaindex.py | 15 ++++++++ .../unit/common/ai/hooks/test_langchain.py | 35 +++++++++++++++++++ .../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 81834a83764b3..b2f4e1a8c222a 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 @@ 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) 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 05e002d86425e..354e12afe4534 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 @@ 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) 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 646f72aa10779..e1af49eef9630 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 @@ 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.""" 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 9d6e71790b3e0..c91866822cec7 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 @@ 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