This is an automated email from the ASF dual-hosted git repository.
dabla 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 6997dc76fe5 Add async aget_conn and acreate_embeddings to OpenAIHook
(#74149)
6997dc76fe5 is described below
commit 6997dc76fe50fc2027081d6c303db44e96bfd57f
Author: David Blain <[email protected]>
AuthorDate: Fri Oct 16 08:24:19 2026 +0200
Add async aget_conn and acreate_embeddings to OpenAIHook (#74149)
* Add async aget_conn and acreate_embeddings to OpenAIHook
get_conn() resolves the connection with get_connection() and reads
extra_dejson, whose secret masking sends to the supervisor
synchronously; inside an async task with another async SDK call in
flight both raise DeadlockImminentError. aget_conn() returns an
AsyncOpenAI client configured exactly like get_conn()'s, from
get_async_connection() with the extra from Connection.aextra_dejson()
(#71890, Airflow 3.3.2+), or the raw extra via json.loads on older
versions (as in #72130). The client settings move into
_client_kwargs(), shared by both paths. acreate_embeddings() is the
async counterpart of create_embeddings().
Co-Authored-By: Claude Opus 5.5 <[email protected]>
* Use get_async_extra_dejson from common.compat
Replace the local aextra_dejson fallback with the common.compat helper
added in #74147, which masks the extra's secrets asynchronously on
Airflow 3.3.2+ and reads extra_dejson from a worker thread on older
versions.
Co-Authored-By: Claude Opus 5.5 <[email protected]>
* Spec the mocked AsyncOpenAI client in test_acreate_embeddings
MagicMock() without a spec accepts any attribute, so a typo on the async
path such as client.embedding would still pass. spec=AsyncOpenAI guards
the client's namespace the way the sync tests' spec=OpenAI does.
* Assert aget_conn reads the extra through aextra_dejson
The connection mock carried an extra string and no test checked that
aextra_dejson() was awaited, so a hook reading json.loads(conn.extra),
which skips the secret masking, kept every test green. Drop extra from
the mock's spec and assert the await in the tests that have it.
* Share the embeddings result shaping between the sync and async paths
acreate_embeddings copied the three lines that pick one vector for a
single input and sort one vector per item for a batch. Move them into
_embeddings_result, used by create_embeddings and acreate_embeddings,
so the two paths cannot drift, like _client_kwargs for the client.
* Resolve the async connection through the hook
get_conn resolves through self.get_connection, so a subclass override
applies there, but aget_conn called get_async_connection without hook=
and the lookup went through BaseHook. Pass hook=self, as
AnthropicHook.get_async_conn does, and test it with a subclass that
overrides aget_connection, through the real common.compat helper.
* Build the async client per call instead of caching it on the hook
_aget_cached_conn checked the cache before two awaits, so concurrent
calls on a fresh hook each built a client, dropped all but one
unclosed, and the kept one was never closed and stayed bound to the
first event loop. aget_conn now builds a client per call and says to
close it, as AnthropicHook.get_async_conn does; acreate_embeddings
closes its client with async with. Tested with three concurrent calls
on one hook.
* State that aget_conn avoids synchronous Task SDK calls from Airflow 3.3.2
Before 3.3.2, core's aget_connection still masks the connection's secrets
synchronously, so the docstring promised more than older cores give.
* Pin the per-call client in the concurrent embeddings test
Patching aget_conn let the test supply distinct clients itself, so it kept
passing if aget_conn cached one client on the hook again.
* Note the token exchange per client for workload identity in aget_conn
The SDK caches the exchanged token on each client, so callers making many
calls should know they can keep one client open and reuse it.
* Require common.compat 1.21.0 for get_async_extra_dejson
The hook imports get_async_extra_dejson at module level, so with an older
common.compat the whole hook, sync get_conn included, fails to import.
* Restore the use-next-version marker on the common.compat dependency
Only release managers raise provider dependency floors; the marker lets
them bump it to the release that ships get_async_extra_dejson.
* Build the embeddings usage as a Usage model in the async test client
mypy checks the annotated _async_client helper and rejects the dict
passed to CreateEmbeddingResponse.usage, which is typed Usage.
* Update providers/openai/src/airflow/providers/openai/hooks/openai.py
Co-authored-by: Kaxil Naik <[email protected]>
* Update providers/openai/src/airflow/providers/openai/hooks/openai.py
Co-authored-by: Kaxil Naik <[email protected]>
* Update providers/openai/src/airflow/providers/openai/hooks/openai.py
Co-authored-by: Kaxil Naik <[email protected]>
* Update providers/openai/tests/unit/openai/hooks/test_openai.py
Co-authored-by: Kaxil Naik <[email protected]>
* Drop test docstrings that repeat the test names
---------
Co-authored-by: Claude Opus 5.5 <[email protected]>
Co-authored-by: Kaxil Naik <[email protected]>
---
providers/openai/pyproject.toml | 2 +-
.../src/airflow/providers/openai/hooks/openai.py | 79 +++++++-
.../openai/tests/unit/openai/hooks/test_openai.py | 202 ++++++++++++++++++++-
3 files changed, 272 insertions(+), 11 deletions(-)
diff --git a/providers/openai/pyproject.toml b/providers/openai/pyproject.toml
index c07c2e7eb8e..d930c21b1d8 100644
--- a/providers/openai/pyproject.toml
+++ b/providers/openai/pyproject.toml
@@ -59,7 +59,7 @@ requires-python = ">=3.11"
# After you modify the dependencies, and rebuild your Breeze CI image with
``breeze ci-image build``
dependencies = [
"apache-airflow>=2.11.0",
- "apache-airflow-providers-common-compat>=1.12.0",
+ "apache-airflow-providers-common-compat>=1.12.0", # use next version
# The ``[datalib]`` extra of the ``openai`` SDK pulls in numpy / pandas /
# pandas-stubs. It was historically used by the SDK's legacy
# ``openai tools fine_tunes.prepare_data`` CLI, which was removed in
diff --git a/providers/openai/src/airflow/providers/openai/hooks/openai.py
b/providers/openai/src/airflow/providers/openai/hooks/openai.py
index 3987999b961..541056859d0 100644
--- a/providers/openai/src/airflow/providers/openai/hooks/openai.py
+++ b/providers/openai/src/airflow/providers/openai/hooks/openai.py
@@ -23,7 +23,7 @@ from functools import cached_property
from typing import TYPE_CHECKING, Any, BinaryIO, Literal, TypeVar, overload
from deprecated import deprecated
-from openai import OpenAI
+from openai import AsyncOpenAI, OpenAI
from openai.auth import (
azure_managed_identity_token_provider,
gcp_id_token_provider,
@@ -33,6 +33,7 @@ from openai.auth import (
if TYPE_CHECKING:
from openai.auth import SubjectTokenProvider, WorkloadIdentity
from openai.types import (
+ CreateEmbeddingResponse,
FileDeleted,
FileObject,
VectorStore,
@@ -53,7 +54,10 @@ if TYPE_CHECKING:
from openai.types.responses import ParsedResponse, Response
from openai.types.vector_stores import VectorStoreFile,
VectorStoreFileBatch, VectorStoreFileDeleted
from pydantic import BaseModel
+
+ from airflow.providers.common.compat.sdk import Connection
from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.connection import get_async_connection,
get_async_extra_dejson
from airflow.providers.common.compat.module_loading import import_string
from airflow.providers.common.compat.sdk import AirflowException, BaseHook
from airflow.providers.openai.exceptions import (
@@ -207,8 +211,25 @@ class OpenAIHook(BaseHook):
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 configured like :meth:`get_conn`.
+
+ From Airflow 3.3.2 the connection is resolved without a blocking call
to the
+ supervisor, so this is safe on an event loop with other async calls in
flight.
+
+ Each call builds a new client: close it with ``async with`` or ``await
client.close()``.
+ With ``workload_identity`` each client exchanges its own token, so
reuse one client
+ for many calls.
+ """
+ conn = await get_async_connection(self.conn_id, hook=self)
+ return AsyncOpenAI(**self._client_kwargs(conn, await
get_async_extra_dejson(conn)))
+
+ def _client_kwargs(self, conn: Connection, extras: dict[str, Any]) ->
dict[str, Any]:
+ """Build the ``OpenAI`` / ``AsyncOpenAI`` constructor kwargs from the
connection."""
+ openai_client_kwargs = dict(extras.get("openai_client_kwargs", {}))
base_url = openai_client_kwargs.pop("base_url", None) or conn.host or
None
# Pop api_key for every path so it is never forwarded alongside
``workload_identity``
# (the OpenAI client rejects both being set at once).
@@ -216,14 +237,14 @@ class OpenAIHook(BaseHook):
auth_type = extras.get("auth_type", "api_key")
if auth_type == "api_key":
- return OpenAI(api_key=api_key or conn.password, base_url=base_url,
**openai_client_kwargs)
+ return {"api_key": api_key or conn.password, "base_url": base_url,
**openai_client_kwargs}
if auth_type == "workload_identity":
- return OpenAI(
- workload_identity=self._build_workload_identity(extras),
- base_url=base_url,
+ return {
+ "workload_identity": self._build_workload_identity(extras),
+ "base_url": base_url,
**openai_client_kwargs,
- )
+ }
raise ValueError(
f"Unsupported auth_type {auth_type!r} for OpenAI connection
{self.conn_id!r}; "
@@ -593,6 +614,48 @@ class OpenAIHook(BaseHook):
:return: One embedding for a single text or token array; one embedding
per item for a batch.
"""
response = self.conn.embeddings.create(model=model, input=text,
**kwargs)
+ return self._embeddings_result(text, response)
+
+ @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]]:
+ """
+ Async version of :meth:`create_embeddings`.
+
+ Each call opens and closes its own client, so for many inputs pass one
batched ``text`` list.
+
+ :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.
+ """
+ async with await self.aget_conn() as client:
+ response = await client.embeddings.create(model=model, input=text,
**kwargs)
+ return self._embeddings_result(text, response)
+
+ @staticmethod
+ def _embeddings_result(
+ text: str | list[str] | list[int] | list[list[int]], response:
CreateEmbeddingResponse
+ ) -> list[float] | list[list[float]]:
+ """Return one vector for a single text or token array, one vector per
item in input order for a batch."""
if isinstance(text, str) or (text and isinstance(text[0], int)):
return response.data[0].embedding
return [item.embedding for item in sorted(response.data, key=lambda
item: item.index)]
diff --git a/providers/openai/tests/unit/openai/hooks/test_openai.py
b/providers/openai/tests/unit/openai/hooks/test_openai.py
index 0b0222198d7..3b8f9d97d59 100644
--- a/providers/openai/tests/unit/openai/hooks/test_openai.py
+++ b/providers/openai/tests/unit/openai/hooks/test_openai.py
@@ -16,12 +16,13 @@
# under the License.
from __future__ import annotations
+import asyncio
import os
from types import SimpleNamespace
-from unittest.mock import MagicMock, mock_open, patch
+from unittest.mock import AsyncMock, MagicMock, Mock, PropertyMock, mock_open,
patch
import pytest
-from openai import OpenAI
+from openai import AsyncOpenAI, OpenAI
from openai.pagination import SyncCursorPage
from openai.types import (
Batch,
@@ -35,6 +36,7 @@ from openai.types import (
from openai.types.beta import Assistant, AssistantDeleted, Thread,
ThreadDeleted
from openai.types.beta.threads import Message, Run
from openai.types.chat import ChatCompletion
+from openai.types.create_embedding_response import Usage
from openai.types.vector_stores import VectorStoreFile, VectorStoreFileBatch,
VectorStoreFileDeleted
from pydantic import BaseModel
@@ -1032,3 +1034,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:
+ """Connection stub whose ``extra_dejson`` fails on the event loop;
``aextra_dejson=False`` mimics Airflow < 3.3.2."""
+ 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):
+ 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():
+ clients = [_async_client([(0, [0.1, 0.2])]) for _ in range(3)]
+ hook = OpenAIHook(conn_id="openai_async")
+ with (
+ patch(
+ "airflow.providers.openai.hooks.openai.get_async_connection",
+ new=AsyncMock(return_value=_async_conn({}, password="api_key")),
+ ),
+ patch("airflow.providers.openai.hooks.openai.AsyncOpenAI",
side_effect=clients) as mock_client,
+ ):
+ results = await asyncio.gather(*(hook.acreate_embeddings(f"text {i}")
for i in range(3)))
+
+ assert results == [[0.1, 0.2]] * 3
+ assert mock_client.call_count == 3
+ for client in clients:
+ client.embeddings.create.assert_awaited_once()
+ client.__aexit__.assert_awaited_once()
+
+
+def _async_client(response_items: list[tuple[int, list[float]]]) -> MagicMock:
+ """An ``AsyncOpenAI`` stand-in whose ``embeddings.create`` answers with
the given (index, vector) items."""
+ client = MagicMock(spec=AsyncOpenAI)
+ client.__aenter__.return_value = client
+ client.embeddings.create = AsyncMock(
+ return_value=CreateEmbeddingResponse(
+ data=[
+ Embedding(embedding=vector, index=index, object="embedding")
+ for index, vector in response_items
+ ],
+ model="text-embedding-3-small",
+ object="list",
+ usage=Usage(prompt_tokens=4, total_tokens=4),
+ )
+ )
+ return client
