kaxil commented on code in PR #69812:
URL: https://github.com/apache/airflow/pull/69812#discussion_r3682396032
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -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(
Review Comment:
Following up on the details added here: for an actual refusal the API
returns `status='completed'` with `error` and `incomplete_details` both `None`
-- the SDK passes `refusal` content items through unparsed and `output_parsed`
comes back `None` (checked on openai 2.46.0, `lib/_parsing/_responses.py`). So
for the headline failure mode this message degrades to just
`status='completed'` with no reason, while the model's explanation sits unread
in `parsed.output[].content[].refusal`. Scanning `parsed.output` for refusal
content and appending that text (and, failing that, the output item types -- a
tools-only response also lands here) would make the message useful exactly when
it matters.
Relatedly, the refusal test fabricates `status='incomplete'` plus string
`error`/`incomplete_details`, a combination the API can't emit -- a true
refusal-shaped test (`status='completed'`, both fields `None`, refusal item in
`output`) would pin this, and the existing one fits better renamed to
`..._incomplete_raises` with realistic `ResponseError`/`IncompleteDetails`
values.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -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:
Review Comment:
One gap left on this path: the SDK parses output regardless of response
status, so a `status='incomplete'` response whose emitted JSON still validates
(a truncated list field with no min-length constraint, or `content_filter`
cutting after a complete object) sails through here and pushes partial data to
XCom as task success. The window is narrow -- truncation usually breaks the
JSON and lands in the `ValidationError` branch -- but the plain-text path below
does check `response.status != "completed"` while this one doesn't. A one-line
status guard before this check, reusing the same details harvesting, closes it;
worth a regression test that `status='incomplete'` with a valid parsed model
raises.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -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,
Review Comment:
Two smaller ones on this param: the SDK also accepts dataclass-like types
(anything with `__pydantic_config__`), so a `@pydantic.dataclass` passed here
completes the billed API call and then hits `AttributeError` on `model_dump`,
uncaught by the `ValidationError` handler. An `issubclass(text_format,
BaseModel)` check in `__init__` fails at parse time instead (needs `BaseModel`
imported at runtime rather than under `TYPE_CHECKING`). And the class docstring
and this `:param:` entry repeat the same `model_dump(mode="json")`/XCom
rationale -- one of the two can go.
##########
providers/openai/tests/unit/openai/operators/test_openai.py:
##########
@@ -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(
Review Comment:
A few test-fidelity items in this block. `Mock(spec=ParsedResponse)` gives
no real spec protection: pydantic v2 keeps declared fields out of the class
namespace, so `id`/`status`/`error`/`incomplete_details` aren't in the spec and
constructor kwargs bypass it anyway -- an SDK field rename would break
production with all four tests still green. Constructing real `ParsedResponse`
instances (`.construct()` with real output items) fixes that and forces
realistic field shapes at the same time.
Below, the ValidationError setup:
`pydantic_core.ValidationError.from_exception_data` is public (the comment
saying no public constructor exists isn't right), but simpler still is `with
pytest.raises(ValidationError) as exc_info:
_StructuredPerson.model_validate({})` in place of the six-line try/except/else.
And on the comment at line 182 -- a pydantic serde helper does exist and does
this exact `model_dump(mode="json")`; the real advantage of returning a dict is
that readers need no `allowed_deserialization_classes` config.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -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(
Review Comment:
This branch raises before a `ParsedResponse` exists, so the message carries
no response id and no `incomplete_details` -- users can't even run
`hook.get_response(id)` to confirm `reason='max_output_tokens'`, and the guide
now promises "``ValueError`` with the API-reported ``status``, ``error`` and
``incomplete_details``" for this case too. Cheapest fix: name truncation
(`max_output_tokens`) as the likely cause in this message, and split the rst
sentence to describe the two branches separately.
##########
providers/openai/tests/unit/openai/operators/test_openai.py:
##########
@@ -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):
Review Comment:
The `str` mixin defeats this regression test: with the default
`mode="python"`, `model_dump` returns the live `_Priority.HIGH`, but str-enum
equality makes `result == {..., "priority": "high"}` pass and
`isinstance(result["priority"], str)` is also true, so reverting
`model_dump(mode="json")` keeps the test green (verified by running it against
a revert). Dropping the mixin (`class _Priority(Enum)`) makes it fail on
revert, and matches the docstring's own point that str/int-mixin enums were
already safe.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
