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

Author: ylnsnv

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,18 +59,18 @@ 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:
         """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 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

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
         )