• Home
• airflow.providers.openai
• airflow.providers.openai.hooks
• airflow.providers.openai.hooks.openai

airflow.providers.openai.hooks.openai

Classes

BatchStatus	Enum for the status of a batch.
OpenAIHook	Use OpenAI SDK to interact with OpenAI APIs.

Module Contents

class airflow.providers.openai.hooks.openai.BatchStatus

Bases: str, enum.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'

__str__()

Return str(self).

classmethod is_in_progress(status)

Check if the batch status is in progress.

class airflow.providers.openai.hooks.openai.OpenAIHook(conn_id=default_conn_name, *args, **kwargs)

Bases: airflow.hooks.base.BaseHook

Use OpenAI SDK to interact with OpenAI APIs.

See also

https://platform.openai.com/docs/introduction/overview

Parameters:

conn_id (str) – OpenAI connection id

conn_name_attr = 'conn_id'

default_conn_name = 'openai_default'

conn_type = 'openai'

hook_name = 'OpenAI'

conn_id = 'openai_default'

classmethod get_ui_field_behaviour()

Return custom field behaviour.

test_connection()

property conn: openai.OpenAI

Return an OpenAI connection object.

get_conn()

Return an OpenAI connection object.

create_chat_completion(messages, model='gpt-3.5-turbo', **kwargs)

Create a model response for the given chat conversation and returns a list of chat completions.

Parameters:

• messages (list[openai.types.chat.ChatCompletionSystemMessageParam | openai.types.chat.ChatCompletionUserMessageParam | openai.types.chat.ChatCompletionAssistantMessageParam | openai.types.chat.ChatCompletionToolMessageParam | openai.types.chat.ChatCompletionFunctionMessageParam]) – A list of messages comprising the conversation so far

• model (str) – ID of the model to use

create_assistant(model='gpt-3.5-turbo', **kwargs)

Create an OpenAI assistant using the given model.

Parameters:

model (str) – The OpenAI model for the assistant to use.

get_assistant(assistant_id)

Get an OpenAI assistant.

Parameters:

assistant_id (str) – The ID of the assistant to retrieve.

get_assistants(**kwargs)

Get a list of Assistant objects.

modify_assistant(assistant_id, **kwargs)

Modify an existing Assistant object.

Parameters:

assistant_id (str) – The ID of the assistant to be modified.

delete_assistant(assistant_id)

Delete an OpenAI Assistant for a given ID.

Parameters:

assistant_id (str) – The ID of the assistant to delete.

create_thread(**kwargs)

Create an OpenAI thread.

modify_thread(thread_id, metadata)

Modify an existing Thread object.

Parameters:

• thread_id (str) – The ID of the thread to modify. Only the metadata can be modified.

• metadata (dict[str, Any]) – Set of 16 key-value pairs that can be attached to an object.

delete_thread(thread_id)

Delete an OpenAI thread for a given thread_id.

Parameters:

thread_id (str) – The ID of the thread to delete.

create_message(thread_id, role, content, **kwargs)

Create a message for a given Thread.

Parameters:

• thread_id (str) – The ID of the thread to create a message for.

• role (Literal['user', 'assistant']) – The role of the entity that is creating the message. Allowed values include: ‘user’, ‘assistant’.

• content (str) – The content of the message.

get_messages(thread_id, **kwargs)

Return a list of messages for a given Thread.

Parameters:

thread_id (str) – The ID of the thread the messages belong to.

modify_message(thread_id, message_id, **kwargs)

Modify an existing message for a given Thread.

Parameters:

• thread_id (str) – The ID of the thread to which this message belongs.

• message_id – The ID of the message to modify.

create_run(thread_id, assistant_id, **kwargs)

Create a run for a given thread and assistant.

Parameters:

• thread_id (str) – The ID of the thread to run.

• assistant_id (str) – The ID of the assistant to use to execute this run.

create_run_and_poll(thread_id, assistant_id, **kwargs)

Create a run for a given thread and assistant and then polls until completion.

Parameters:

• thread_id (str) – The ID of the thread to run.

• assistant_id (str) – The ID of the assistant to use to execute this run.

Returns:

An OpenAI Run object

Return type:

openai.types.beta.threads.Run

get_run(thread_id, run_id)

Retrieve a run for a given thread and run.

Parameters:

• thread_id (str) – The ID of the thread that was run.

• run_id (str) – The ID of the run to retrieve.

get_runs(thread_id, **kwargs)

Return a list of runs belonging to a thread.

Parameters:

thread_id (str) – The ID of the thread the run belongs to.

modify_run(thread_id, run_id, **kwargs)

Modify a run on a given thread.

Parameters:

• thread_id (str) – The ID of the thread that was run.

• run_id (str) – The ID of the run to modify.

create_embeddings(text, model='text-embedding-ada-002', **kwargs)

Generate embeddings for the given text using the given model.

Parameters:

• text (str | list[str] | list[int] | list[list[int]]) – The text to generate embeddings for.

• model (str) – The model to use for generating embeddings.

upload_file(file, purpose)

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.

Parameters:

• file (str) – The File object (not file name) to be uploaded.

• purpose (Literal['fine-tune', 'assistants', 'batch']) – The intended purpose of the uploaded file. Use “fine-tune” for Fine-tuning, “assistants” for Assistants and Messages, and “batch” for Batch API.

get_file(file_id)

Return information about a specific file.

Parameters:

file_id (str) – The ID of the file to use for this request.

get_files()

Return a list of files that belong to the user’s organization.

delete_file(file_id)

Delete a file.

Parameters:

file_id (str) – The ID of the file to be deleted.

create_vector_store(**kwargs)

Create a vector store.

get_vector_stores(**kwargs)

Return a list of vector stores.

get_vector_store(vector_store_id)

Retrieve a vector store.

Parameters:

vector_store_id (str) – The ID of the vector store to retrieve.

modify_vector_store(vector_store_id, **kwargs)

Modify a vector store.

Parameters:

vector_store_id (str) – The ID of the vector store to modify.

delete_vector_store(vector_store_id)

Delete a vector store.

Parameters:

vector_store_id (str) – The ID of the vector store to delete.

upload_files_to_vector_store(vector_store_id, files)

Upload files to a vector store and poll until completion.

Parameters:

• vector_store_id (str) – The ID of the vector store the files are to be uploaded to.

• files (list[BinaryIO]) – A list of binary files to upload.

get_vector_store_files(vector_store_id)

Return a list of vector store files.

Parameters:

vector_store_id (str)

delete_vector_store_file(vector_store_id, file_id)

Delete a vector store file. This will remove the file from the vector store but the file itself will not be deleted. To delete the file, use delete_file.

Parameters:

• vector_store_id (str) – The ID of the vector store that the file belongs to.

• file_id (str) – The ID of the file to delete.

create_batch(file_id, endpoint, metadata=None, completion_window='24h')

Create a batch for a given model and files.

Parameters:

• file_id (str) – The ID of the file to be used for this batch.

• endpoint (Literal['/v1/chat/completions', '/v1/embeddings', '/v1/completions']) – The endpoint to use for this batch. Allowed values include: ‘/v1/chat/completions’, ‘/v1/embeddings’, ‘/v1/completions’.

• metadata (dict[str, str] | None) – A set of key-value pairs that can be attached to an object.

• completion_window (Literal['24h']) – The time window for the batch to complete. Default is 24 hours.

get_batch(batch_id)

Get the status of a batch.

Parameters:

batch_id (str) – The ID of the batch to get the status of.

wait_for_batch(batch_id, wait_seconds=3, timeout=3600)

Poll a batch to check if it finishes.

Parameters:

• batch_id (str) – Id of the Batch to wait for.

• wait_seconds (float) – Optional. Number of seconds between checks.

• timeout (float) – Optional. How many seconds wait for batch to be ready. Used only if not ran in deferred operator.

cancel_batch(batch_id)

Cancel a batch.

Parameters:

batch_id (str) – The ID of the batch to delete.

Previous Next

Was this entry helpful?