Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 18 additions & 1 deletion providers/openai/docs/operators/openai.rst
Original file line number Diff line number Diff line change
Expand Up @@ -58,13 +58,30 @@ specify the OpenAI connection to use, and ``response_kwargs`` to pass through op
:start-after: [START howto_operator_openai_response]
:end-before: [END howto_operator_openai_response]

Structured outputs (Pydantic models)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

To request a structured response, pass a Pydantic ``BaseModel`` subclass as ``text_format``. The
operator then calls the Responses API's structured-output path (``responses.parse``) and returns
the parsed model as a ``dict`` (via ``model_dump(mode="json")``), which is safe to push to XCom
regardless of the field types the model uses (enums, dates, etc. are rendered as their JSON
representations). If the model refuses, the response is incomplete, or the returned JSON cannot
be coerced into ``text_format``, the operator raises ``ValueError`` with the API-reported
``status``, ``error`` and ``incomplete_details``.

.. exampleinclude:: /../../openai/tests/system/openai/example_openai.py
:language: python
:start-after: [START howto_operator_openai_response_structured]
:end-before: [END howto_operator_openai_response_structured]

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``
- Responses: ``create_response``, ``parse_response`` (structured-output wrapper),
``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
Expand Down
33 changes: 31 additions & 2 deletions providers/openai/src/airflow/providers/openai/hooks/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
import time
from enum import Enum
from functools import cached_property
from typing import TYPE_CHECKING, Any, BinaryIO, Literal
from typing import TYPE_CHECKING, Any, BinaryIO, Literal, TypeVar

from deprecated import deprecated
from openai import OpenAI
Expand Down Expand Up @@ -50,8 +50,9 @@
ChatCompletionUserMessageParam,
)
from openai.types.conversations import Conversation, ConversationDeletedResource
from openai.types.responses import Response
from openai.types.responses import ParsedResponse, Response
from openai.types.vector_stores import VectorStoreFile, VectorStoreFileBatch, VectorStoreFileDeleted
from pydantic import BaseModel
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.common.compat.module_loading import import_string
from airflow.providers.common.compat.sdk import BaseHook
Expand All @@ -66,6 +67,11 @@
"See https://platform.openai.com/docs/guides/migrate-to-responses."
)

#: Generic type variable for the Pydantic model used as the ``text_format`` in structured-output
#: Responses API calls. Mirrors the SDK's ``TextFormatT`` so ``parse_response`` returns a
#: ``ParsedResponse[T]`` — callers get ``output_parsed`` typed as ``T | None``.
_TextFormatT = TypeVar("_TextFormatT", bound="BaseModel")


class BatchStatus(str, Enum):
"""Enum for the status of a batch."""
Expand Down Expand Up @@ -248,6 +254,29 @@ def create_response(self, input: Any, model: str = "gpt-4o-mini", **kwargs: Any)
"""
return self.conn.responses.create(model=model, input=input, **kwargs)

def parse_response(
self,
input: Any,
text_format: type[_TextFormatT],
model: str = "gpt-4o-mini",
**kwargs: Any,
) -> ParsedResponse[_TextFormatT]:
"""
Create a model response and parse it into a Pydantic model via the Responses API.

Wraps :py:meth:`openai.resources.responses.Responses.parse`. The SDK converts
``text_format`` into a JSON schema, sends it as a structured-output request, and
returns a :class:`~openai.types.responses.ParsedResponse` whose ``output_parsed``
attribute is an instance of ``text_format`` (or ``None`` if the model refused).

:param input: Text, image, or file input(s) to the model.
:param text_format: A Pydantic ``BaseModel`` subclass describing the expected
structured output. The SDK converts it to a JSON schema and sends the
structured-output request.
:param model: ID of the model to use.
"""
return self.conn.responses.parse(input=input, model=model, text_format=text_format, **kwargs)

def get_response(self, response_id: str, **kwargs: Any) -> Response:
"""
Retrieve a previously created model response.
Expand Down
57 changes: 52 additions & 5 deletions providers/openai/src/airflow/providers/openai/operators/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,16 @@
from functools import cached_property
from typing import TYPE_CHECKING, Any, Literal

from pydantic import ValidationError

from airflow.providers.common.compat.sdk import BaseOperator, conf
from airflow.providers.openai.exceptions import OpenAIBatchJobException
from airflow.providers.openai.hooks.openai import OpenAIHook
from airflow.providers.openai.triggers.openai import OpenAIBatchTrigger

if TYPE_CHECKING:
from pydantic import BaseModel

from airflow.providers.common.compat.sdk import Context


Expand Down Expand Up @@ -84,16 +88,24 @@ 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.
By default the operator is synchronous and returns the response's aggregated output text.
Pass ``text_format`` (a Pydantic ``BaseModel`` subclass) to request a structured output; the
operator then returns the parsed model as a ``dict`` (via ``model_dump(mode="json")``), which
is safe to push to XCom regardless of the field types the model uses (enums, dates, etc. are
rendered as their JSON representations). For ``background=True`` responses, 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``).
(or ``parse_response`` when ``text_format`` is set) method — for example ``instructions``,
``tools``, ``conversation`` or ``previous_response_id``.
:param text_format: Optional. A Pydantic ``BaseModel`` subclass describing the expected
structured output. When set, the operator calls ``parse_response`` and returns
``output_parsed.model_dump(mode="json")``; otherwise it calls ``create_response`` and
returns ``output_text``.

.. seealso::
For more information on how to use this operator, take a look at the guide:
Expand All @@ -110,20 +122,55 @@ def __init__(
input_text: str | list[Any],
model: str = "gpt-4o-mini",
response_kwargs: dict | None = None,
text_format: type[BaseModel] | 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 {}
self.text_format = text_format

@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:
def execute(self, context: Context) -> str | dict[str, Any]:
if self.text_format is not None:
try:
parsed = self.hook.parse_response(
input=self.input_text,
model=self.model,
text_format=self.text_format,
**self.response_kwargs,
)
except ValidationError as exc:
# ``responses.parse`` raises ``ValidationError`` when the model's JSON output
# can't be coerced into ``text_format`` — most commonly because the response
# was truncated (e.g. ``max_output_tokens`` hit) mid-JSON. Convert to a clean
# ``ValueError`` so callers see a consistent shape across all parse failures.
raise ValueError(
f"OpenAI Responses API returned a payload that does not match "
f"{self.text_format.__name__!r}: {exc}"
) from exc

self.log.info("Generated response %s", parsed.id)
if parsed.output_parsed is None:
# No structured output — the model refused, the request errored, or the response
# was incomplete. Surface a clear error so downstream tasks don't get ``None``,
# and include what the API already told us so the user does not need a follow-up
# ``OpenAIHook.get_response`` call.
details: list[str] = [f"status={parsed.status!r}"]
if parsed.error is not None:
details.append(f"error={parsed.error!r}")
if parsed.incomplete_details is not None:
details.append(f"incomplete_details={parsed.incomplete_details!r}")
raise ValueError(
f"Response {parsed.id} did not return a structured output ({', '.join(details)})."
)
return parsed.output_parsed.model_dump(mode="json")
response = self.hook.create_response(input=self.input_text, model=self.model, **self.response_kwargs)
if response.status != "completed":
self.log.warning(
Expand Down
14 changes: 14 additions & 0 deletions providers/openai/tests/system/openai/example_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from __future__ import annotations

import pendulum
from pydantic import BaseModel

# This example uses common.compat for Airflow 2.x/3.x compatibility.
# If you only need Airflow 3+, you can use: from airflow.sdk import dag, task
Expand Down Expand Up @@ -109,6 +110,19 @@ def task_to_store_input_text_in_xcom():
)
# [END howto_operator_openai_response]

# [START howto_operator_openai_response_structured]
class Person(BaseModel):
name: str
age: int

OpenAIResponseOperator(
task_id="openai_response_structured",
conn_id="openai_default",
input_text="Extract the name and age from: 'Alice is 30 years old'.",
text_format=Person,
)
# [END howto_operator_openai_response_structured]

create_embeddings_using_hook()


Expand Down
21 changes: 21 additions & 0 deletions providers/openai/tests/unit/openai/hooks/test_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
from openai.types.beta.threads import Message, Run
from openai.types.chat import ChatCompletion
from openai.types.vector_stores import VectorStoreFile, VectorStoreFileBatch, VectorStoreFileDeleted
from pydantic import BaseModel

from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.models import Connection
Expand Down Expand Up @@ -315,6 +316,26 @@ def test_create_response(mock_openai_hook):
assert result is expected


def test_parse_response(mock_openai_hook):
class Person(BaseModel):
name: str

expected = mock_openai_hook.conn.responses.parse.return_value
result = mock_openai_hook.parse_response(
input="Extract: Alice",
text_format=Person,
model=MODEL,
instructions="Be precise.",
)
mock_openai_hook.conn.responses.parse.assert_called_once_with(
model=MODEL,
input="Extract: Alice",
text_format=Person,
instructions="Be precise.",
)
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")
Expand Down
Loading