airflow.providers.openai.hooks.openai

Module Contents

Classes

OpenAIHook

Use OpenAI SDK to interact with OpenAI APIs.
class airflow.providers.openai.hooks.openai.OpenAIHook(conn_id=default_conn_name, *args, **kwargs)[source]

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'[source]
default_conn_name = 'openai_default'[source]
conn_type = 'openai'[source]
hook_name = 'OpenAI'[source]
classmethod get_ui_field_behaviour()[source]

Return custom field behaviour.

test_connection()[source]
conn()[source]

Return an OpenAI connection object.

get_conn()[source]

Return an OpenAI connection object.

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

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)[source]

Create an OpenAI assistant using the given model.

Parameters

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

get_assistant(assistant_id)[source]

Get an OpenAI assistant.

Parameters

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

get_assistants(**kwargs)[source]

Get a list of Assistant objects.

modify_assistant(assistant_id, **kwargs)[source]

Modify an existing Assistant object.

Parameters

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

delete_assistant(assistant_id)[source]

Delete an OpenAI Assistant for a given ID.

Parameters

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

create_thread(**kwargs)[source]

Create an OpenAI thread.

modify_thread(thread_id, metadata)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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]) – The intended purpose of the uploaded file. Use “fine-tune” for Fine-tuning and “assistants” for Assistants and Messages.

get_file(file_id)[source]

Return information about a specific file.

Parameters

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

get_files()[source]

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

delete_file(file_id)[source]

Delete a file.

Parameters

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

create_vector_store(**kwargs)[source]

Create a vector store.

get_vector_stores(**kwargs)[source]

Return a list of vector stores.

get_vector_store(vector_store_id)[source]

Retrieve a vector store.

Parameters

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

modify_vector_store(vector_store_id, **kwargs)[source]

Modify a vector store.

Parameters

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

delete_vector_store(vector_store_id)[source]

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)[source]

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)[source]

Return a list of vector store files.

Parameters

vector_store_id (str) –

delete_vector_store_file(vector_store_id, file_id)[source]

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.

Was this entry helpful?