From 35bb928a82633ee5b2a1e2b79d5c191640825197 Mon Sep 17 00:00:00 2001 From: Yoel Nisanov Date: Thu, 9 Nov 2023 16:30:49 +0700 Subject: [PATCH 01/10] Changed type of input_text for openai operator as well as added validation that it matches the type we need, and aligned docstring for the openai operator accordingly --- airflow/providers/openai/hooks/openai.py | 5 +++- airflow/providers/openai/operators/openai.py | 30 +++++++++++++------- 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/airflow/providers/openai/hooks/openai.py b/airflow/providers/openai/hooks/openai.py index ee21b5a3c3f08..fac725b5bea40 100644 --- a/airflow/providers/openai/hooks/openai.py +++ b/airflow/providers/openai/hooks/openai.py @@ -74,7 +74,10 @@ def _get_api_base(self) -> None | str: return conn.host def create_embeddings( - self, text: str | list[Any], model: str = "text-embedding-ada-002", **kwargs: Any + self, + text: str | list[str] | list[int] | list[list[int]], + model: str = "text-embedding-ada-002", + **kwargs: Any, ) -> list[float]: """Generate embeddings for the given text using the given model. diff --git a/airflow/providers/openai/operators/openai.py b/airflow/providers/openai/operators/openai.py index 1e585b5d0a262..9b593c0d34bf6 100644 --- a/airflow/providers/openai/operators/openai.py +++ b/airflow/providers/openai/operators/openai.py @@ -20,6 +20,7 @@ from functools import cached_property from typing import TYPE_CHECKING, Any, Sequence +from airflow.exceptions import AirflowException from airflow.models import BaseOperator from airflow.providers.openai.hooks.openai import OpenAIHook @@ -31,16 +32,18 @@ class OpenAIEmbeddingOperator(BaseOperator): """ Operator that accepts input text to generate OpenAI embeddings using the specified model. + :param conn_id: The OpenAI connection ID to use. + :param input_text: The text to generate OpenAI embeddings for. This can be a string, a list of strings, + a list of integers, or a list of lists of integers. + :param model: The OpenAI model to be used for generating the embeddings. Defaults to 'text-embedding-ada-002'. + :param embedding_kwargs: Additional keyword arguments to pass to the OpenAI `create_embeddings` method. + :raises AirflowException: Raises an exception if `input_text` is empty or not a string or list. + .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:OpenAIEmbeddingOperator` - - :param conn_id: The OpenAI connection. - :param input_text: The text to generate OpenAI embeddings on. Either input_text or input_callable - should be provided. - :param model: The OpenAI model to be used for generating the embeddings. - :param embedding_kwargs: For possible option check - .. seealso:: https://platform.openai.com/docs/api-reference/embeddings/create + For possible options for `embedding_kwargs`, see: + https://platform.openai.com/docs/api-reference/embeddings/create """ template_fields: Sequence[str] = ("input_text",) @@ -48,16 +51,21 @@ class OpenAIEmbeddingOperator(BaseOperator): def __init__( self, conn_id: str, - input_text: str | list[Any], + input_text: str | list[str] | list[int] | list[list[int]], model: str = "text-embedding-ada-002", embedding_kwargs: dict | None = None, **kwargs: Any, ): - self.embedding_kwargs = embedding_kwargs or {} super().__init__(**kwargs) self.conn_id = conn_id self.input_text = input_text self.model = model + self.embedding_kwargs = embedding_kwargs or {} + + if not self.input_text or not isinstance(self.input_text, (str, list)): + raise AirflowException( + "The 'input_text' must be a non-empty string, list of strings, list of integers, or list of lists of integers." + ) @cached_property def hook(self) -> OpenAIHook: @@ -65,7 +73,7 @@ def hook(self) -> OpenAIHook: return OpenAIHook(conn_id=self.conn_id) def execute(self, context: Context) -> list[float]: - self.log.info("Input text: %s", self.input_text) + self.log.info("Generating embeddings for the input text: %s", self.input_text) embeddings = self.hook.create_embeddings(self.input_text, model=self.model, **self.embedding_kwargs) - self.log.info("Embeddings: %s", embeddings) + self.log.info("Generated embeddings: %s", embeddings) return embeddings From bfb9b16069deca1fc0ba61c2819af250d045a542 Mon Sep 17 00:00:00 2001 From: Yoel Nisanov Date: Thu, 9 Nov 2023 16:45:44 +0700 Subject: [PATCH 02/10] Added tests for invalid inputs for OpenAI operator --- tests/providers/openai/operators/test_openai.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/providers/openai/operators/test_openai.py b/tests/providers/openai/operators/test_openai.py index 11a6483802790..e124e2db58679 100644 --- a/tests/providers/openai/operators/test_openai.py +++ b/tests/providers/openai/operators/test_openai.py @@ -18,6 +18,8 @@ from unittest.mock import Mock +import pytest + from airflow.providers.openai.operators.openai import OpenAIEmbeddingOperator from airflow.utils.context import Context @@ -34,3 +36,18 @@ def test_execute_with_input_text(): embeddings = operator.execute(context) assert embeddings == [1.0, 2.0, 3.0] + + +def test_execute_with_invalid_input_empty_string(): + with pytest.raises(ValueError): + OpenAIEmbeddingOperator(task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text="") + + +def test_execute_with_invalid_input_none(): + with pytest.raises(ValueError): + OpenAIEmbeddingOperator(task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=None) + + +def test_execute_with_invalid_input_wrong_type(): + with pytest.raises(ValueError): + OpenAIEmbeddingOperator(task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=123) From 62f788ad8b970f6130f74eb9ba080b09a65034b6 Mon Sep 17 00:00:00 2001 From: Yoel Nisanov Date: Thu, 9 Nov 2023 17:00:53 +0700 Subject: [PATCH 03/10] Fixed exception type when providing invalid input on OpenAI operators --- tests/providers/openai/operators/test_openai.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/providers/openai/operators/test_openai.py b/tests/providers/openai/operators/test_openai.py index e124e2db58679..09e7fa731fd87 100644 --- a/tests/providers/openai/operators/test_openai.py +++ b/tests/providers/openai/operators/test_openai.py @@ -20,6 +20,7 @@ import pytest +from airflow.exceptions import AirflowException from airflow.providers.openai.operators.openai import OpenAIEmbeddingOperator from airflow.utils.context import Context @@ -39,15 +40,15 @@ def test_execute_with_input_text(): def test_execute_with_invalid_input_empty_string(): - with pytest.raises(ValueError): + with pytest.raises(AirflowException): OpenAIEmbeddingOperator(task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text="") def test_execute_with_invalid_input_none(): - with pytest.raises(ValueError): + with pytest.raises(AirflowException): OpenAIEmbeddingOperator(task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=None) def test_execute_with_invalid_input_wrong_type(): - with pytest.raises(ValueError): + with pytest.raises(AirflowException): OpenAIEmbeddingOperator(task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=123) From eb76e6c23d9d0ec0ee90b0327efd8d709d5915b7 Mon Sep 17 00:00:00 2001 From: Yoel Nisanov Date: Thu, 9 Nov 2023 17:04:29 +0700 Subject: [PATCH 04/10] Removed from docstring lines that caused failing due to misspelling --- airflow/providers/openai/operators/openai.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/airflow/providers/openai/operators/openai.py b/airflow/providers/openai/operators/openai.py index 9b593c0d34bf6..c11f8017c5adc 100644 --- a/airflow/providers/openai/operators/openai.py +++ b/airflow/providers/openai/operators/openai.py @@ -35,9 +35,8 @@ class OpenAIEmbeddingOperator(BaseOperator): :param conn_id: The OpenAI connection ID to use. :param input_text: The text to generate OpenAI embeddings for. This can be a string, a list of strings, a list of integers, or a list of lists of integers. - :param model: The OpenAI model to be used for generating the embeddings. Defaults to 'text-embedding-ada-002'. + :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. - :raises AirflowException: Raises an exception if `input_text` is empty or not a string or list. .. seealso:: For more information on how to use this operator, take a look at the guide: From 0a58706564725727a0dc78c2c68fee64528777f0 Mon Sep 17 00:00:00 2001 From: Yoel Nisanov Date: Thu, 9 Nov 2023 18:20:37 +0700 Subject: [PATCH 05/10] Removed validation of input_text to the execution of the OpenAI operator --- airflow/providers/openai/operators/openai.py | 9 ++++----- tests/providers/openai/operators/test_openai.py | 15 ++++++++++++--- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/airflow/providers/openai/operators/openai.py b/airflow/providers/openai/operators/openai.py index c11f8017c5adc..bd2e8f3deb71c 100644 --- a/airflow/providers/openai/operators/openai.py +++ b/airflow/providers/openai/operators/openai.py @@ -61,17 +61,16 @@ def __init__( self.model = model self.embedding_kwargs = embedding_kwargs or {} - if not self.input_text or not isinstance(self.input_text, (str, list)): - raise AirflowException( - "The 'input_text' must be a non-empty string, list of strings, list of integers, or list of lists of integers." - ) - @cached_property def hook(self) -> OpenAIHook: """Return an instance of the OpenAIHook.""" return OpenAIHook(conn_id=self.conn_id) def execute(self, context: Context) -> list[float]: + if not self.input_text or not isinstance(self.input_text, (str, list)): + raise AirflowException( + "The 'input_text' must be a non-empty string, list of strings, list of integers, or list of lists of integers." + ) self.log.info("Generating embeddings for the input text: %s", self.input_text) embeddings = self.hook.create_embeddings(self.input_text, model=self.model, **self.embedding_kwargs) self.log.info("Generated embeddings: %s", embeddings) diff --git a/tests/providers/openai/operators/test_openai.py b/tests/providers/openai/operators/test_openai.py index 09e7fa731fd87..130f4ed2a20d8 100644 --- a/tests/providers/openai/operators/test_openai.py +++ b/tests/providers/openai/operators/test_openai.py @@ -41,14 +41,23 @@ def test_execute_with_input_text(): def test_execute_with_invalid_input_empty_string(): with pytest.raises(AirflowException): - OpenAIEmbeddingOperator(task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text="") + operator = OpenAIEmbeddingOperator( + task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text="" + ) + operator.execute() def test_execute_with_invalid_input_none(): with pytest.raises(AirflowException): - OpenAIEmbeddingOperator(task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=None) + operator = OpenAIEmbeddingOperator( + task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=None + ) + operator.execute() def test_execute_with_invalid_input_wrong_type(): with pytest.raises(AirflowException): - OpenAIEmbeddingOperator(task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=123) + operator = OpenAIEmbeddingOperator( + task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=123 + ) + operator.execute() From 2da6005740a587cfa4b84750f07119196e0264a9 Mon Sep 17 00:00:00 2001 From: Yoel Nisanov Date: Thu, 9 Nov 2023 18:39:24 +0700 Subject: [PATCH 06/10] Added empty context to the execution of OpenAI operator on the invalid input tests --- tests/providers/openai/operators/test_openai.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/providers/openai/operators/test_openai.py b/tests/providers/openai/operators/test_openai.py index 130f4ed2a20d8..b5530bc8d5256 100644 --- a/tests/providers/openai/operators/test_openai.py +++ b/tests/providers/openai/operators/test_openai.py @@ -44,7 +44,8 @@ def test_execute_with_invalid_input_empty_string(): operator = OpenAIEmbeddingOperator( task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text="" ) - operator.execute() + context = Context() + operator.execute(context) def test_execute_with_invalid_input_none(): @@ -52,7 +53,8 @@ def test_execute_with_invalid_input_none(): operator = OpenAIEmbeddingOperator( task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=None ) - operator.execute() + context = Context() + operator.execute(context) def test_execute_with_invalid_input_wrong_type(): @@ -60,4 +62,5 @@ def test_execute_with_invalid_input_wrong_type(): operator = OpenAIEmbeddingOperator( task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=123 ) - operator.execute() + context = Context() + operator.execute(context) From c986a936fbed9a90bc071529f812511f2495ffbb Mon Sep 17 00:00:00 2001 From: Yoel Nisanov Date: Mon, 13 Nov 2023 14:11:43 +0700 Subject: [PATCH 07/10] OpenAI operator- Changed execution log for generation of embeddings to log the length of input_text instead of the whole text --- airflow/providers/openai/operators/openai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/openai/operators/openai.py b/airflow/providers/openai/operators/openai.py index bd2e8f3deb71c..34824ba00bd40 100644 --- a/airflow/providers/openai/operators/openai.py +++ b/airflow/providers/openai/operators/openai.py @@ -71,7 +71,7 @@ def execute(self, context: Context) -> list[float]: raise AirflowException( "The 'input_text' must be a non-empty string, list of strings, list of integers, or list of lists of integers." ) - self.log.info("Generating embeddings for the input text: %s", self.input_text) + self.log.info("Generating embeddings for the input text of length: %d", len(self.input_text)) embeddings = self.hook.create_embeddings(self.input_text, model=self.model, **self.embedding_kwargs) self.log.info("Generated embeddings: %s", embeddings) return embeddings From 4f0982d514b6072cc4e4171e35b7d2174896e9eb Mon Sep 17 00:00:00 2001 From: Yoel Nisanov Date: Mon, 13 Nov 2023 14:53:44 +0700 Subject: [PATCH 08/10] OpenAI operator- Changed generation of embeddings, moved validation of input text to the constructor of the operator, changed value of invalid value to ValueError instead of AirflowException --- airflow/providers/openai/operators/openai.py | 11 +++++------ tests/providers/openai/operators/test_openai.py | 7 +++---- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/airflow/providers/openai/operators/openai.py b/airflow/providers/openai/operators/openai.py index 34824ba00bd40..c7f7ae64a18f0 100644 --- a/airflow/providers/openai/operators/openai.py +++ b/airflow/providers/openai/operators/openai.py @@ -20,7 +20,6 @@ from functools import cached_property from typing import TYPE_CHECKING, Any, Sequence -from airflow.exceptions import AirflowException from airflow.models import BaseOperator from airflow.providers.openai.hooks.openai import OpenAIHook @@ -60,6 +59,10 @@ def __init__( self.input_text = input_text self.model = model self.embedding_kwargs = embedding_kwargs or {} + 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." + ) @cached_property def hook(self) -> OpenAIHook: @@ -67,11 +70,7 @@ def hook(self) -> OpenAIHook: return OpenAIHook(conn_id=self.conn_id) def execute(self, context: Context) -> list[float]: - if not self.input_text or not isinstance(self.input_text, (str, list)): - raise AirflowException( - "The 'input_text' must be a non-empty string, list of strings, list of integers, or list of lists of integers." - ) self.log.info("Generating embeddings for the input text of length: %d", len(self.input_text)) embeddings = self.hook.create_embeddings(self.input_text, model=self.model, **self.embedding_kwargs) - self.log.info("Generated embeddings: %s", embeddings) + self.log.info("Generated embeddings for %d items", len(embeddings)) return embeddings diff --git a/tests/providers/openai/operators/test_openai.py b/tests/providers/openai/operators/test_openai.py index b5530bc8d5256..16ba81c31f675 100644 --- a/tests/providers/openai/operators/test_openai.py +++ b/tests/providers/openai/operators/test_openai.py @@ -20,7 +20,6 @@ import pytest -from airflow.exceptions import AirflowException from airflow.providers.openai.operators.openai import OpenAIEmbeddingOperator from airflow.utils.context import Context @@ -40,7 +39,7 @@ def test_execute_with_input_text(): def test_execute_with_invalid_input_empty_string(): - with pytest.raises(AirflowException): + with pytest.raises(ValueError): operator = OpenAIEmbeddingOperator( task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text="" ) @@ -49,7 +48,7 @@ def test_execute_with_invalid_input_empty_string(): def test_execute_with_invalid_input_none(): - with pytest.raises(AirflowException): + with pytest.raises(ValueError): operator = OpenAIEmbeddingOperator( task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=None ) @@ -58,7 +57,7 @@ def test_execute_with_invalid_input_none(): def test_execute_with_invalid_input_wrong_type(): - with pytest.raises(AirflowException): + with pytest.raises(ValueError): operator = OpenAIEmbeddingOperator( task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=123 ) From f728376d94b01d82db3df9a6804110e0ecb59e93 Mon Sep 17 00:00:00 2001 From: Yoel Nisanov Date: Mon, 13 Nov 2023 23:09:07 +0700 Subject: [PATCH 09/10] OpenAI operator, removed validation to execution function instead of constructor --- airflow/providers/openai/operators/openai.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/airflow/providers/openai/operators/openai.py b/airflow/providers/openai/operators/openai.py index c7f7ae64a18f0..1697e88b98371 100644 --- a/airflow/providers/openai/operators/openai.py +++ b/airflow/providers/openai/operators/openai.py @@ -59,10 +59,6 @@ def __init__( self.input_text = input_text self.model = model self.embedding_kwargs = embedding_kwargs or {} - 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." - ) @cached_property def hook(self) -> OpenAIHook: @@ -70,6 +66,10 @@ def hook(self) -> OpenAIHook: return OpenAIHook(conn_id=self.conn_id) def execute(self, context: Context) -> 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." + ) self.log.info("Generating embeddings for the input text of length: %d", len(self.input_text)) embeddings = self.hook.create_embeddings(self.input_text, model=self.model, **self.embedding_kwargs) self.log.info("Generated embeddings for %d items", len(embeddings)) From 948903e33966a9c1f74f8b577ed957eba67aa52d Mon Sep 17 00:00:00 2001 From: Yoel Nisanov Date: Mon, 13 Nov 2023 23:17:50 +0700 Subject: [PATCH 10/10] OpenAI Operator- Consolidate 3 invalid input types tests into a single parameterized test --- .../providers/openai/operators/test_openai.py | 23 +++---------------- 1 file changed, 3 insertions(+), 20 deletions(-) diff --git a/tests/providers/openai/operators/test_openai.py b/tests/providers/openai/operators/test_openai.py index 16ba81c31f675..12b1cf0b13c34 100644 --- a/tests/providers/openai/operators/test_openai.py +++ b/tests/providers/openai/operators/test_openai.py @@ -38,28 +38,11 @@ def test_execute_with_input_text(): assert embeddings == [1.0, 2.0, 3.0] -def test_execute_with_invalid_input_empty_string(): +@pytest.mark.parametrize("invalid_input", ["", None, 123]) +def test_execute_with_invalid_input(invalid_input): with pytest.raises(ValueError): operator = OpenAIEmbeddingOperator( - task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text="" - ) - context = Context() - operator.execute(context) - - -def test_execute_with_invalid_input_none(): - with pytest.raises(ValueError): - operator = OpenAIEmbeddingOperator( - task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=None - ) - context = Context() - operator.execute(context) - - -def test_execute_with_invalid_input_wrong_type(): - with pytest.raises(ValueError): - operator = OpenAIEmbeddingOperator( - task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=123 + task_id="TaskId", conn_id="test_conn_id", model="test_model", input_text=invalid_input ) context = Context() operator.execute(context)