Lee-W commented on code in PR #39248:
URL: https://github.com/apache/airflow/pull/39248#discussion_r1582501444
##########
tests/providers/openai/hooks/test_openai.py:
##########
@@ -296,6 +397,103 @@ def test_create_embeddings(mock_openai_hook,
mock_embeddings_response):
assert embeddings == [0.1, 0.2, 0.3]
+@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
+ mock_file.purpose = "assistants"
+ mock_openai_hook.conn.files.create.return_value = mock_file
+ file = mock_openai_hook.upload_file(file=mock_file_open(),
purpose="assistants")
+ assert file.name == FILE_NAME
+ assert file.purpose == "assistants"
+
+
+def test_get_file(mock_openai_hook, mock_file):
+ mock_openai_hook.conn.files.retrieve.return_value = mock_file
+ file = mock_openai_hook.get_file(file_id=FILE_ID)
+ assert file.id == FILE_ID
+ assert file.filename == FILE_NAME
+
+
+def test_get_files(mock_openai_hook, mock_file_list):
+ mock_openai_hook.conn.files.list.return_value = mock_file_list
+ files = mock_openai_hook.get_files()
+ assert isinstance(files, list)
+
+
+def test_delete_file(mock_openai_hook):
+ delete_response = FileDeleted(id=FILE_ID, object="file", deleted=True)
+ mock_openai_hook.conn.files.delete.return_value = delete_response
+ file_deleted = mock_openai_hook.delete_file(file_id=FILE_ID)
+ assert file_deleted.deleted
+
+
+def test_create_vector_store(mock_openai_hook, mock_vector_store):
+ mock_openai_hook.conn.beta.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
+ 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
+ vector_stores = mock_openai_hook.get_vectors_stores()
Review Comment:
```suggestion
vector_stores = mock_openai_hook.get_vector_stores()
```
##########
airflow/providers/openai/hooks/openai.py:
##########
@@ -257,11 +278,118 @@ def create_embeddings(
model: str = "text-embedding-ada-002",
**kwargs: Any,
) -> list[float]:
- """Generate embeddings for the given text using the given model.
+ """
+ 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.
"""
response = self.conn.embeddings.create(model=model, input=text,
**kwargs)
embeddings: list[float] = response.data[0].embedding
return embeddings
+
+ def upload_file(self, file: str, purpose: Literal["fine-tune",
"assistants"]) -> FileObject:
+ """
+ Upload a file that can be used across various endpoints. The size of
all the files uploaded by one organization can be up to 100 GB.
+
+ :param file: The File object (not file name) to be uploaded.
+ :param purpose: The intended purpose of the uploaded file. Use
"fine-tune" for
+ Fine-tuning and "assistants" for Assistants and Messages.
+ """
+ with open(file, "rb") as file_stream:
+ file_object = self.conn.files.create(file=file_stream,
purpose=purpose)
+ return file_object
+
+ def get_file(self, file_id: str) -> FileObject:
+ """
+ Return information about a specific file.
+
+ :param file_id: The ID of the file to use for this request.
+ """
+ file = self.conn.files.retrieve(file_id=file_id)
+ return file
+
+ def get_files(self) -> list[FileObject]:
+ """Return a list of files that belong to the user's organization."""
+ files = self.conn.files.list()
+ return files.data
+
+ def delete_file(self, file_id: str) -> FileDeleted:
+ """
+ Delete a file.
+
+ :param file_id: The ID of the file to be deleted.
+ """
+ response = self.conn.files.delete(file_id=file_id)
+ return response
+
+ def create_vector_store(self, **kwargs: Any) -> VectorStore:
+ """Create a vector store."""
+ vector_store = self.conn.beta.vector_stores.create(**kwargs)
+ return vector_store
+
+ def get_vectors_stores(self, **kwargs: Any) -> list[VectorStore]:
Review Comment:
```suggestion
def get_vector_stores(self, **kwargs: Any) -> list[VectorStore]:
```
--
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]
