kaxil commented on code in PR #74149:
URL: https://github.com/apache/airflow/pull/74149#discussion_r4284770918
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -593,6 +619,49 @@ def create_embeddings(
: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]]:
+ """
+ Generate embeddings for the given text using the given model,
asynchronously.
+
+ The async counterpart of :meth:`create_embeddings`, on a client from
:meth:`aget_conn`
+ that is closed when the call returns.
Review Comment:
The first two sentences say the same thing. The useful fact is the per-call
client, since the reuse advice on `aget_conn` can't be applied through this
method:
```suggestion
Async version of :meth:`create_embeddings`.
Each call opens and closes its own client, so for many inputs pass
one batched ``text`` list.
```
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -207,23 +211,45 @@ 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.
``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. On
+ Airflow 3.3.2+ the connection is resolved without any synchronous call
to the Task SDK.
+ On older versions, fetching the connection with ``aget_connection``
still masks its
+ secrets synchronously, so that guarantee does not hold there.
+
+ Each call builds a new client. Close it when done, with ``async with``
or
+ ``await client.close()``, to release its HTTP connections. With
``auth_type``
+ ``workload_identity`` the client caches the exchanged token, so every
new client
+ exchanges a token again: for many calls, keep one client open and
reuse it.
Review Comment:
This docstring has grown to 13 lines, and the deadlock mechanics are already
explained in `get_async_extra_dejson` in common.compat. Something like this
keeps what a caller needs:
```suggestion
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.
```
##########
providers/openai/tests/unit/openai/hooks/test_openai.py:
##########
@@ -1032,3 +1034,206 @@ 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.
+ """
Review Comment:
One line is enough here. The docstrings on
`test_aget_conn_without_aextra_dejson_reads_extra_dejson_off_the_event_loop`
and `test_acreate_embeddings_concurrent_calls_each_close_their_own_client`
repeat the test names, so they can go too.
```suggestion
"""Connection stub whose ``extra_dejson`` fails on the event loop;
``aextra_dejson=False`` mimics Airflow < 3.3.2."""
```
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -207,23 +211,45 @@ 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.
``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. On
+ Airflow 3.3.2+ the connection is resolved without any synchronous call
to the Task SDK.
+ On older versions, fetching the connection with ``aget_connection``
still masks its
+ secrets synchronously, so that guarantee does not hold there.
+
+ Each call builds a new client. Close it when done, with ``async with``
or
+ ``await client.close()``, to release its HTTP connections. With
``auth_type``
+ ``workload_identity`` the client caches the exchanged token, so every
new client
+ exchanges a token again: for many calls, keep one client open and
reuse it.
+ """
+ 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]:
+ """Return the keyword arguments of the OpenAI client, for
:meth:`get_conn` and :meth:`aget_conn`."""
Review Comment:
Naming the callers here goes stale as soon as a third one appears.
```suggestion
"""Build the ``OpenAI`` / ``AsyncOpenAI`` constructor kwargs from
the connection."""
```
--
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]
