From 1d333583183015e72b30a97f11d0be1384a282c5 Mon Sep 17 00:00:00 2001 From: Yash Jain <23dec512@lnmiit.ac.in> Date: Mon, 13 Jul 2026 15:36:10 +0530 Subject: [PATCH 1/4] Support Pydantic structured outputs in OpenAIResponseOperator The OpenAIResponseOperator currently returns only the aggregated output text of a Responses API call. Users who want a structured JSON output -- the pattern every AI-eng team uses for reliable extraction, tool selection, and downstream DAG chaining -- have to bypass the operator and call the hook directly, as the operator's own docstring instructs. Pass a Pydantic BaseModel subclass as text_format and the operator now uses the SDK's structured-output path (responses.parse), returning output_parsed.model_dump() -- a plain dict, XCom-safe, no manual serialization dance. When text_format is None the operator's existing plain-text behavior is preserved bit-for-bit. If the model refuses or produces a non-parseable response, the operator raises a clear ValueError so downstream tasks do not silently receive None. --- .../airflow/providers/openai/hooks/openai.py | 26 ++++++++++++- .../providers/openai/operators/openai.py | 39 ++++++++++++++++--- 2 files changed, 59 insertions(+), 6 deletions(-) diff --git a/providers/openai/src/airflow/providers/openai/hooks/openai.py b/providers/openai/src/airflow/providers/openai/hooks/openai.py index 8315a907ffa6c..2e27ac5f3dd3e 100644 --- a/providers/openai/src/airflow/providers/openai/hooks/openai.py +++ b/providers/openai/src/airflow/providers/openai/hooks/openai.py @@ -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 @@ -248,6 +249,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[BaseModel], + model: str = "gpt-4o-mini", + **kwargs: Any, + ) -> ParsedResponse[Any]: + """ + 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. Requires a model that supports structured outputs + (``gpt-4o-2024-08-06`` and later). + :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. diff --git a/providers/openai/src/airflow/providers/openai/operators/openai.py b/providers/openai/src/airflow/providers/openai/operators/openai.py index 0f040a03b5505..3a86898e3d0d4 100644 --- a/providers/openai/src/airflow/providers/openai/operators/openai.py +++ b/providers/openai/src/airflow/providers/openai/operators/openai.py @@ -28,6 +28,8 @@ from airflow.providers.openai.triggers.openai import OpenAIBatchTrigger if TYPE_CHECKING: + from pydantic import BaseModel + from airflow.providers.common.compat.sdk import Context @@ -84,16 +86,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()``), which is safe to + push to XCom. Requires a model that supports structured outputs (``gpt-4o-2024-08-06`` and + later). For ``previous_response_id`` chaining or ``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()``; 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: @@ -110,6 +120,7 @@ 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) @@ -117,13 +128,31 @@ def __init__( 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: + parsed = self.hook.parse_response( + input=self.input_text, + model=self.model, + text_format=self.text_format, + **self.response_kwargs, + ) + 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 truncated. Surface a clear error so downstream tasks don't get None. + raise ValueError( + f"Response {parsed.id} did not produce a parseable structured output " + f"(status={parsed.status!r}). Inspect the response via OpenAIHook.get_response " + f"for refusal / error details." + ) + return parsed.output_parsed.model_dump() response = self.hook.create_response(input=self.input_text, model=self.model, **self.response_kwargs) if response.status != "completed": self.log.warning( From fbe7c8f3746c00bc2abf37b1e8b23decef69f736 Mon Sep 17 00:00:00 2001 From: Yash Jain <23dec512@lnmiit.ac.in> Date: Mon, 13 Jul 2026 15:36:37 +0530 Subject: [PATCH 2/4] Add unit tests for OpenAIResponseOperator structured outputs Covers the hook.parse_response wrapper (arg forwarding), the operator happy path (parsed model returned as dict, create_response not called), and the refusal path (output_parsed=None raises ValueError). The existing plain-text test is left untouched to guard the backward-compat path. --- .../tests/system/openai/example_openai.py | 16 +++++ .../tests/unit/openai/hooks/test_openai.py | 22 +++++++ .../unit/openai/operators/test_openai.py | 60 ++++++++++++++++++- 3 files changed, 97 insertions(+), 1 deletion(-) diff --git a/providers/openai/tests/system/openai/example_openai.py b/providers/openai/tests/system/openai/example_openai.py index e03bc6397456f..3c7d5ab654c3e 100644 --- a/providers/openai/tests/system/openai/example_openai.py +++ b/providers/openai/tests/system/openai/example_openai.py @@ -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 @@ -109,6 +110,21 @@ 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", + # Structured outputs require a compatible model (gpt-4o-2024-08-06 or later). + model="gpt-4o-2024-08-06", + 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() diff --git a/providers/openai/tests/unit/openai/hooks/test_openai.py b/providers/openai/tests/unit/openai/hooks/test_openai.py index 5e5882e354869..e9353696dda42 100644 --- a/providers/openai/tests/unit/openai/hooks/test_openai.py +++ b/providers/openai/tests/unit/openai/hooks/test_openai.py @@ -315,6 +315,28 @@ def test_create_response(mock_openai_hook): assert result is expected +def test_parse_response(mock_openai_hook): + from pydantic import BaseModel + + 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") diff --git a/providers/openai/tests/unit/openai/operators/test_openai.py b/providers/openai/tests/unit/openai/operators/test_openai.py index 954306d42a5eb..aa84a720ba524 100644 --- a/providers/openai/tests/unit/openai/operators/test_openai.py +++ b/providers/openai/tests/unit/openai/operators/test_openai.py @@ -20,7 +20,8 @@ import pytest from openai.types.batch import Batch -from openai.types.responses import Response +from openai.types.responses import ParsedResponse, Response +from pydantic import BaseModel from airflow.providers.common.compat.sdk import Context, TaskDeferred from airflow.providers.openai.hooks.openai import OpenAIHook @@ -104,6 +105,63 @@ def test_openai_response_operator_execute(): ) +class _StructuredPerson(BaseModel): + """Pydantic model used by the structured-output operator tests.""" + + name: str + + +def test_openai_response_operator_structured_output_returns_dict(): + operator = OpenAIResponseOperator( + task_id=TASK_ID, + conn_id=CONN_ID, + input_text="Extract: Alice", + model="test_model", + text_format=_StructuredPerson, + response_kwargs={"instructions": "Be precise."}, + ) + mock_hook_instance = Mock(spec=OpenAIHook) + mock_hook_instance.parse_response.return_value = Mock( + spec=ParsedResponse, + id="resp_str_1", + status="completed", + output_parsed=_StructuredPerson(name="Alice"), + ) + operator.hook = mock_hook_instance + + result = operator.execute(Context()) + + assert result == {"name": "Alice"} + mock_hook_instance.parse_response.assert_called_once_with( + input="Extract: Alice", + model="test_model", + text_format=_StructuredPerson, + instructions="Be precise.", + ) + mock_hook_instance.create_response.assert_not_called() + + +def test_openai_response_operator_structured_output_refusal_raises(): + operator = OpenAIResponseOperator( + task_id=TASK_ID, + conn_id=CONN_ID, + input_text="Extract: Alice", + model="test_model", + text_format=_StructuredPerson, + ) + mock_hook_instance = Mock(spec=OpenAIHook) + mock_hook_instance.parse_response.return_value = Mock( + spec=ParsedResponse, + id="resp_refused", + status="incomplete", + output_parsed=None, + ) + operator.hook = mock_hook_instance + + with pytest.raises(ValueError, match="did not produce a parseable structured output"): + operator.execute(Context()) + + @pytest.mark.parametrize("wait_for_completion", [True, False]) def test_openai_trigger_batch_operator_not_deferred(mock_batch, wait_for_completion): operator = OpenAITriggerBatchOperator( From a6c015210012f88e120ab4cf813925c24e8885fa Mon Sep 17 00:00:00 2001 From: Yash Jain <23dec512@lnmiit.ac.in> Date: Mon, 13 Jul 2026 15:36:55 +0530 Subject: [PATCH 3/4] Document structured outputs in OpenAIResponseOperator how-to guide Adds a Structured outputs (Pydantic models) subsection under the existing OpenAIResponseOperator how-to, with an exampleinclude pointing at the new [START/END] howto_operator_openai_response_structured markers in the example DAG. Calls out the compatible-model requirement (gpt-4o-2024-08-06 or later) and the ValueError raised on refusal. --- providers/openai/docs/operators/openai.rst | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/providers/openai/docs/operators/openai.rst b/providers/openai/docs/operators/openai.rst index bfb2fd8dee94b..d26910e7fd32f 100644 --- a/providers/openai/docs/operators/openai.rst +++ b/providers/openai/docs/operators/openai.rst @@ -58,6 +58,20 @@ 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()``), which is safe to push to XCom. This requires +a model that supports structured outputs (``gpt-4o-2024-08-06`` or later). If the model refuses or +returns non-parseable content, the operator raises ``ValueError``. + +.. 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 ===================================================== From a56f7041001c44ac51c348d4501a49d3417da2d9 Mon Sep 17 00:00:00 2001 From: Yash Jain <23dec512@lnmiit.ac.in> Date: Tue, 14 Jul 2026 10:08:19 +0530 Subject: [PATCH 4/4] Address kaxil review: JSON-mode dump, wrap ValidationError, richer refusal Five fixes from the review of #69812: (1) OpenAIResponseOperator now calls model_dump(mode='json'), so non-JSON field types (plain enum.Enum, datetime, UUID, ...) come back as their JSON representations instead of live Python objects. The old default (mode='python') would break XCom push for models with a plain enum field, since Airflow's serde only unwraps enums that mix in str/int -- and the API call has already been paid for by then. mode='json' is what airflow.serialization.serde.serializers.pydantic already uses, so this alignment holds the operator's XCom-safe claim for any model. (2) parse_response is now wrapped in try/except ValidationError, re-raising as ValueError. responses.parse() raises pydantic.ValidationError internally when the model's JSON output can't be coerced into text_format -- most commonly when the response is truncated mid-JSON on max_output_tokens. Without the wrap, users got a raw pydantic traceback; now they get a single ValueError shape across all parse failures. (3) The output_parsed=None ValueError now includes the API-reported error and incomplete_details in the message, so users don't need a follow-up OpenAIHook.get_response round-trip to diagnose a refusal. (4) Dropped the docstring claim that structured outputs require 'gpt-4o-2024-08-06 or later' -- the default gpt-4o-mini already supports them (has since its initial release), and specific model claims rot as new models ship. Also removed the contradictory 'use hook directly for previous_response_id chaining' line, since response_kwargs already forwards previous_response_id to the SDK. (5) OpenAIHook.parse_response is now generic (TypeVar bound to BaseModel), returning ParsedResponse[T] instead of ParsedResponse[Any]. Callers get output_parsed typed as T | None, matching the SDK's own signature. Also fixes the failing docs --spellcheck-only CI job by rephrasing 'parseable' -> 'return a structured output' throughout, moves the pydantic import in test_openai.py to module top, adds parse_response to the hook-methods list in openai.rst, drops the redundant model= override from the example DAG, and adds a regression test for both the enum-field dump-mode and the ValidationError-to-ValueError conversion. --- providers/openai/docs/operators/openai.rst | 11 ++- .../airflow/providers/openai/hooks/openai.py | 15 ++-- .../providers/openai/operators/openai.py | 50 +++++++---- .../tests/system/openai/example_openai.py | 2 - .../tests/unit/openai/hooks/test_openai.py | 3 +- .../unit/openai/operators/test_openai.py | 83 ++++++++++++++++++- 6 files changed, 133 insertions(+), 31 deletions(-) diff --git a/providers/openai/docs/operators/openai.rst b/providers/openai/docs/operators/openai.rst index d26910e7fd32f..48d8cda9976bc 100644 --- a/providers/openai/docs/operators/openai.rst +++ b/providers/openai/docs/operators/openai.rst @@ -63,9 +63,11 @@ 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()``), which is safe to push to XCom. This requires -a model that supports structured outputs (``gpt-4o-2024-08-06`` or later). If the model refuses or -returns non-parseable content, the operator raises ``ValueError``. +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 @@ -78,7 +80,8 @@ 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 diff --git a/providers/openai/src/airflow/providers/openai/hooks/openai.py b/providers/openai/src/airflow/providers/openai/hooks/openai.py index 2e27ac5f3dd3e..4b731235eb6da 100644 --- a/providers/openai/src/airflow/providers/openai/hooks/openai.py +++ b/providers/openai/src/airflow/providers/openai/hooks/openai.py @@ -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 @@ -67,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.""" @@ -252,10 +257,10 @@ def create_response(self, input: Any, model: str = "gpt-4o-mini", **kwargs: Any) def parse_response( self, input: Any, - text_format: type[BaseModel], + text_format: type[_TextFormatT], model: str = "gpt-4o-mini", **kwargs: Any, - ) -> ParsedResponse[Any]: + ) -> ParsedResponse[_TextFormatT]: """ Create a model response and parse it into a Pydantic model via the Responses API. @@ -266,8 +271,8 @@ def parse_response( :param input: Text, image, or file input(s) to the model. :param text_format: A Pydantic ``BaseModel`` subclass describing the expected - structured output. Requires a model that supports structured outputs - (``gpt-4o-2024-08-06`` and later). + 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) diff --git a/providers/openai/src/airflow/providers/openai/operators/openai.py b/providers/openai/src/airflow/providers/openai/operators/openai.py index 3a86898e3d0d4..531aa5e1622b9 100644 --- a/providers/openai/src/airflow/providers/openai/operators/openai.py +++ b/providers/openai/src/airflow/providers/openai/operators/openai.py @@ -22,6 +22,8 @@ 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 @@ -88,9 +90,9 @@ class OpenAIResponseOperator(BaseOperator): 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()``), which is safe to - push to XCom. Requires a model that supports structured outputs (``gpt-4o-2024-08-06`` and - later). For ``previous_response_id`` chaining or ``background=True`` responses, use + 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. @@ -102,8 +104,8 @@ class OpenAIResponseOperator(BaseOperator): ``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()``; otherwise it calls ``create_response`` and returns - ``output_text``. + ``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: @@ -137,22 +139,38 @@ def hook(self) -> OpenAIHook: def execute(self, context: Context) -> str | dict[str, Any]: if self.text_format is not None: - parsed = self.hook.parse_response( - input=self.input_text, - model=self.model, - text_format=self.text_format, - **self.response_kwargs, - ) + 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 truncated. Surface a clear error so downstream tasks don't get None. + # 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 produce a parseable structured output " - f"(status={parsed.status!r}). Inspect the response via OpenAIHook.get_response " - f"for refusal / error details." + f"Response {parsed.id} did not return a structured output ({', '.join(details)})." ) - return parsed.output_parsed.model_dump() + 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( diff --git a/providers/openai/tests/system/openai/example_openai.py b/providers/openai/tests/system/openai/example_openai.py index 3c7d5ab654c3e..bb1b63ebedd95 100644 --- a/providers/openai/tests/system/openai/example_openai.py +++ b/providers/openai/tests/system/openai/example_openai.py @@ -118,8 +118,6 @@ class Person(BaseModel): OpenAIResponseOperator( task_id="openai_response_structured", conn_id="openai_default", - # Structured outputs require a compatible model (gpt-4o-2024-08-06 or later). - model="gpt-4o-2024-08-06", input_text="Extract the name and age from: 'Alice is 30 years old'.", text_format=Person, ) diff --git a/providers/openai/tests/unit/openai/hooks/test_openai.py b/providers/openai/tests/unit/openai/hooks/test_openai.py index e9353696dda42..59e103b90f811 100644 --- a/providers/openai/tests/unit/openai/hooks/test_openai.py +++ b/providers/openai/tests/unit/openai/hooks/test_openai.py @@ -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 @@ -316,8 +317,6 @@ def test_create_response(mock_openai_hook): def test_parse_response(mock_openai_hook): - from pydantic import BaseModel - class Person(BaseModel): name: str diff --git a/providers/openai/tests/unit/openai/operators/test_openai.py b/providers/openai/tests/unit/openai/operators/test_openai.py index aa84a720ba524..d116d72aae2da 100644 --- a/providers/openai/tests/unit/openai/operators/test_openai.py +++ b/providers/openai/tests/unit/openai/operators/test_openai.py @@ -16,12 +16,13 @@ # under the License. from __future__ import annotations +from enum import Enum from unittest.mock import Mock import pytest from openai.types.batch import Batch from openai.types.responses import ParsedResponse, Response -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from airflow.providers.common.compat.sdk import Context, TaskDeferred from airflow.providers.openai.hooks.openai import OpenAIHook @@ -111,6 +112,23 @@ class _StructuredPerson(BaseModel): name: str +class _Priority(str, Enum): + LOW = "low" + HIGH = "high" + + +class _StructuredTask(BaseModel): + """Pydantic model with an enum field — exercises ``model_dump(mode="json")``. + + A plain ``enum.Enum`` field returned as a live enum instance (``model_dump``'s default + ``mode="python"``) is not XCom-safe: Airflow's serde only unwraps enums that mix in + ``str``/``int``. Using ``mode="json"`` renders the enum via its JSON representation. + """ + + title: str + priority: _Priority + + def test_openai_response_operator_structured_output_returns_dict(): operator = OpenAIResponseOperator( task_id=TASK_ID, @@ -141,6 +159,31 @@ def test_openai_response_operator_structured_output_returns_dict(): mock_hook_instance.create_response.assert_not_called() +def test_openai_response_operator_structured_output_dumps_enum_as_json(): + operator = OpenAIResponseOperator( + task_id=TASK_ID, + conn_id=CONN_ID, + input_text="Classify", + model="test_model", + text_format=_StructuredTask, + ) + mock_hook_instance = Mock(spec=OpenAIHook) + mock_hook_instance.parse_response.return_value = Mock( + spec=ParsedResponse, + id="resp_str_2", + status="completed", + output_parsed=_StructuredTask(title="Deploy", priority=_Priority.HIGH), + ) + operator.hook = mock_hook_instance + + result = operator.execute(Context()) + + # ``mode="json"`` renders the enum as its ``.value`` string, not the ``_Priority`` instance; + # this is what makes the returned dict safe to push to XCom without a serde helper. + assert result == {"title": "Deploy", "priority": "high"} + assert isinstance(result["priority"], str) + + def test_openai_response_operator_structured_output_refusal_raises(): operator = OpenAIResponseOperator( task_id=TASK_ID, @@ -155,10 +198,46 @@ def test_openai_response_operator_structured_output_refusal_raises(): id="resp_refused", status="incomplete", output_parsed=None, + error="model refused", + incomplete_details="max_output_tokens", ) operator.hook = mock_hook_instance - with pytest.raises(ValueError, match="did not produce a parseable structured output"): + with pytest.raises(ValueError, match="did not return a structured output") as excinfo: + operator.execute(Context()) + # API-reported details are surfaced so users don't need a follow-up ``get_response`` call. + message = str(excinfo.value) + assert "status='incomplete'" in message + assert "error='model refused'" in message + assert "incomplete_details='max_output_tokens'" in message + + +def test_openai_response_operator_structured_output_validation_error_raises(): + # ``responses.parse`` raises ``pydantic.ValidationError`` when the model's JSON output + # can't be coerced into ``text_format`` (e.g. truncated mid-JSON on ``max_output_tokens``). + # The operator converts it to a ``ValueError`` so callers see one exception type across + # all parse failures. + operator = OpenAIResponseOperator( + task_id=TASK_ID, + conn_id=CONN_ID, + input_text="Extract: Alice", + model="test_model", + text_format=_StructuredPerson, + ) + # Build a real ValidationError instance the same way the SDK's internal parse would -- + # by feeding a payload that violates the model's schema. No public constructor exists. + try: + _StructuredPerson.model_validate({}) + except ValidationError as real_exc: + validation_error = real_exc + else: + raise AssertionError("expected ValidationError when validating empty dict") + + mock_hook_instance = Mock(spec=OpenAIHook) + mock_hook_instance.parse_response.side_effect = validation_error + operator.hook = mock_hook_instance + + with pytest.raises(ValueError, match="does not match '_StructuredPerson'"): operator.execute(Context())