josix commented on code in PR #41554:
URL: https://github.com/apache/airflow/pull/41554#discussion_r1723161345
##########
airflow/providers/openai/hooks/openai.py:
##########
@@ -43,9 +46,26 @@
ChatCompletionUserMessageParam,
)
+from airflow.exceptions import AirflowException
from airflow.hooks.base import BaseHook
+class BatchStatus(str, Enum):
+ """Enum for the status of a batch."""
+
+ VALIDATING = "validating"
+ FAILED = "failed"
+ IN_PROGRESS = "in_progress"
+ FINALIZING = "finalizing"
+ COMPLETED = "completed"
+ EXPIRED = "expired"
+ CANCELLING = "cancelling"
+ CANCELLED = "cancelled"
+
+ def __str__(self) -> str:
+ return str(self.value)
Review Comment:
I think it is better to keep it for easier compatible with `Literal` type in
[`Batch.status`](https://github.com/openai/openai-python/blob/main/src/openai/types/batch.py#L39),
which is derived from Pydantic BaseModel. With that, it would be easier to
pass the enum member to the model when we want to build a mock batch like
[this](https://github.com/apache/airflow/pull/41554/files#diff-e0e2387fb039e0c87d941ba74fbf721c61ed06ddd9e877444cabab494bab4485R119)
since it would be automatically casting to string from `BatchStatus` Enum
type. Removing that would cause some annoying validation errors like
```console
FAILED
tests/providers/openai/triggers/test_openai.py::TestOpenAIBatchTrigger::test_openai_batch_for_timeout[in_progress]
- pydantic_core._pydantic_core.ValidationError: 1 validation error for Batch
status
Input should be 'validating', 'failed', 'in_progress', 'finalizing',
'completed', 'expired', 'cancelling' or 'cancelled' [type=literal_error,
input_value=<BatchStatus.IN_PROGRESS: 'in_progress'>, input_type=BatchStatus]
For further information visit
https://errors.pydantic.dev/2.8/v/literal_error
```
Even passing string type like `str(BatchStatus.VALIDATING)` it would still
get invalid string value, as shown below
```console
FAILED
tests/providers/openai/triggers/test_openai.py::TestOpenAIBatchTrigger::test_openai_batch_for_timeout[BatchStatus.VALIDATING]
- pydantic_core._pydantic_core.ValidationError: 1 validation error for Batch
status
Input should be 'validating', 'failed', 'in_progress', 'finalizing',
'completed', 'expired', 'cancelling' or 'cancelled' [type=literal_error,
input_value='BatchStatus.VALIDATING', input_type=str]
For further information visit
https://errors.pydantic.dev/2.8/v/literal_error
```
--
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]