Return all embeddings for batched OpenAI inputs

providers/openai/docs/operators/openai.rst

@@ -30,6 +30,9 @@ Using the Operator
 The OpenAIEmbeddingOperator requires the ``input_text`` as an input to embedding API. Use the ``conn_id`` parameter to specify the OpenAI connection to use to
 connect to your account.
 
+A single string or token array returns one embedding vector. A list of strings or token arrays returns
+one vector per input item in the same order.
+
 An example of using the operator:
 
 .. exampleinclude:: /../../openai/tests/system/openai/example_openai.py

providers/openai/src/airflow/providers/openai/hooks/openai.py

@@ -20,7 +20,7 @@
 import time
 from enum import Enum
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, BinaryIO, Literal
+from typing import TYPE_CHECKING, Any, BinaryIO, Literal, overload
 
 from deprecated import deprecated
 from openai import OpenAI
@@ -494,21 +494,39 @@ def modify_run(self, thread_id: str, run_id: str, **kwargs: Any) -> Run:
         run = self.conn.beta.threads.runs.update(thread_id=thread_id, run_id=run_id, **kwargs)
         return run
 
+    @overload
+    def create_embeddings(
+        self,
+        text: str | list[int],
+        model: str = "text-embedding-3-small",
+        **kwargs: Any,
+    ) -> list[float]: ...
+
+    @overload
+    def create_embeddings(
+        self,
+        text: list[str] | list[list[int]],
+        model: str = "text-embedding-3-small",
+        **kwargs: Any,
+    ) -> list[list[float]]: ...
+
     def create_embeddings(
         self,
         text: str | list[str] | list[int] | list[list[int]],
         model: str = "text-embedding-3-small",
         **kwargs: Any,
-    ) -> list[float]:
+    ) -> list[float] | list[list[float]]:
         """
         Generate embeddings for the given text using the given model.
 
         :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.
         """
         response = self.conn.embeddings.create(model=model, input=text, **kwargs)
-        embeddings: list[float] = response.data[0].embedding
-        return embeddings
+        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)]
 
     def upload_file(self, file: str, purpose: Literal["fine-tune", "assistants", "batch"]) -> FileObject:
         """

providers/openai/src/airflow/providers/openai/operators/openai.py

@@ -40,6 +40,8 @@ class OpenAIEmbeddingOperator(BaseOperator):
     :param model: The OpenAI model to be used for generating the embeddings.
     :param embedding_kwargs: Additional keyword arguments to pass to the OpenAI `create_embeddings` method.
 
+    Returns one embedding for a single string or token array, and one embedding per item for a batch.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:OpenAIEmbeddingOperator`
@@ -68,7 +70,7 @@ def hook(self) -> OpenAIHook:
         """Return an instance of the OpenAIHook."""
         return OpenAIHook(conn_id=self.conn_id)
 
-    def execute(self, context: Context) -> list[float]:
+    def execute(self, context: Context) -> list[float] | list[list[float]]:
         if not self.input_text or not isinstance(self.input_text, (str, list)):
             raise ValueError(
                 "The 'input_text' must be a non-empty string, list of strings, list of integers, or list of lists of integers."

providers/openai/tests/unit/openai/hooks/test_openai.py

@@ -509,6 +509,41 @@ def test_create_embeddings(mock_openai_hook, mock_embeddings_response):
     assert embeddings == [0.1, 0.2, 0.3]
 
 
+@pytest.mark.parametrize(
+    ("input_text", "response_items", "expected"),
+    [
+        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",
+        ),
+        pytest.param(
+            [[1, 2], [3, 4]],
+            [(1, [0.3, 0.4]), (0, [0.1, 0.2])],
+            [[0.1, 0.2], [0.3, 0.4]],
+            id="token-batch",
+        ),
+    ],
+)
+def test_create_batched_embeddings(input_text, response_items, expected):
+    hook = OpenAIHook(conn_id="unused")
+    conn = MagicMock(spec=OpenAI)
+    conn.embeddings.create.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={"prompt_tokens": 4, "total_tokens": 4},
+    )
+    hook.__dict__["conn"] = conn
+
+    embeddings = hook.create_embeddings(input_text)
+
+    assert embeddings == expected
+
+
 @patch("builtins.open", new_callable=mock_open, read_data="test-data")
 def test_upload_file(mock_file_open, mock_openai_hook, mock_file):
     mock_file.name = FILE_NAME