diff --git a/providers/openai/docs/operators/openai.rst b/providers/openai/docs/operators/openai.rst index bfb2fd8dee94b..8eeba975b95d6 100644 --- a/providers/openai/docs/operators/openai.rst +++ b/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 diff --git a/providers/openai/src/airflow/providers/openai/hooks/openai.py b/providers/openai/src/airflow/providers/openai/hooks/openai.py index 97dcc4ceec8dc..e046881ea4099 100644 --- a/providers/openai/src/airflow/providers/openai/hooks/openai.py +++ b/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: """ diff --git a/providers/openai/src/airflow/providers/openai/operators/openai.py b/providers/openai/src/airflow/providers/openai/operators/openai.py index dfed6e48d51a9..3c8e99304d930 100644 --- a/providers/openai/src/airflow/providers/openai/operators/openai.py +++ b/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." diff --git a/providers/openai/tests/unit/openai/hooks/test_openai.py b/providers/openai/tests/unit/openai/hooks/test_openai.py index d132a6a3d9582..60b4b5585810b 100644 --- a/providers/openai/tests/unit/openai/hooks/test_openai.py +++ b/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