diff --git a/providers/openai/docs/operators/openai.rst b/providers/openai/docs/operators/openai.rst index bfb2fd8dee94b..773dcd25e710a 100644 --- a/providers/openai/docs/operators/openai.rst +++ b/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 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..36aa0c9cd3d9e 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 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,6 +141,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 +149,44 @@ def __init__( 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( 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..c462cf5ecb403 100644 --- a/providers/openai/tests/unit/openai/operators/test_openai.py +++ b/providers/openai/tests/unit/openai/operators/test_openai.py @@ -16,11 +16,24 @@ # under the License. from __future__ import annotations +from enum import Enum +from typing import Any 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, + ParsedResponseOutputMessage, + ParsedResponseOutputText, + Response, + ResponseError, + ResponseFunctionToolCall, + ResponseOutputRefusal, +) +from openai.types.responses.response import IncompleteDetails +from pydantic import BaseModel, ValidationError +from pydantic.dataclasses import dataclass as pydantic_dataclass from airflow.providers.common.compat.sdk import Context, TaskDeferred from airflow.providers.openai.hooks.openai import OpenAIHook @@ -104,6 +117,260 @@ def test_openai_response_operator_execute(): ) +class _StructuredPerson(BaseModel): + """Pydantic model used by the structured-output operator tests.""" + + name: str + + +class _Priority(Enum): + LOW = "low" + HIGH = "high" + + +class _StructuredTask(BaseModel): + title: str + priority: _Priority + + +def _build_parsed_response( + output_parsed: BaseModel | None = None, + *, + response_id: str = "resp_structured", + status: str = "completed", + error: ResponseError | None = None, + incomplete_details: IncompleteDetails | None = None, + refusal: str | None = None, + output_items: list[Any] | None = None, +) -> ParsedResponse: + content: list[ParsedResponseOutputText[BaseModel] | ResponseOutputRefusal] + if output_items is not None: + output = output_items + elif output_parsed is not None: + content = [ + ParsedResponseOutputText[BaseModel]( + annotations=[], + text=output_parsed.model_dump_json(), + type="output_text", + parsed=output_parsed, + ) + ] + output = [ + ParsedResponseOutputMessage[BaseModel]( + id=f"msg_{response_id}", + content=content, + role="assistant", + status="completed", + type="message", + ) + ] + elif refusal is not None: + content = [ResponseOutputRefusal(refusal=refusal, type="refusal")] + output = [ + ParsedResponseOutputMessage[BaseModel]( + id=f"msg_{response_id}", + content=content, + role="assistant", + status="completed", + type="message", + ) + ] + else: + output = [] + return ParsedResponse[BaseModel].model_construct( + id=response_id, + status=status, + output=output, + error=error, + incomplete_details=incomplete_details, + ) + + +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 = _build_parsed_response( + _StructuredPerson(name="Alice"), response_id="resp_str_1" + ) + 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 = _build_parsed_response( + _StructuredTask(title="Deploy", priority=_Priority.HIGH), response_id="resp_str_2" + ) + operator.hook = mock_hook_instance + + result = operator.execute(Context()) + + 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 = _build_parsed_response( + response_id="resp_refused", + refusal="I cannot help with that request.", + ) + operator.hook = mock_hook_instance + + with pytest.raises(ValueError, match="did not return a structured output") as excinfo: + operator.execute(Context()) + message = str(excinfo.value) + assert "status='completed'" in message + assert "refusal='I cannot help with that request.'" in message + + +def test_openai_response_operator_structured_output_tools_only_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 = _build_parsed_response( + response_id="resp_tool_call", + output_items=[ + ResponseFunctionToolCall( + arguments='{"name": "Alice"}', + call_id="call_1", + name="extract_person", + type="function_call", + status="completed", + ) + ], + ) + operator.hook = mock_hook_instance + + with pytest.raises(ValueError, match="did not return a structured output") as excinfo: + operator.execute(Context()) + + assert "output_types=['function_call']" in str(excinfo.value) + + +def test_openai_response_operator_structured_output_incomplete_raises_with_valid_model(): + 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 = _build_parsed_response( + _StructuredPerson(name="Alice"), + response_id="resp_incomplete", + status="incomplete", + incomplete_details=IncompleteDetails(reason="max_output_tokens"), + ) + operator.hook = mock_hook_instance + + with pytest.raises(ValueError, match="did not complete") as excinfo: + operator.execute(Context()) + + message = str(excinfo.value) + assert "status='incomplete'" in message + assert "reason='max_output_tokens'" in message + + +def test_openai_response_operator_structured_output_failed_raises_with_error(): + 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 = _build_parsed_response( + response_id="resp_failed", + status="failed", + error=ResponseError(code="server_error", message="The model failed."), + ) + operator.hook = mock_hook_instance + + with pytest.raises(ValueError, match="did not complete") as excinfo: + operator.execute(Context()) + + message = str(excinfo.value) + assert "status='failed'" in message + assert "code='server_error'" in message + assert "message='The model failed.'" 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, + ) + with pytest.raises(ValidationError) as exc_info: + _StructuredPerson.model_validate({}) + + mock_hook_instance = Mock(spec=OpenAIHook) + mock_hook_instance.parse_response.side_effect = exc_info.value + operator.hook = mock_hook_instance + + with pytest.raises(ValueError, match="max_output_tokens"): + operator.execute(Context()) + + +def test_openai_response_operator_rejects_non_base_model_text_format(): + @pydantic_dataclass + class StructuredPerson: + name: str + + with pytest.raises(TypeError, match="Pydantic BaseModel subclass"): + OpenAIResponseOperator( + task_id=TASK_ID, + conn_id=CONN_ID, + input_text="Extract: Alice", + text_format=StructuredPerson, + ) + + @pytest.mark.parametrize("wait_for_completion", [True, False]) def test_openai_trigger_batch_operator_not_deferred(mock_batch, wait_for_completion): operator = OpenAITriggerBatchOperator(