diff --git a/tests/providers/openai/hooks/test_openai.py b/tests/providers/openai/hooks/test_openai.py index aa7a479bbbe0b1..42e161129cb744 100644 --- a/tests/providers/openai/hooks/test_openai.py +++ b/tests/providers/openai/hooks/test_openai.py @@ -23,10 +23,20 @@ openai = pytest.importorskip("openai") +from unittest.mock import mock_open + from openai.pagination import SyncCursorPage -from openai.types import CreateEmbeddingResponse, Embedding -from openai.types.beta import Assistant, AssistantDeleted, Thread, ThreadDeleted +from openai.types import CreateEmbeddingResponse, Embedding, FileDeleted, FileObject +from openai.types.beta import ( + Assistant, + AssistantDeleted, + Thread, + ThreadDeleted, + VectorStore, + VectorStoreDeleted, +) 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 airflow.models import Connection @@ -39,7 +49,12 @@ MESSAGE_ID = "test_message_abc123" RUN_ID = "test_run_abc123" MODEL = "gpt-4" +FILE_ID = "test_file_abc123" +FILE_NAME = "test_file.pdf" METADATA = {"modified": "true", "user": "abc123"} +VECTOR_STORE_ID = "test_vs_abc123" +VECTOR_STORE_NAME = "Test Vector Store" +VECTOR_FILE_STORE_BATCH_ID = "test_vfsb_abc123" @pytest.fixture @@ -161,6 +176,86 @@ def mock_run_list(mock_run): return SyncCursorPage[Run](data=[mock_run]) +@pytest.fixture +def mock_file(): + return FileObject( + id=FILE_ID, + object="file", + bytes=120000, + created_at=1677610602, + filename=FILE_NAME, + purpose="assistants", + status="processed", + ) + + +@pytest.fixture +def mock_file_list(mock_file): + return SyncCursorPage[FileObject](data=[mock_file]) + + +@pytest.fixture +def mock_vector_store(): + return VectorStore( + id=VECTOR_STORE_ID, + object="vector_store", + created_at=1698107661, + usage_bytes=123456, + last_active_at=1698107661, + name=VECTOR_STORE_NAME, + bytes=123456, + status="completed", + file_counts={"in_progress": 0, "completed": 100, "cancelled": 0, "failed": 0, "total": 100}, + metadata={}, + last_used_at=1698107661, + ) + + +@pytest.fixture +def mock_vector_store_list(mock_vector_store): + return SyncCursorPage[VectorStore](data=[mock_vector_store]) + + +@pytest.fixture +def mock_vector_file_store_batch(): + return VectorStoreFileBatch( + id=VECTOR_FILE_STORE_BATCH_ID, + object="vector_store.files_batch", + created_at=1699061776, + vector_store_id=VECTOR_STORE_ID, + status="completed", + file_counts={ + "in_progress": 0, + "completed": 3, + "failed": 0, + "cancelled": 0, + "total": 0, + }, + ) + + +@pytest.fixture +def mock_vector_file_store_list(): + return SyncCursorPage[VectorStoreFile]( + data=[ + VectorStoreFile( + id="test-file-abc123", + object="vector_store.file", + created_at=1699061776, + vector_store_id=VECTOR_STORE_ID, + status="completed", + ), + VectorStoreFile( + id="test-file-abc456", + object="vector_store.file", + created_at=1699061776, + vector_store_id=VECTOR_STORE_ID, + status="completed", + ), + ] + ) + + def test_create_chat_completion(mock_openai_hook, mock_completion): messages = [ {"role": "system", "content": "You are a helpful assistant."}, @@ -269,6 +364,14 @@ def test_create_run(mock_openai_hook, mock_run): assert run.id == RUN_ID +def test_create_run_and_poll(mock_openai_hook, mock_run): + thread_id = THREAD_ID + assistant_id = ASSISTANT_ID + mock_openai_hook.conn.beta.threads.runs.create_and_poll.return_value = mock_run + run = mock_openai_hook.create_run_and_poll(thread_id=thread_id, assistant_id=assistant_id) + assert run.id == RUN_ID + + def test_get_runs(mock_openai_hook, mock_run_list): mock_openai_hook.conn.beta.threads.runs.list.return_value = mock_run_list runs = mock_openai_hook.get_runs(thread_id=THREAD_ID) @@ -296,6 +399,117 @@ 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_get_file_by_name(mock_openai_hook, mock_file_list): + mock_openai_hook.conn.files.list.return_value = mock_file_list + file = mock_openai_hook.get_file_by_name(file_name=FILE_NAME) + assert file.id == FILE_ID + assert file.filename == FILE_NAME + + +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() + assert isinstance(vector_stores, list) + + +def test_get_vector_store_by_name(mock_openai_hook, mock_vector_store_list): + mock_openai_hook.conn.beta.vector_stores.list.return_value = mock_vector_store_list + vector_store = mock_openai_hook.get_vector_store_by_name(vector_store_name=VECTOR_STORE_NAME) + assert vector_store.id == VECTOR_STORE_ID + assert vector_store.name == VECTOR_STORE_NAME + + +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 + vector_store = mock_openai_hook.modify_vector_store( + vector_store_id=VECTOR_STORE_ID, name=new_vector_store_name + ) + assert vector_store.name == new_vector_store_name + + +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 + 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_vector_file_store_batch + ) + vector_file_store_batch = mock_openai_hook.upload_files_to_vector_store( + vector_store_id=VECTOR_STORE_ID, files=files + ) + assert vector_file_store_batch.id == VECTOR_FILE_STORE_BATCH_ID + assert vector_file_store_batch.file_counts.completed == len(files) + + +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 + vector_file_store_list = mock_openai_hook.get_vector_store_files(vector_store_id=VECTOR_STORE_ID) + assert isinstance(vector_file_store_list, list) + + +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 + vector_store_file_deleted = mock_openai_hook.delete_vector_store_file( + vector_store_id=VECTOR_STORE_ID, file_id=FILE_ID + ) + assert vector_store_file_deleted.id == FILE_ID + assert vector_store_file_deleted.deleted + + def test_openai_hook_test_connection(mock_openai_hook): result, message = mock_openai_hook.test_connection() assert result is True