Make open ai provider compatible with openai>=1.66.0 (apache#47726)

closes: apache#47642

Co-authored-by: Simon Prydden <[email protected]>

Author: simonprydden

generated/provider_dependencies.json

@@ -970,7 +970,7 @@
   "openai": {
     "deps": [
       "apache-airflow>=2.9.0",
-      "openai[datalib]>=1.32.0,<1.66.0"
+      "openai[datalib]>=1.66.0"
     ],
     "devel-deps": [],
     "plugins": [],

providers/openai/README.rst

@@ -50,12 +50,12 @@ The package supports the following python versions: 3.9,3.10,3.11,3.12
 Requirements
 ------------
 
-===================  ====================
+===================  ==================
 PIP package          Version required
-===================  ====================
+===================  ==================
 ``apache-airflow``   ``>=2.9.0``
-``openai[datalib]``  ``>=1.32.0,<1.66.0``
-===================  ====================
+``openai[datalib]``  ``>=1.66.0``
+===================  ==================
 
 The changelog for the provider package can be found in the
 `changelog <https://airflow.apache.org/docs/apache-airflow-providers-openai/1.5.2/changelog.html>`_.

providers/openai/pyproject.toml

@@ -58,7 +58,7 @@ requires-python = "~=3.9"
 # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
 dependencies = [
     "apache-airflow>=2.9.0",
-    "openai[datalib]>=1.32.0,<1.66.0",  # Pinned due to https://github.com/apache/airflow/issues/47642
+    "openai[datalib]>=1.66.0",
 ]
 
 [dependency-groups]

providers/openai/src/airflow/providers/openai/get_provider_info.py

@@ -64,6 +64,6 @@ def get_provider_info():
                 "connection-type": "openai",
             }
         ],
-        "dependencies": ["apache-airflow>=2.9.0", "openai[datalib]>=1.32.0,<1.66.0"],
+        "dependencies": ["apache-airflow>=2.9.0", "openai[datalib]>=1.66.0"],
         "devel-dependencies": [],
     }

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

@@ -25,18 +25,15 @@
 from openai import OpenAI
 
 if TYPE_CHECKING:
-    from openai.types import FileDeleted, FileObject
-    from openai.types.batch import Batch
-    from openai.types.beta import (
-        Assistant,
-        AssistantDeleted,
-        Thread,
-        ThreadDeleted,
+    from openai.types import (
+        FileDeleted,
+        FileObject,
         VectorStore,
         VectorStoreDeleted,
     )
+    from openai.types.batch import Batch
+    from openai.types.beta import Assistant, AssistantDeleted, Thread, ThreadDeleted
     from openai.types.beta.threads import Message, Run
-    from openai.types.beta.vector_stores import VectorStoreFile, VectorStoreFileBatch, VectorStoreFileDeleted
     from openai.types.chat import (
         ChatCompletionAssistantMessageParam,
         ChatCompletionFunctionMessageParam,
@@ -45,6 +42,7 @@
         ChatCompletionToolMessageParam,
         ChatCompletionUserMessageParam,
     )
+    from openai.types.vector_stores import VectorStoreFile, VectorStoreFileBatch, VectorStoreFileDeleted
 from airflow.hooks.base import BaseHook
 from airflow.providers.openai.exceptions import OpenAIBatchJobException, OpenAIBatchTimeout
 
@@ -349,12 +347,12 @@ def delete_file(self, file_id: str) -> FileDeleted:
 
     def create_vector_store(self, **kwargs: Any) -> VectorStore:
         """Create a vector store."""
-        vector_store = self.conn.beta.vector_stores.create(**kwargs)
+        vector_store = self.conn.vector_stores.create(**kwargs)
         return vector_store
 
     def get_vector_stores(self, **kwargs: Any) -> list[VectorStore]:
         """Return a list of vector stores."""
-        vector_stores = self.conn.beta.vector_stores.list(**kwargs)
+        vector_stores = self.conn.vector_stores.list(**kwargs)
         return vector_stores.data
 
     def get_vector_store(self, vector_store_id: str) -> VectorStore:
@@ -363,7 +361,7 @@ def get_vector_store(self, vector_store_id: str) -> VectorStore:
 
         :param vector_store_id: The ID of the vector store to retrieve.
         """
-        vector_store = self.conn.beta.vector_stores.retrieve(vector_store_id=vector_store_id)
+        vector_store = self.conn.vector_stores.retrieve(vector_store_id=vector_store_id)
         return vector_store
 
     def modify_vector_store(self, vector_store_id: str, **kwargs: Any) -> VectorStore:
@@ -372,7 +370,7 @@ def modify_vector_store(self, vector_store_id: str, **kwargs: Any) -> VectorStor
 
         :param vector_store_id: The ID of the vector store to modify.
         """
-        vector_store = self.conn.beta.vector_stores.update(vector_store_id=vector_store_id, **kwargs)
+        vector_store = self.conn.vector_stores.update(vector_store_id=vector_store_id, **kwargs)
         return vector_store
 
     def delete_vector_store(self, vector_store_id: str) -> VectorStoreDeleted:
@@ -381,7 +379,7 @@ def delete_vector_store(self, vector_store_id: str) -> VectorStoreDeleted:
 
         :param vector_store_id: The ID of the vector store to delete.
         """
-        response = self.conn.beta.vector_stores.delete(vector_store_id=vector_store_id)
+        response = self.conn.vector_stores.delete(vector_store_id=vector_store_id)
         return response
 
     def upload_files_to_vector_store(
@@ -394,7 +392,7 @@ def upload_files_to_vector_store(
             to.
         :param files: A list of binary files to upload.
         """
-        file_batch = self.conn.beta.vector_stores.file_batches.upload_and_poll(
+        file_batch = self.conn.vector_stores.file_batches.upload_and_poll(
             vector_store_id=vector_store_id, files=files
         )
         return file_batch
@@ -405,7 +403,7 @@ def get_vector_store_files(self, vector_store_id: str) -> list[VectorStoreFile]:
 
         :param vector_store_id:
         """
-        vector_store_files = self.conn.beta.vector_stores.files.list(vector_store_id=vector_store_id)
+        vector_store_files = self.conn.vector_stores.files.list(vector_store_id=vector_store_id)
         return vector_store_files.data
 
     def delete_vector_store_file(self, vector_store_id: str, file_id: str) -> VectorStoreFileDeleted:
@@ -415,7 +413,7 @@ def delete_vector_store_file(self, vector_store_id: str, file_id: str) -> Vector
         :param vector_store_id: The ID of the vector store that the file belongs to.
         :param file_id: The ID of the file to delete.
         """
-        response = self.conn.beta.vector_stores.files.delete(vector_store_id=vector_store_id, file_id=file_id)
+        response = self.conn.vector_stores.files.delete(vector_store_id=vector_store_id, file_id=file_id)
         return response
 
     def create_batch(

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

@@ -21,18 +21,19 @@
 
 import pytest
 from openai.pagination import SyncCursorPage
-from openai.types import Batch, CreateEmbeddingResponse, Embedding, FileDeleted, FileObject
-from openai.types.beta import (
-    Assistant,
-    AssistantDeleted,
-    Thread,
-    ThreadDeleted,
+from openai.types import (
+    Batch,
+    CreateEmbeddingResponse,
+    Embedding,
+    FileDeleted,
+    FileObject,
     VectorStore,
     VectorStoreDeleted,
 )
+from openai.types.beta import Assistant, AssistantDeleted, Thread, ThreadDeleted
 from openai.types.beta.threads import Message, Run
-from openai.types.beta.vector_stores import VectorStoreFile, VectorStoreFileBatch, VectorStoreFileDeleted
 from openai.types.chat import ChatCompletion
+from openai.types.vector_stores import VectorStoreFile, VectorStoreFileBatch, VectorStoreFileDeleted
 
 from airflow.models import Connection
 from airflow.providers.openai.exceptions import OpenAIBatchJobException, OpenAIBatchTimeout
@@ -458,29 +459,29 @@ def test_delete_file(mock_openai_hook):
 
 
 def test_create_vector_store(mock_openai_hook, mock_vector_store):
-    mock_openai_hook.conn.beta.vector_stores.create.return_value = mock_vector_store
+    mock_openai_hook.conn.vector_stores.create.return_value = mock_vector_store
     vector_store = mock_openai_hook.create_vector_store(name=VECTOR_STORE_NAME)
     assert vector_store.id == VECTOR_STORE_ID
     assert vector_store.name == VECTOR_STORE_NAME
 
 
 def test_get_vector_store(mock_openai_hook, mock_vector_store):
-    mock_openai_hook.conn.beta.vector_stores.retrieve.return_value = mock_vector_store
+    mock_openai_hook.conn.vector_stores.retrieve.return_value = mock_vector_store
     vector_store = mock_openai_hook.get_vector_store(vector_store_id=VECTOR_STORE_ID)
     assert vector_store.id == VECTOR_STORE_ID
     assert vector_store.name == VECTOR_STORE_NAME
 
 
 def test_get_vector_stores(mock_openai_hook, mock_vector_store_list):
-    mock_openai_hook.conn.beta.vector_stores.list.return_value = mock_vector_store_list
+    mock_openai_hook.conn.vector_stores.list.return_value = mock_vector_store_list
     vector_stores = mock_openai_hook.get_vector_stores()
     assert isinstance(vector_stores, list)
 
 
 def test_modify_vector_store(mock_openai_hook, mock_vector_store):
     new_vector_store_name = "New Vector Store"
     mock_vector_store.name = new_vector_store_name
-    mock_openai_hook.conn.beta.vector_stores.update.return_value = mock_vector_store
+    mock_openai_hook.conn.vector_stores.update.return_value = mock_vector_store
     vector_store = mock_openai_hook.modify_vector_store(
         vector_store_id=VECTOR_STORE_ID, name=new_vector_store_name
     )
@@ -489,14 +490,14 @@ def test_modify_vector_store(mock_openai_hook, mock_vector_store):
 
 def test_delete_vector_store(mock_openai_hook):
     delete_response = VectorStoreDeleted(id=VECTOR_STORE_ID, object="vector_store.deleted", deleted=True)
-    mock_openai_hook.conn.beta.vector_stores.delete.return_value = delete_response
+    mock_openai_hook.conn.vector_stores.delete.return_value = delete_response
     vector_store_deleted = mock_openai_hook.delete_vector_store(vector_store_id=VECTOR_STORE_ID)
     assert vector_store_deleted.deleted
 
 
 def test_upload_files_to_vector_store(mock_openai_hook, mock_vector_file_store_batch):
     files = ["file1.txt", "file2.txt", "file3.txt"]
-    mock_openai_hook.conn.beta.vector_stores.file_batches.upload_and_poll.return_value = (
+    mock_openai_hook.conn.vector_stores.file_batches.upload_and_poll.return_value = (
         mock_vector_file_store_batch
     )
     vector_file_store_batch = mock_openai_hook.upload_files_to_vector_store(
@@ -507,7 +508,7 @@ def test_upload_files_to_vector_store(mock_openai_hook, mock_vector_file_store_b
 
 
 def test_get_vector_store_files(mock_openai_hook, mock_vector_file_store_list):
-    mock_openai_hook.conn.beta.vector_stores.files.list.return_value = mock_vector_file_store_list
+    mock_openai_hook.conn.vector_stores.files.list.return_value = mock_vector_file_store_list
     vector_file_store_list = mock_openai_hook.get_vector_store_files(vector_store_id=VECTOR_STORE_ID)
     assert isinstance(vector_file_store_list, list)
 
@@ -516,7 +517,7 @@ def test_delete_vector_store_file(mock_openai_hook):
     delete_response = VectorStoreFileDeleted(
         id="test_file_abc123", object="vector_store.file.deleted", deleted=True
     )
-    mock_openai_hook.conn.beta.vector_stores.files.delete.return_value = delete_response
+    mock_openai_hook.conn.vector_stores.files.delete.return_value = delete_response
     vector_store_file_deleted = mock_openai_hook.delete_vector_store_file(
         vector_store_id=VECTOR_STORE_ID, file_id=FILE_ID
     )