kaxil commented on code in PR #74149:
URL: https://github.com/apache/airflow/pull/74149#discussion_r4247999778
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -207,23 +211,41 @@ def get_conn(self) -> OpenAI:
OpenAI client's workload identity support. See
:meth:`_build_workload_identity`.
"""
conn = self.get_connection(self.conn_id)
- extras = conn.extra_dejson
- openai_client_kwargs = extras.get("openai_client_kwargs", {})
+ return OpenAI(**self._client_kwargs(conn, conn.extra_dejson))
+
+ async def aget_conn(self) -> AsyncOpenAI:
+ """
+ Return an ``AsyncOpenAI`` client, for use inside an async task or
trigger.
+
+ Configured exactly as :meth:`get_conn` configures its client, but the
connection is
+ resolved without any synchronous call to the Task SDK.
``get_connection`` and the
Review Comment:
This only holds on Airflow 3.3.2+. `get_async_connection` takes
`aget_connection` whenever it exists (3.1+), and `_async_get_connection` still
masks synchronously on older cores: `_mask_connection_secrets` calls the sync
`mask_secret` on 3.1 and 3.2, and on 3.3.0/3.3.1 the secrets-backend branch
calls `conn.get_uri()`, which reads `extra_dejson`. Could the docstring and the
PR description say the guarantee starts at 3.3.2? Covering older cores would be
a common.compat change, not something for this PR.
##########
providers/openai/tests/unit/openai/hooks/test_openai.py:
##########
@@ -1032,3 +1033,199 @@ def
test_managed_agents_reports_sdk_upgrade_without_breaking_hook():
hook.conn = SimpleNamespace(beta=SimpleNamespace())
with pytest.raises(OpenAIAgentSessionError, match="requires
openai>=3.13.0"):
hook.create_agent_session(input="Hello", environment={"type": "none"},
agent_id="agent")
+
+
+def _async_conn(
+ extra: dict, password: str | None = None, host: str | None = None, *,
aextra_dejson: bool = True
+) -> Mock:
+ """
+ A connection whose ``extra_dejson`` raises on the event loop, where its
secret masking deadlocks.
+
+ With ``aextra_dejson=False`` it lacks ``Connection.aextra_dejson()``, like
on Airflow < 3.3.2,
+ where ``get_async_extra_dejson`` reads ``extra_dejson`` from a worker
thread instead.
+ """
+ conn = Mock(spec=["password", "host", "extra_dejson"] + (["aextra_dejson"]
if aextra_dejson else []))
+ conn.password = password
+ conn.host = host
+ if aextra_dejson:
+ conn.aextra_dejson = AsyncMock(return_value=extra)
+ type(conn).extra_dejson = PropertyMock(side_effect=lambda:
_read_off_the_event_loop(extra))
+ return conn
+
+
+def _read_off_the_event_loop(extra: dict) -> dict:
+ try:
+ asyncio.get_running_loop()
+ except RuntimeError:
+ return extra
+ raise AssertionError("extra_dejson sends mask_secret synchronously on the
event loop")
+
+
[email protected]
[email protected](
+ ("password", "extra", "expected_kwargs"),
+ [
+ pytest.param(
+ "api_key_in_password",
+ {},
+ {"api_key": "api_key_in_password", "base_url": "https://host/v1"},
+ id="api-key-in-password",
+ ),
+ pytest.param(
+ None,
+ {"openai_client_kwargs": {"api_key": "api_key_in_extra",
"base_url": "https://other/v1"}},
+ {"api_key": "api_key_in_extra", "base_url": "https://other/v1"},
+ id="api-key-and-base-url-in-extra",
+ ),
+ pytest.param(
+ "api_key",
+ {"openai_client_kwargs": {"timeout": 30}},
+ {"api_key": "api_key", "base_url": "https://host/v1";, "timeout":
30},
+ id="openai-client-kwargs",
+ ),
+ ],
+)
+async def test_aget_conn_api_key(password, extra, expected_kwargs):
+ conn = _async_conn(extra, password=password, host="https://host/v1";)
+ with (
+ patch(
+ "airflow.providers.openai.hooks.openai.get_async_connection",
new=AsyncMock(return_value=conn)
+ ) as mock_get_async_connection,
+ patch.object(OpenAIHook, "get_connection",
side_effect=AssertionError("sync get_connection")),
+ patch("airflow.providers.openai.hooks.openai.AsyncOpenAI") as
mock_client,
+ ):
+ hook = OpenAIHook(conn_id="openai_async")
+ client = await hook.aget_conn()
+
+ mock_get_async_connection.assert_awaited_once_with("openai_async",
hook=hook)
+ conn.aextra_dejson.assert_awaited_once_with()
+ mock_client.assert_called_once_with(**expected_kwargs)
+ assert client is mock_client.return_value
+
+
[email protected]
+@patch("airflow.providers.openai.hooks.openai.AsyncOpenAI")
+async def
test_aget_conn_without_aextra_dejson_reads_extra_dejson_off_the_event_loop(mock_client):
+ """On Airflow < 3.3.2 the connection has no ``aextra_dejson()``:
``extra_dejson`` runs in a worker thread."""
+ conn = _async_conn({"openai_client_kwargs": {"api_key":
"api_key_in_extra"}}, aextra_dejson=False)
+ with patch(
+ "airflow.providers.openai.hooks.openai.get_async_connection",
new=AsyncMock(return_value=conn)
+ ):
+ await OpenAIHook(conn_id="openai_async").aget_conn()
+
+ mock_client.assert_called_once_with(api_key="api_key_in_extra",
base_url=None)
+
+
[email protected]
+@patch("airflow.providers.openai.hooks.openai.AsyncOpenAI")
+async def test_aget_conn_workload_identity(mock_client):
+ extra = {
+ "auth_type": "workload_identity",
+ "identity_provider_id": "idp",
+ "service_account_id": "sa",
+ "workload_identity_provider": "custom",
+ "token_provider": "some.module.get_token",
+ "openai_client_kwargs": {"api_key": "ignored"},
+ }
+ conn = _async_conn(extra, host="https://host/v1";)
+ with (
+ patch("airflow.providers.openai.hooks.openai.get_async_connection",
new=AsyncMock(return_value=conn)),
+ patch.object(OpenAIHook, "_build_workload_identity",
return_value={"provider": "custom"}) as build,
+ ):
+ await OpenAIHook(conn_id="openai_async").aget_conn()
+
+ conn.aextra_dejson.assert_awaited_once_with()
+ build.assert_called_once_with(extra)
+ mock_client.assert_called_once_with(workload_identity={"provider":
"custom"}, base_url="https://host/v1";)
+
+
[email protected]
+@patch("airflow.providers.openai.hooks.openai.AsyncOpenAI")
+async def test_aget_conn_resolves_the_connection_through_the_hook(mock_client):
+ """A subclass overriding ``aget_connection`` supplies the connection, as
``get_connection`` does for ``get_conn``."""
+ conn = _async_conn({}, password="api_key_from_subclass")
+
+ class SubclassedOpenAIHook(OpenAIHook):
+ @classmethod
+ async def aget_connection(cls, conn_id: str):
+ assert conn_id == "openai_async"
+ return conn
+
+ with patch.object(OpenAIHook, "get_connection",
side_effect=AssertionError("sync get_connection")):
+ await SubclassedOpenAIHook(conn_id="openai_async").aget_conn()
+
+ conn.aextra_dejson.assert_awaited_once_with()
+ mock_client.assert_called_once_with(api_key="api_key_from_subclass",
base_url=None)
+
+
[email protected]
+async def test_aget_conn_invalid_auth_type():
+ conn = _async_conn({"auth_type": "magic"})
+ with (
+ patch("airflow.providers.openai.hooks.openai.get_async_connection",
new=AsyncMock(return_value=conn)),
+ pytest.raises(ValueError, match="Unsupported auth_type 'magic'"),
+ ):
+ await OpenAIHook(conn_id="openai_async").aget_conn()
+
+ conn.aextra_dejson.assert_awaited_once_with()
+
+
[email protected]
[email protected](
+ ("input_text", "response_items", "expected"),
+ [
+ pytest.param("Sample text", [(0, [0.1, 0.2])], [0.1, 0.2], id="text"),
+ pytest.param([1, 2, 3], [(0, [0.1, 0.2])], [0.1, 0.2], id="tokens"),
+ pytest.param(
+ ["First text", "Second text"],
+ [(1, [0.3, 0.4]), (0, [0.1, 0.2])],
+ [[0.1, 0.2], [0.3, 0.4]],
+ id="text-batch",
+ ),
+ ],
+)
+async def test_acreate_embeddings(input_text, response_items, expected):
+ client = _async_client(response_items)
+ with patch.object(OpenAIHook, "aget_conn",
new=AsyncMock(return_value=client)) as mock_aget_conn:
+ assert (
+ await
OpenAIHook(conn_id="openai_async").acreate_embeddings(input_text,
dimensions=1024)
+ == expected
+ )
+
+ mock_aget_conn.assert_awaited_once_with()
+ client.embeddings.create.assert_awaited_once_with(
+ model="text-embedding-3-small", input=input_text, dimensions=1024
+ )
+ client.__aexit__.assert_awaited_once()
+
+
[email protected]
+async def
test_acreate_embeddings_concurrent_calls_each_close_their_own_client():
+ """Concurrent calls on one hook each build a client and close it; none is
shared or left open."""
+ clients = [_async_client([(0, [0.1, 0.2])]) for _ in range(3)]
+ hook = OpenAIHook(conn_id="openai_async")
+ with patch.object(OpenAIHook, "aget_conn",
new=AsyncMock(side_effect=clients)):
Review Comment:
Since `aget_conn` itself is patched with `side_effect=clients`, the test
supplies the distinct clients, so it would still pass if `aget_conn` went back
to caching one client on the hook. Patching `get_async_connection` and
`AsyncOpenAI` (`side_effect=clients`) instead and asserting three constructions
would pin the per-call behaviour.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
