This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ff471253f18 Add OpenAI Responses and Conversations support (#69070)
ff471253f18 is described below
commit ff471253f18f7652988f029cfece652d1cfaf440
Author: Kaxil Naik <[email protected]>
AuthorDate: Sat Jun 27 03:39:42 2026 +0100
Add OpenAI Responses and Conversations support (#69070)
---
providers/openai/docs/operators/openai.rst | 41 ++++++++++++++
.../src/airflow/providers/openai/hooks/openai.py | 64 ++++++++++++++++++++++
.../airflow/providers/openai/operators/openai.py | 55 +++++++++++++++++++
.../openai/tests/system/openai/example_openai.py | 11 +++-
.../openai/tests/unit/openai/hooks/test_openai.py | 62 ++++++++++++++++++++-
.../tests/unit/openai/operators/test_openai.py | 39 +++++++++++--
.../openai/tests/unit/openai/test_exceptions.py | 3 +-
7 files changed, 267 insertions(+), 8 deletions(-)
diff --git a/providers/openai/docs/operators/openai.rst
b/providers/openai/docs/operators/openai.rst
index 10d220f6d8e..d7f88d14d1c 100644
--- a/providers/openai/docs/operators/openai.rst
+++ b/providers/openai/docs/operators/openai.rst
@@ -37,6 +37,47 @@ An example using the operator is in way:
:start-after: [START howto_operator_openai_embedding]
:end-before: [END howto_operator_openai_embedding]
+.. _howto/operator:OpenAIResponseOperator:
+
+OpenAIResponseOperator
+=======================
+
+Use the
:class:`~airflow.providers.openai.operators.openai.OpenAIResponseOperator` to
generate a
+model response with the OpenAI Responses API, OpenAI's recommended interface
for text generation and
+tool use. The operator returns the response's aggregated output text.
+
+Using the Operator
+^^^^^^^^^^^^^^^^^^^
+
+The OpenAIResponseOperator requires the ``input_text`` prompt. Use the
``conn_id`` parameter to
+specify the OpenAI connection to use, and ``response_kwargs`` to pass through
options such as
+``tools``, ``conversation`` or ``previous_response_id``.
+
+.. exampleinclude:: /../../openai/tests/system/openai/example_openai.py
+ :language: python
+ :start-after: [START howto_operator_openai_response]
+ :end-before: [END howto_operator_openai_response]
+
+Using the OpenAIHook for Responses and Conversations
+=====================================================
+
+The :class:`~airflow.providers.openai.hooks.openai.OpenAIHook` exposes the
Responses and
+Conversations APIs directly for use inside ``@task`` functions or custom
operators:
+
+- Responses: ``create_response``, ``get_response``, ``delete_response`` and
``cancel_response``
+ (the last cancels a response created with ``background=True``).
+- Conversations: ``create_conversation``, ``get_conversation``,
``update_conversation`` and
+ ``delete_conversation``. Pass the conversation id to ``create_response``
(via the operator's
+ ``response_kwargs`` or the hook) to persist state across responses.
+
+For example, to create a conversation and continue it across responses:
+
+.. code-block:: python
+
+ hook = OpenAIHook()
+ conversation = hook.create_conversation()
+ hook.create_response(input="Hello", conversation=conversation.id)
+
.. _howto/operator:OpenAITriggerBatchOperator:
OpenAITriggerBatchOperator
diff --git a/providers/openai/src/airflow/providers/openai/hooks/openai.py
b/providers/openai/src/airflow/providers/openai/hooks/openai.py
index ca80b4d3761..fbebadd493e 100644
--- a/providers/openai/src/airflow/providers/openai/hooks/openai.py
+++ b/providers/openai/src/airflow/providers/openai/hooks/openai.py
@@ -48,6 +48,8 @@ if TYPE_CHECKING:
ChatCompletionToolMessageParam,
ChatCompletionUserMessageParam,
)
+ from openai.types.conversations import Conversation,
ConversationDeletedResource
+ from openai.types.responses import Response
from openai.types.vector_stores import VectorStoreFile,
VectorStoreFileBatch, VectorStoreFileDeleted
from airflow.providers.common.compat.module_loading import import_string
from airflow.providers.common.compat.sdk import BaseHook
@@ -226,6 +228,68 @@ class OpenAIHook(BaseHook):
response = self.conn.chat.completions.create(model=model,
messages=messages, **kwargs)
return response.choices
+ def create_response(self, input: Any, model: str = "gpt-4o-mini",
**kwargs: Any) -> Response:
+ """
+ Create a model response using the Responses API.
+
+ :param input: Text, image, or file input(s) to the model.
+ :param model: ID of the model to use.
+ """
+ return self.conn.responses.create(model=model, input=input, **kwargs)
+
+ def get_response(self, response_id: str, **kwargs: Any) -> Response:
+ """
+ Retrieve a previously created model response.
+
+ :param response_id: The ID of the response to retrieve.
+ """
+ return self.conn.responses.retrieve(response_id, **kwargs)
+
+ def delete_response(self, response_id: str) -> None:
+ """
+ Delete a model response.
+
+ :param response_id: The ID of the response to delete.
+ """
+ self.conn.responses.delete(response_id)
+
+ def cancel_response(self, response_id: str) -> Response:
+ """
+ Cancel an in-progress response created with ``background=True``.
+
+ :param response_id: The ID of the response to cancel.
+ """
+ return self.conn.responses.cancel(response_id)
+
+ def create_conversation(self, **kwargs: Any) -> Conversation:
+ """Create a conversation that can be reused across responses to
persist state."""
+ return self.conn.conversations.create(**kwargs)
+
+ def get_conversation(self, conversation_id: str) -> Conversation:
+ """
+ Retrieve a conversation.
+
+ :param conversation_id: The ID of the conversation to retrieve.
+ """
+ return self.conn.conversations.retrieve(conversation_id)
+
+ def update_conversation(self, conversation_id: str, metadata: dict[str,
str]) -> Conversation:
+ """
+ Update a conversation's metadata.
+
+ :param conversation_id: The ID of the conversation to update.
+ :param metadata: Set of key-value pairs to attach to the conversation.
+ """
+ return self.conn.conversations.update(conversation_id,
metadata=metadata)
+
+ def delete_conversation(self, conversation_id: str) ->
ConversationDeletedResource:
+ """
+ Delete a conversation.
+
+ :param conversation_id: The ID of the conversation to delete.
+ """
+ return self.conn.conversations.delete(conversation_id)
+
def create_assistant(self, model: str = "gpt-4o-mini", **kwargs: Any) ->
Assistant:
"""
Create an OpenAI assistant using the given model.
diff --git a/providers/openai/src/airflow/providers/openai/operators/openai.py
b/providers/openai/src/airflow/providers/openai/operators/openai.py
index 192986ab456..0f040a03b55 100644
--- a/providers/openai/src/airflow/providers/openai/operators/openai.py
+++ b/providers/openai/src/airflow/providers/openai/operators/openai.py
@@ -80,6 +80,61 @@ class OpenAIEmbeddingOperator(BaseOperator):
return embeddings
+class OpenAIResponseOperator(BaseOperator):
+ """
+ Operator that generates a model response using the OpenAI Responses API.
+
+ The operator is synchronous and returns the response's aggregated output
text. For
+ ``previous_response_id`` chaining, ``background=True`` responses, or
access to the full
+ structured response, use
:class:`~airflow.providers.openai.hooks.openai.OpenAIHook` directly.
+
+ :param conn_id: The OpenAI connection ID to use.
+ :param input_text: The input prompt for the model. This can be a string or
a structured list of
+ input items.
+ :param model: The OpenAI model to use.
+ :param response_kwargs: Additional keyword arguments to pass to the OpenAI
``create_response``
+ method (for example ``instructions``, ``tools``, ``conversation`` or
``previous_response_id``).
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:OpenAIResponseOperator`
+ For possible options, see:
+ https://platform.openai.com/docs/api-reference/responses/create
+ """
+
+ template_fields: Sequence[str] = ("input_text",)
+
+ def __init__(
+ self,
+ conn_id: str,
+ input_text: str | list[Any],
+ model: str = "gpt-4o-mini",
+ response_kwargs: dict | None = None,
+ **kwargs: Any,
+ ):
+ super().__init__(**kwargs)
+ self.conn_id = conn_id
+ self.input_text = input_text
+ self.model = model
+ self.response_kwargs = response_kwargs or {}
+
+ @cached_property
+ def hook(self) -> OpenAIHook:
+ """Return an instance of the OpenAIHook."""
+ return OpenAIHook(conn_id=self.conn_id)
+
+ def execute(self, context: Context) -> str:
+ response = self.hook.create_response(input=self.input_text,
model=self.model, **self.response_kwargs)
+ if response.status != "completed":
+ self.log.warning(
+ "Response %s ended with status %s; the returned output text
may be empty.",
+ response.id,
+ response.status,
+ )
+ self.log.info("Generated response %s", response.id)
+ return response.output_text
+
+
class OpenAITriggerBatchOperator(BaseOperator):
"""
Operator that triggers an OpenAI Batch API endpoint and waits for the
batch to complete.
diff --git a/providers/openai/tests/system/openai/example_openai.py
b/providers/openai/tests/system/openai/example_openai.py
index 6cbbe2b47f3..e03bc639745 100644
--- a/providers/openai/tests/system/openai/example_openai.py
+++ b/providers/openai/tests/system/openai/example_openai.py
@@ -22,7 +22,7 @@ import pendulum
# If you only need Airflow 3+, you can use: from airflow.sdk import dag, task
from airflow.providers.common.compat.sdk import dag, task
from airflow.providers.openai.hooks.openai import OpenAIHook
-from airflow.providers.openai.operators.openai import OpenAIEmbeddingOperator
+from airflow.providers.openai.operators.openai import OpenAIEmbeddingOperator,
OpenAIResponseOperator
def input_text_callable(
@@ -100,6 +100,15 @@ def example_openai_dag():
)
# [END howto_operator_openai_embedding]
+ # [START howto_operator_openai_response]
+ OpenAIResponseOperator(
+ task_id="openai_response",
+ conn_id="openai_default",
+ input_text="Write a haiku about data pipelines.",
+ response_kwargs={"instructions": "You are a helpful assistant."},
+ )
+ # [END howto_operator_openai_response]
+
create_embeddings_using_hook()
diff --git a/providers/openai/tests/unit/openai/hooks/test_openai.py
b/providers/openai/tests/unit/openai/hooks/test_openai.py
index 6cee1244e2c..9b4760953d4 100644
--- a/providers/openai/tests/unit/openai/hooks/test_openai.py
+++ b/providers/openai/tests/unit/openai/hooks/test_openai.py
@@ -17,9 +17,10 @@
from __future__ import annotations
import os
-from unittest.mock import mock_open, patch
+from unittest.mock import MagicMock, mock_open, patch
import pytest
+from openai import OpenAI
from openai.pagination import SyncCursorPage
from openai.types import (
Batch,
@@ -82,7 +83,10 @@ def mock_openai_connection():
@pytest.fixture
def mock_openai_hook(mock_openai_connection):
- with patch("airflow.providers.openai.hooks.openai.OpenAI"):
+ # spec=OpenAI guards top-level namespace access (an unknown attribute on
the client fails);
+ # the method surface (``conn.responses.create``) is type-checked by mypy
against the SDK stubs.
+ with patch("airflow.providers.openai.hooks.openai.OpenAI") as mock_client:
+ mock_client.return_value = MagicMock(spec=OpenAI)
yield OpenAIHook(conn_id=mock_openai_connection.conn_id)
@@ -303,6 +307,60 @@ def test_create_chat_completion(mock_openai_hook,
mock_completion):
assert choice.message.content == "Hello there, how may I assist you today?"
+def test_create_response(mock_openai_hook):
+ expected = mock_openai_hook.conn.responses.create.return_value
+ result = mock_openai_hook.create_response(input="Hello", model=MODEL)
+
mock_openai_hook.conn.responses.create.assert_called_once_with(model=MODEL,
input="Hello")
+ assert result is expected
+
+
+def test_get_response(mock_openai_hook):
+ expected = mock_openai_hook.conn.responses.retrieve.return_value
+ result = mock_openai_hook.get_response("resp_123")
+
mock_openai_hook.conn.responses.retrieve.assert_called_once_with("resp_123")
+ assert result is expected
+
+
+def test_delete_response(mock_openai_hook):
+ mock_openai_hook.delete_response("resp_123")
+ mock_openai_hook.conn.responses.delete.assert_called_once_with("resp_123")
+
+
+def test_cancel_response(mock_openai_hook):
+ expected = mock_openai_hook.conn.responses.cancel.return_value
+ result = mock_openai_hook.cancel_response("resp_123")
+ mock_openai_hook.conn.responses.cancel.assert_called_once_with("resp_123")
+ assert result is expected
+
+
+def test_create_conversation(mock_openai_hook):
+ expected = mock_openai_hook.conn.conversations.create.return_value
+ result = mock_openai_hook.create_conversation(metadata={"topic": "demo"})
+
mock_openai_hook.conn.conversations.create.assert_called_once_with(metadata={"topic":
"demo"})
+ assert result is expected
+
+
+def test_get_conversation(mock_openai_hook):
+ expected = mock_openai_hook.conn.conversations.retrieve.return_value
+ result = mock_openai_hook.get_conversation("conv_123")
+
mock_openai_hook.conn.conversations.retrieve.assert_called_once_with("conv_123")
+ assert result is expected
+
+
+def test_update_conversation(mock_openai_hook):
+ expected = mock_openai_hook.conn.conversations.update.return_value
+ result = mock_openai_hook.update_conversation("conv_123",
metadata={"topic": "demo"})
+
mock_openai_hook.conn.conversations.update.assert_called_once_with("conv_123",
metadata={"topic": "demo"})
+ assert result is expected
+
+
+def test_delete_conversation(mock_openai_hook):
+ expected = mock_openai_hook.conn.conversations.delete.return_value
+ result = mock_openai_hook.delete_conversation("conv_123")
+
mock_openai_hook.conn.conversations.delete.assert_called_once_with("conv_123")
+ assert result is expected
+
+
def test_create_assistant(mock_openai_hook, mock_assistant):
mock_openai_hook.conn.beta.assistants.create.return_value = mock_assistant
assistant = mock_openai_hook.create_assistant(
diff --git a/providers/openai/tests/unit/openai/operators/test_openai.py
b/providers/openai/tests/unit/openai/operators/test_openai.py
index f85f3dd3d70..954306d42a5 100644
--- a/providers/openai/tests/unit/openai/operators/test_openai.py
+++ b/providers/openai/tests/unit/openai/operators/test_openai.py
@@ -20,9 +20,15 @@ from unittest.mock import Mock
import pytest
from openai.types.batch import Batch
+from openai.types.responses import Response
from airflow.providers.common.compat.sdk import Context, TaskDeferred
-from airflow.providers.openai.operators.openai import OpenAIEmbeddingOperator,
OpenAITriggerBatchOperator
+from airflow.providers.openai.hooks.openai import OpenAIHook
+from airflow.providers.openai.operators.openai import (
+ OpenAIEmbeddingOperator,
+ OpenAIResponseOperator,
+ OpenAITriggerBatchOperator,
+)
from airflow.providers.openai.triggers.openai import OpenAIBatchTrigger
openai = pytest.importorskip("openai")
@@ -50,7 +56,7 @@ def test_execute_with_input_text():
operator = OpenAIEmbeddingOperator(
task_id=TASK_ID, conn_id=CONN_ID, model="test_model", input_text="Test
input text"
)
- mock_hook_instance = Mock()
+ mock_hook_instance = Mock(spec=OpenAIHook)
mock_hook_instance.create_embeddings.return_value = [1.0, 2.0, 3.0]
operator.hook = mock_hook_instance
@@ -73,6 +79,31 @@ def test_execute_with_invalid_input(invalid_input):
operator.execute(context)
+def test_openai_response_operator_execute():
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ model="test_model",
+ response_kwargs={"instructions": "Be concise.",
"previous_response_id": "resp_prev"},
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_hook_instance.create_response.return_value = Mock(
+ spec=Response, output_text="haiku text", id="resp_123",
status="completed"
+ )
+ operator.hook = mock_hook_instance
+
+ result = operator.execute(Context())
+
+ assert result == "haiku text"
+ mock_hook_instance.create_response.assert_called_once_with(
+ input="Write a haiku.",
+ model="test_model",
+ instructions="Be concise.",
+ previous_response_id="resp_prev",
+ )
+
+
@pytest.mark.parametrize("wait_for_completion", [True, False])
def test_openai_trigger_batch_operator_not_deferred(mock_batch,
wait_for_completion):
operator = OpenAITriggerBatchOperator(
@@ -83,7 +114,7 @@ def
test_openai_trigger_batch_operator_not_deferred(mock_batch, wait_for_complet
wait_for_completion=wait_for_completion,
deferrable=False,
)
- mock_hook_instance = Mock()
+ mock_hook_instance = Mock(spec=OpenAIHook)
mock_hook_instance.get_batch.return_value = mock_batch
mock_hook_instance.create_batch.return_value = mock_batch
operator.hook = mock_hook_instance
@@ -103,7 +134,7 @@ def
test_openai_trigger_batch_operator_with_deferred(mock_batch, wait_for_comple
deferrable=True,
wait_for_completion=wait_for_completion,
)
- mock_hook_instance = Mock()
+ mock_hook_instance = Mock(spec=OpenAIHook)
mock_hook_instance.get_batch.return_value = mock_batch
mock_hook_instance.create_batch.return_value = mock_batch
operator.hook = mock_hook_instance
diff --git a/providers/openai/tests/unit/openai/test_exceptions.py
b/providers/openai/tests/unit/openai/test_exceptions.py
index 38f6d38a509..fabaad35343 100644
--- a/providers/openai/tests/unit/openai/test_exceptions.py
+++ b/providers/openai/tests/unit/openai/test_exceptions.py
@@ -22,6 +22,7 @@ from unittest.mock import Mock
import pytest
from airflow.providers.openai.exceptions import OpenAIBatchJobException,
OpenAIBatchTimeout
+from airflow.providers.openai.hooks.openai import OpenAIHook
@pytest.mark.parametrize(
@@ -32,7 +33,7 @@ from airflow.providers.openai.exceptions import
OpenAIBatchJobException, OpenAIB
],
)
def test_wait_for_batch_raise_exception(exception_class):
- mock_hook_instance = Mock()
+ mock_hook_instance = Mock(spec=OpenAIHook)
mock_hook_instance.wait_for_batch.side_effect = exception_class
hook = mock_hook_instance
with pytest.raises(exception_class):
