diff --git a/providers/openai/docs/operators/openai.rst b/providers/openai/docs/operators/openai.rst index bfb2fd8dee94b..48d8cda9976bc 100644 --- a/providers/openai/docs/operators/openai.rst +++ b/providers/openai/docs/operators/openai.rst @@ -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 diff --git a/providers/openai/src/airflow/providers/openai/hooks/openai.py b/providers/openai/src/airflow/providers/openai/hooks/openai.py index 8315a907ffa6c..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 @@ -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 @@ -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.""" @@ -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. diff --git a/providers/openai/src/airflow/providers/openai/operators/openai.py b/providers/openai/src/airflow/providers/openai/operators/openai.py index 0f040a03b5505..531aa5e1622b9 100644 --- a/providers/openai/src/airflow/providers/openai/operators/openai.py +++ b/providers/openai/src/airflow/providers/openai/operators/openai.py @@ -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 @@ -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: @@ -110,6 +122,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 +130,47 @@ 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: + 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( diff --git a/providers/openai/tests/system/openai/example_openai.py b/providers/openai/tests/system/openai/example_openai.py index e03bc6397456f..bb1b63ebedd95 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,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() diff --git a/providers/openai/tests/unit/openai/hooks/test_openai.py b/providers/openai/tests/unit/openai/hooks/test_openai.py index 5e5882e354869..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 @@ -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") diff --git a/providers/openai/tests/unit/openai/operators/test_openai.py b/providers/openai/tests/unit/openai/operators/test_openai.py index 954306d42a5eb..d116d72aae2da 100644 --- a/providers/openai/tests/unit/openai/operators/test_openai.py +++ b/providers/openai/tests/unit/openai/operators/test_openai.py @@ -16,11 +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 Response +from openai.types.responses import ParsedResponse, Response +from pydantic import BaseModel, ValidationError from airflow.providers.common.compat.sdk import Context, TaskDeferred from airflow.providers.openai.hooks.openai import OpenAIHook @@ -104,6 +106,141 @@ def test_openai_response_operator_execute(): ) +class _StructuredPerson(BaseModel): + """Pydantic model used by the structured-output operator tests.""" + + 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, + 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_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, + 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, + error="model refused", + incomplete_details="max_output_tokens", + ) + operator.hook = mock_hook_instance + + 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()) + + @pytest.mark.parametrize("wait_for_completion", [True, False]) def test_openai_trigger_batch_operator_not_deferred(mock_batch, wait_for_completion): operator = OpenAITriggerBatchOperator(