Support Pydantic structured outputs in OpenAIResponseOperator

providers/openai/docs/operators/openai.rst

@@ -58,13 +58,35 @@ 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).
+
+The operator rejects incomplete and failed responses even if the partial output happens to match
+the Pydantic model. The resulting ``ValueError`` includes the response id and available API details,
+such as ``status``, ``error``, ``incomplete_details``, refusal text, or output item types. If the SDK
+cannot parse the model output, it raises ``ValidationError`` before returning a response object; the
+operator converts that to ``ValueError`` naming the requested model and notes that reaching
+``max_output_tokens`` is a likely cause. In that case, a response id and API details are unavailable.
+
+.. 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

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.

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 BaseModel, 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
@@ -31,6 +33,30 @@
     from airflow.providers.common.compat.sdk import Context
 
 
+def _get_structured_response_details(response: Any) -> str:
+    """Return API-reported context for a structured-response failure."""
+    details = [f"status={response.status!r}"]
+    if response.error is not None:
+        details.append(f"error={response.error!r}")
+    if response.incomplete_details is not None:
+        details.append(f"incomplete_details={response.incomplete_details!r}")
+
+    refusals = [
+        content.refusal
+        for output in response.output
+        if output.type == "message"
+        for content in output.content
+        if content.type == "refusal"
+    ]
+    if refusals:
+        details.append(f"refusal={'; '.join(refusals)!r}")
+    else:
+        output_types = [output.type for output in response.output]
+        if output_types:
+            details.append(f"output_types={output_types!r}")
+    return ", ".join(details)
+
+
 class OpenAIEmbeddingOperator(BaseOperator):
     """
     Operator that accepts input text to generate OpenAI embeddings using the specified model.
@@ -84,16 +110,21 @@ 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 an XCom-safe ``dict``. 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``; 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,20 +141,52 @@ 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 {}
+        if text_format is not None and (
+            not isinstance(text_format, type) or not issubclass(text_format, BaseModel)
+        ):
+            raise TypeError("text_format must be a Pydantic BaseModel subclass.")
+        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}. The response may have been truncated because "
+                    f"max_output_tokens was reached: {exc}"
+                ) from exc
+
+            self.log.info("Generated response %s", parsed.id)
+            details = _get_structured_response_details(parsed)
+            if parsed.status != "completed":
+                raise ValueError(f"Response {parsed.id} did not complete ({details}).")
+            if parsed.output_parsed is None:
+                raise ValueError(f"Response {parsed.id} did not return a structured output ({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(

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()
 
 

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")