• Home
• airflow.providers.openai
• airflow.providers.openai.operators
• airflow.providers.openai.operators.openai

airflow.providers.openai.operators.openai

Classes

OpenAIEmbeddingOperator	Operator that accepts input text to generate OpenAI embeddings using the specified model.
OpenAITriggerBatchOperator	Operator that triggers an OpenAI Batch API endpoint and waits for the batch to complete.

Module Contents

class airflow.providers.openai.operators.openai.OpenAIEmbeddingOperator(conn_id, input_text, model='text-embedding-ada-002', embedding_kwargs=None, **kwargs)

Bases: airflow.models.BaseOperator

Operator that accepts input text to generate OpenAI embeddings using the specified model.

Parameters:

• conn_id (str) – The OpenAI connection ID to use.

• input_text (str | list[str] | list[int] | list[list[int]]) – The text to generate OpenAI embeddings for. This can be a string, a list of strings, a list of integers, or a list of lists of integers.

• model (str) – The OpenAI model to be used for generating the embeddings.

• embedding_kwargs (dict | None) – Additional keyword arguments to pass to the OpenAI create_embeddings method.

See also

For more information on how to use this operator, take a look at the guide: OpenAIEmbeddingOperator For possible options for embedding_kwargs, see: https://platform.openai.com/docs/api-reference/embeddings/create

template_fields: collections.abc.Sequence[str] = ('input_text',)

conn_id

input_text

model = 'text-embedding-ada-002'

embedding_kwargs

property hook: airflow.providers.openai.hooks.openai.OpenAIHook

Return an instance of the OpenAIHook.

execute(context)

Derive when creating an operator.

Context is the same dictionary used as when rendering jinja templates.

Refer to get_template_context for more context.

class airflow.providers.openai.operators.openai.OpenAITriggerBatchOperator(file_id, endpoint, conn_id=OpenAIHook.default_conn_name, deferrable=conf.getboolean('operators', 'default_deferrable', fallback=False), wait_seconds=3, timeout=24 * 60 * 60, wait_for_completion=True, **kwargs)

Bases: airflow.models.BaseOperator

Operator that triggers an OpenAI Batch API endpoint and waits for the batch to complete.

Parameters:

• file_id (str) – Required. The ID of the batch file to trigger.

• endpoint (Literal['/v1/chat/completions', '/v1/embeddings', '/v1/completions']) – Required. The OpenAI Batch API endpoint to trigger.

• conn_id (str) – Optional. The OpenAI connection ID to use. Defaults to ‘openai_default’.

• deferrable (bool) – Optional. Run operator in the deferrable mode.

• wait_seconds (float) – Optional. Number of seconds between checks. Only used when deferrable is False. Defaults to 3 seconds.

• timeout (float) – Optional. The amount of time, in seconds, to wait for the request to complete. Only used when deferrable is False. Defaults to 24 hour, which is the SLA for OpenAI Batch API.

• wait_for_completion (bool) – Optional. Whether to wait for the batch to complete. If set to False, the operator will return immediately after triggering the batch. Defaults to True.

See also

For more information on how to use this operator, please take a look at the guide: OpenAITriggerBatchOperator

template_fields: collections.abc.Sequence[str] = ('file_id',)

conn_id = 'openai_default'

file_id

endpoint

deferrable = True

wait_seconds = 3

timeout = 86400

wait_for_completion = True

batch_id: str | None = None

property hook: airflow.providers.openai.hooks.openai.OpenAIHook

Return an instance of the OpenAIHook.

execute(context)

Derive when creating an operator.

Context is the same dictionary used as when rendering jinja templates.

Refer to get_template_context for more context.

execute_complete(context, event=None)

Invoke this callback when the trigger fires; return immediately.

Relies on trigger to throw an exception, otherwise it assumes execution was successful.

on_kill()

Cancel the batch if task is cancelled.

Previous Next

Was this entry helpful?