kaxil commented on code in PR #69812:
URL: https://github.com/apache/airflow/pull/69812#discussion_r3574140539
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -110,20 +120,39 @@ 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:
+ 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.
Review Comment:
The "response was truncated" case isn't actually covered by the `None`
check: when output is cut off mid-JSON (e.g. `max_output_tokens` hit),
`responses.parse()` raises `pydantic.ValidationError` from its internal parse
step (checked on openai 2.37.0: `parse_text` calls `model_parse_json` with no
try/except). That happens inside `self.hook.parse_response(...)` before
`parsed` is assigned, so the response id is never logged and this message never
fires; users get a raw pydantic traceback instead. Wrapping the
`parse_response` call in `try/except ValidationError` and re-raising the same
shaped `ValueError` would make this comment and the guide's claim hold.
And for the cases that do reach this raise, the response already carries the
detail (`parsed.error`, `parsed.incomplete_details`, refusal text in
`parsed.output`) -- including it in the message would save users the
`get_response` round trip.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -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
Review Comment:
This reads as if the default model doesn't qualify, but `gpt-4o-mini` (the
default here and in `parse_response`) has supported structured outputs since
its initial release -- per OpenAI's structured outputs guide it's the
gpt-4o-mini, gpt-4o-mini-2024-07-18, and gpt-4o-2024-08-06 snapshots and later.
The same claim repeats in the hook docstring, the rst guide, and the example
DAG (which overrides `model="gpt-4o-2024-08-06"` with a comment saying it's
required, when the default already works). I'd drop the specific-model claim
(it will rot as new models ship) or note that the default already supports it.
One more thing in this docstring: this line says to use the hook directly
for `previous_response_id` chaining, but the `response_kwargs` description
below lists `previous_response_id` as a supported passthrough. One of the two
should give.
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -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]:
Review Comment:
The SDK's `parse` is generic (`ParsedResponse[TextFormatT]`), so you can
keep that instead of hardcoding `ParsedResponse[Any]`: `T = TypeVar("T",
bound=BaseModel)`, `text_format: type[T]`, return `ParsedResponse[T]` -- then
`output_parsed` is typed as `T | None` at call sites. Also, can you add
`parse_response` to the hook-methods list in `docs/operators/openai.rst`
alongside the other Responses methods?
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -110,20 +120,39 @@ 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:
+ 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()
Review Comment:
`model_dump()` defaults to python mode, so non-JSON field types come back as
live Python objects. A model with a plain `enum.Enum` field (a common shape for
classification-style structured outputs) then fails at XCom push with
`TypeError: cannot serialize object of type <enum 'Priority'>`, because
Airflow's serde only unwraps enums that mix in `str`/`int` -- and by then the
API call has already been paid for. `model_dump(mode="json")` fixes this and is
what Airflow's own pydantic serializer does
([serde/serializers/pydantic.py](https://github.com/apache/airflow/blob/a1cec525543c048b492f617e174b6c862664e55e/task-sdk/src/airflow/sdk/serde/serializers/pydantic.py#L49)),
making the docstring's XCom-safe claim hold for any model.
##########
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
Review Comment:
Can you move this import to the module top? pydantic is a hard dependency
(the openai SDK requires it), and the operator test file in this PR already
imports it top-level.
--
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]
