OpenAIEmbeddingOperator

Use the OpenAIEmbeddingOperator to interact with Open APIs to create embeddings for given text.

Using the Operator

The OpenAIEmbeddingOperator requires the input_text as an input to embedding API. Use the conn_id parameter to specify the OpenAI connection to use to connect to your account.

An example using the operator is in way:

tests/system/openai/example_openai.py[source]

    OpenAIEmbeddingOperator(
        task_id="embedding_using_xcom_data",
        conn_id="openai_default",
        input_text=task_to_store_input_text_in_xcom(),
        model="text-embedding-ada-002",
    )

    OpenAIEmbeddingOperator(
        task_id="embedding_using_callable",
        conn_id="openai_default",
        input_text=input_text_callable(
            "input_arg1_value",
            "input2_value",
            input_kwarg1="input_kwarg1_value",
            input_kwarg2="input_kwarg2_value",
        ),
        model="text-embedding-ada-002",
    )
    OpenAIEmbeddingOperator(
        task_id="embedding_using_text",
        conn_id="openai_default",
        input_text=texts,
        model="text-embedding-ada-002",
    )

OpenAITriggerBatchOperator

Use the OpenAITriggerBatchOperator to interact with Open APIs to trigger a batch job. This operator is used to trigger a batch job and wait for the job to complete.

Using the Operator

The OpenAITriggerBatchOperator requires the prepared batch file as an input to trigger the batch job. Provide the file_id and the endpoint to trigger the batch job. Use the conn_id parameter to specify the OpenAI connection to use to

The OpenAITriggerBatchOperator

An example using the operator is in way:

tests/system/openai/example_trigger_batch_operator.py[source]

    from airflow.providers.openai.operators.openai import OpenAITriggerBatchOperator

    batch_id = OpenAITriggerBatchOperator(
        task_id="batch_operator_deferred",
        conn_id=OPENAI_CONN_ID,
        file_id=batch_file_id,
        endpoint="/v1/chat/completions",
        deferrable=True,
    )

Was this entry helpful?