dabla commented on code in PR #74149:
URL: https://github.com/apache/airflow/pull/74149#discussion_r4222917802
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -201,23 +205,44 @@ 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
+ secret masking of ``extra_dejson`` send to the supervisor
synchronously, which raises
+ ``DeadlockImminentError`` on an event loop with another async call in
flight.
+ """
+ conn = await get_async_connection(self.conn_id)
Review Comment:
Fixed in 5834e14dd1: `get_async_connection(self.conn_id, hook=self)`, plus a
test with a subclass overriding `aget_connection()` that goes through the real
common.compat helper (it failed before the change: the lookup went to
`BaseHook`).
Drafted-by: Claude Fable 5.1; reviewed by @dabla before posting
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -201,23 +205,44 @@ 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
+ secret masking of ``extra_dejson`` send to the supervisor
synchronously, which raises
+ ``DeadlockImminentError`` on an event loop with another async call in
flight.
+ """
+ conn = await get_async_connection(self.conn_id)
+ return AsyncOpenAI(**self._client_kwargs(conn, await
get_async_extra_dejson(conn)))
+
+ async def _aget_cached_conn(self) -> AsyncOpenAI:
+ """Return the hook's ``AsyncOpenAI`` client, created on first use like
:attr:`conn`."""
+ if self._async_conn is None:
Review Comment:
Fixed in 47f89bb45b: the cache is gone. `aget_conn()` builds a client per
call and its docstring says to close it, like `AnthropicHook.get_async_conn`;
`acreate_embeddings()` uses `async with await self.aget_conn() as client:`.
Added a test with three concurrent calls on one hook, each getting and closing
its own client.
Drafted-by: Claude Fable 5.1; reviewed by @dabla before posting
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -566,6 +591,43 @@ def create_embeddings(
return response.data[0].embedding
return [item.embedding for item in sorted(response.data, key=lambda
item: item.index)]
+ @overload
+ async def acreate_embeddings(
+ self,
+ text: str | list[int],
+ model: str = "text-embedding-3-small",
+ **kwargs: Any,
+ ) -> list[float]: ...
+
+ @overload
+ async def acreate_embeddings(
+ self,
+ text: list[str] | list[list[int]],
+ model: str = "text-embedding-3-small",
+ **kwargs: Any,
+ ) -> list[list[float]]: ...
+
+ async def acreate_embeddings(
+ self,
+ text: str | list[str] | list[int] | list[list[int]],
+ model: str = "text-embedding-3-small",
+ **kwargs: Any,
+ ) -> list[float] | list[list[float]]:
+ """
+ Generate embeddings for the given text using the given model,
asynchronously.
+
+ The async counterpart of :meth:`create_embeddings`, through
:meth:`aget_conn`.
+
+ :param text: The text to generate embeddings for.
+ :param model: The model to use for generating embeddings.
+ :return: One embedding for a single text or token array; one embedding
per item for a batch.
+ """
+ client = await self._aget_cached_conn()
+ response = await client.embeddings.create(model=model, input=text,
**kwargs)
+ if isinstance(text, str) or (text and isinstance(text[0], int)):
Review Comment:
Fixed in a9c790045b: `_embeddings_result(text, response)`, used by both
`create_embeddings` and `acreate_embeddings`.
Drafted-by: Claude Fable 5.1; reviewed by @dabla before posting
##########
providers/openai/tests/unit/openai/hooks/test_openai.py:
##########
@@ -998,3 +1000,156 @@ 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", "extra_dejson"] +
(["aextra_dejson"] if aextra_dejson else [])
+ )
+ conn.password = password
+ conn.host = host
+ conn.extra = json.dumps(extra)
Review Comment:
Fixed in 7d14c8e1d6: `extra` dropped from the mock's spec and
`conn.aextra_dejson.assert_awaited_once_with()` in the tests that have it. With
the hook reading `json.loads(conn.extra)` instead, six of the `aget_conn` tests
fail now.
Drafted-by: Claude Fable 5.1; reviewed by @dabla before posting
##########
providers/openai/tests/unit/openai/hooks/test_openai.py:
##########
@@ -998,3 +1000,156 @@ 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", "extra_dejson"] +
(["aextra_dejson"] if aextra_dejson else [])
+ )
+ conn.password = password
+ conn.host = host
+ conn.extra = json.dumps(extra)
+ 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,
+ ):
+ client = await OpenAIHook(conn_id="openai_async").aget_conn()
+
+ mock_get_async_connection.assert_awaited_once_with("openai_async")
+ 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()
+
+ build.assert_called_once_with(extra)
+ mock_client.assert_called_once_with(workload_identity={"provider":
"custom"}, base_url="https://host/v1";)
+
+
[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()
+
+
[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 = MagicMock()
Review Comment:
Fixed in fdbd00aee8: `MagicMock(spec=AsyncOpenAI)`.
Drafted-by: Claude Fable 5.1; reviewed by @dabla before posting
--
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]
