kaxil commented on code in PR #72151:
URL: https://github.com/apache/airflow/pull/72151#discussion_r4032445669
##########
providers/openai/tests/unit/openai/operators/test_openai.py:
##########
@@ -94,24 +104,103 @@ def test_openai_response_operator_execute():
response_kwargs={"instructions": "Be concise.",
"previous_response_id": "resp_prev"},
)
mock_hook_instance = Mock(spec=OpenAIHook)
- mock_hook_instance.create_response.return_value = Mock(
- spec=Response, output_text="haiku text", id="resp_123",
status="completed"
+ usage = ResponseUsage(
+ input_tokens=5,
+ input_tokens_details=InputTokensDetails(cached_tokens=1,
cache_write_tokens=0),
+ output_tokens=7,
+ output_tokens_details=OutputTokensDetails(reasoning_tokens=2),
+ total_tokens=12,
)
+ mock_response = Mock(
+ spec=Response, output_text="haiku text", id="resp_123",
status="completed", usage=usage
+ )
+ mock_hook_instance.create_response.return_value = mock_response
operator.hook = mock_hook_instance
- result = operator.execute(Context())
+ context = _build_execute_context()
+ result = operator.execute(context)
+ # Backward compat: the return value is still the aggregated output text,
unchanged
+ # by the new XCom pushes below.
assert result == "haiku text"
mock_hook_instance.create_response.assert_called_once_with(
input="Write a haiku.",
model="test_model",
instructions="Be concise.",
previous_response_id="resp_prev",
)
+ context["ti"].xcom_push.assert_any_call(key="response_id",
value="resp_123")
+ context["ti"].xcom_push.assert_any_call(
+ key="usage",
+ value={
+ "input_tokens": 5,
+ "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens":
1},
+ "output_tokens": 7,
+ "output_tokens_details": {"reasoning_tokens": 2},
+ "total_tokens": 12,
+ },
+ )
+
+
+def test_openai_response_operator_execute_without_usage():
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID, conn_id=CONN_ID, input_text="Write a haiku.",
model="test_model"
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_response = Mock(
+ spec=Response, output_text="haiku text", id="resp_123",
status="completed", usage=None
+ )
+ mock_hook_instance.create_response.return_value = mock_response
+ operator.hook = mock_hook_instance
+
+ context = _build_execute_context()
+ result = operator.execute(context)
+
+ assert result == "haiku text"
+ context["ti"].xcom_push.assert_any_call(key="usage", value=None)
+
+
+def test_openai_response_operator_execute_skips_xcom_push_when_disabled():
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ model="test_model",
+ do_xcom_push=False,
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_response = Mock(spec=Response, output_text="haiku text",
id="resp_123", status="completed")
+ mock_hook_instance.create_response.return_value = mock_response
+ operator.hook = mock_hook_instance
+
+ context = _build_execute_context()
+ result = operator.execute(context)
+
+ assert result == "haiku text"
+ context["ti"].xcom_push.assert_not_called()
+
+
+def test_openai_response_operator_templates_input_text_and_response_kwargs():
Review Comment:
This covers one shape -- a flat one-key dict whose value is a plain `{{
params.x }}` string -- but the changelog warns about a different one: "a
defined variable is substituted silently, sending a different prompt or tool
schema to the paid API than before". A tool schema is a list of dicts, and the
templater recurses into lists, dicts and tuples and rewrites every string it
finds, so that is the shape that bites. Nothing here binds it, nor the `{% raw
%}` escape the docs promise, nor a non-string value surviving untouched --
narrow the field to top-level strings tomorrow and every assertion in this test
still passes. To satisfy this, parametrize over the three: `{"instructions":
"{% raw %}{{ not_a_variable }}{% endraw %}"}` -> unchanged literal; `{"tools":
[{"type": "function", "parameters": {"k": "{{ params.input_text }}"}}],
"max_retries": 3}` -> nested string rendered and the int untouched; plus the
existing case.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -134,9 +137,23 @@ class OpenAIResponseOperator(BaseOperator):
:ref:`howto/operator:OpenAIResponseOperator`
For possible options, see:
https://platform.openai.com/docs/api-reference/responses/create
+
+ ``execute`` also pushes two XCom keys: ``response_id`` (the response's ID,
usable as
+ a downstream call's ``previous_response_id``) and ``usage`` (the
response's token
+ usage, or ``None`` when the API omits it). ``usage`` is the nested dict
returned by
+ ``ResponseUsage.model_dump()``: top-level ``input_tokens``,
``output_tokens`` and
+ ``total_tokens`` counts, plus the nested ``input_tokens_details`` and
+ ``output_tokens_details`` dicts. ``input_tokens_details.cached_tokens`` is
part of
+ the ``input_tokens`` total, not additional to it, so pricing a run
correctly means
+ reading the breakdown rather than treating ``input_tokens`` as a single
uniformly
+ priced count -- see OpenAI's `prompt caching guide
+ <https://platform.openai.com/docs/guides/prompt-caching>`_ for how cached
tokens are
+ priced. Beyond that, ``usage`` reports token counts only -- the OpenAI
response
+ carries no cost field, so turning any of these counts into a price means
multiplying
+ by your own per-token rate.
"""
- template_fields: Sequence[str] = ("input_text", "max_output_tokens",
"max_tool_calls")
+ template_fields: Sequence[str] = ("input_text", "response_kwargs",
"max_output_tokens", "max_tool_calls")
Review Comment:
The changelog note scopes the hazard to a literal `{{ ... }}` string, but
templating this dict changes value *types* too, and in one direction it breaks
Dags that contain no Jinja at all. Under `render_template_as_native_obj=True`,
`native_concat` runs `ast.literal_eval` on every string value, so an existing
`response_kwargs={"metadata": {"run": "20240115"}}` starts sending `{"run":
20240115}` after upgrade -- and OpenAI's `metadata` is a string-to-string map.
Under the default environment the reverse: `{"max_output_tokens": "{{
params.tokens }}"}` renders to `"500"` and goes to the API as a string, because
`_build_response_kwargs` only coerces names in `self._supplied_ceilings` (set
from the operator arguments in `__init__`), so a ceiling living only in
`response_kwargs` now renders but skips `_coerce_token_ceiling` entirely. To
satisfy this: extend the note and `:param response_kwargs:` to cover both type
effects, and either run `_coerce_token_ceiling` over any `_TOKEN_CEILING_P
ARAM_NAMES` key found in the rendered dict, or say plainly that a ceiling
placed there is neither coerced nor validated.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -134,9 +137,23 @@ class OpenAIResponseOperator(BaseOperator):
:ref:`howto/operator:OpenAIResponseOperator`
For possible options, see:
https://platform.openai.com/docs/api-reference/responses/create
+
+ ``execute`` also pushes two XCom keys: ``response_id`` (the response's ID,
usable as
Review Comment:
These thirteen lines are near-verbatim with
`docs/operators/openai.rst:47-59` -- an 86-word identical run -- and the two
copies have already drifted before merge: the rst opens with "When
`do_xcom_push` is enabled (the default)" and closes with "Setting
`do_xcom_push=False` skips both pushes", and the docstring says neither. So the
class doc currently promises two XCom pushes with no hint that the operator's
own flag suppresses them, which is the one fact that makes the keys safe. The
convention is a cross-reference, not a copy -- `BashOperator` shares 182
characters with its howto and `PythonOperator` 37, just the `:ref:`. To satisfy
this: cut this block to the mechanical facts (the two keys, that `usage` is
`ResponseUsage.model_dump()` and `None` when omitted, that `do_xcom_push=False`
skips them) and let the `.. seealso::` already at :135-137 carry the pricing
discussion.
##########
providers/openai/tests/system/openai/example_openai.py:
##########
@@ -101,13 +101,26 @@ def task_to_store_input_text_in_xcom():
# [END howto_operator_openai_embedding]
# [START howto_operator_openai_response]
- OpenAIResponseOperator(
+ openai_response = OpenAIResponseOperator(
task_id="openai_response",
conn_id="openai_default",
input_text="Write a haiku about data pipelines.",
response_kwargs={"instructions": "You are a helpful assistant."},
)
+ # Chains onto the previous response via its response_id XCom, continuing
+ # the same conversation without resending prior turns.
+ openai_response_follow_up = OpenAIResponseOperator(
Review Comment:
A Jinja `ti.xcom_pull` string creates no task dependency, so the `>>` at
line 122 is load-bearing -- without it these two race and the follow-up renders
`previous_response_id` to the literal string `"None"`, which reaches the paid
API and 400s. The comment above explains what the chaining does but not that,
so someone copying this snippet can reasonably delete the `>>`. `XComArg` fixes
both halves: `response_kwargs={"previous_response_id": XComArg(openai_response,
key="response_id")}` resolves through the same dict recursion, creates the edge
automatically so line 122 can go, and raises `XComNotFound` instead of
rendering `"None"` when the key is absent. It works on the 2.11 floor too. To
satisfy this: either switch to `XComArg`, or keep the string and extend the
comment to say the Jinja pull creates no dependency so the edge is declared
explicitly below.
##########
providers/openai/docs/operators/openai.rst:
##########
@@ -44,14 +44,30 @@ OpenAIResponseOperator
Use the
:class:`~airflow.providers.openai.operators.openai.OpenAIResponseOperator` to
generate a
model response with the OpenAI Responses API, OpenAI's recommended interface
for text generation and
-tool use. The operator returns the response's aggregated output text.
+tool use. The operator returns the response's aggregated output text. When
``do_xcom_push`` is
Review Comment:
This paragraph announces the chaining, but the `store` bullet further down
the same page (lines 118-122, outside this diff) still says `execute` "only
passes `response.id` to the task log" and that "nothing downstream of this
operator's task can retrieve a stored response's id" -- the exact limitation
this PR removes. A reader who scrolls to the options list is told the feature
does not exist. To satisfy this: rewrite that bullet to say the id is pushed as
`response_id` when `do_xcom_push` is enabled, and add the caveat that now
becomes real -- `store=False` makes the pushed id useless as a
`previous_response_id`. Separately, line 59's "skips both pushes" is
incomplete: `do_xcom_push=False` also suppresses the task's own `return_value`
XCom, so `openai_response.output` stops resolving.
##########
providers/openai/docs/changelog.rst:
##########
@@ -20,6 +20,25 @@
Changelog
---------
+.. Behavior note
Review Comment:
`.. Behavior note` has no `::`, so docutils parses it as an RST comment
rather than a directive -- the rendered page emits `<!-- Behavior note -->` and
lines 25-40 become three unlabelled body paragraphs with no admonition styling.
No docutils warning is emitted for this form, so no Sphinx build will ever
catch it. The placement is right and I am not asking you to move it; the
sibling note at :45 shows the form. To satisfy this: `.. note::` with the body
indented four spaces, and open it with the version that ships this change,
since as written it sits directly above the released 1.8.2 heading and reads as
applying to that. While editing, the hazard list is `{{ }}`-only: `{# ... #}`
is silently deleted from the string and `{% ... %}` raises
`TemplateSyntaxError` -- `{% raw %}` already covers all three.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -272,6 +289,13 @@ def execute(self, context: Context) -> str:
response.status,
)
self.log.info("Generated response %s", response.id)
+ if self.do_xcom_push:
Review Comment:
XComs are cleared at the start of every attempt, so this key only ever holds
the final attempt's counts. Every earlier attempt of a retried task already
made a billed `create_response` call whose tokens OpenAI charges but which this
value no longer shows, so a user summing `usage.total_tokens` to price a Dag
run gets a number that is low exactly on the runs that cost the most -- and the
PR sells these counts as the run's cost signal. The sibling provider already
solved this: `providers/anthropic/.../operators/agent.py:284-293` pushes the
same `usage` key and stamps `try_number` with a comment saying it is precisely
to avoid silently under-reporting spend across retries. To satisfy this: stamp
`try_number` into the dict the same way, or add one sentence to the docstring
saying `usage` reflects the final attempt only.
##########
providers/openai/tests/unit/openai/operators/test_openai.py:
##########
@@ -94,24 +104,103 @@ def test_openai_response_operator_execute():
response_kwargs={"instructions": "Be concise.",
"previous_response_id": "resp_prev"},
)
mock_hook_instance = Mock(spec=OpenAIHook)
- mock_hook_instance.create_response.return_value = Mock(
- spec=Response, output_text="haiku text", id="resp_123",
status="completed"
+ usage = ResponseUsage(
+ input_tokens=5,
+ input_tokens_details=InputTokensDetails(cached_tokens=1,
cache_write_tokens=0),
+ output_tokens=7,
+ output_tokens_details=OutputTokensDetails(reasoning_tokens=2),
+ total_tokens=12,
)
+ mock_response = Mock(
+ spec=Response, output_text="haiku text", id="resp_123",
status="completed", usage=usage
+ )
+ mock_hook_instance.create_response.return_value = mock_response
operator.hook = mock_hook_instance
- result = operator.execute(Context())
+ context = _build_execute_context()
+ result = operator.execute(context)
+ # Backward compat: the return value is still the aggregated output text,
unchanged
+ # by the new XCom pushes below.
assert result == "haiku text"
mock_hook_instance.create_response.assert_called_once_with(
input="Write a haiku.",
model="test_model",
instructions="Be concise.",
previous_response_id="resp_prev",
)
+ context["ti"].xcom_push.assert_any_call(key="response_id",
value="resp_123")
+ context["ti"].xcom_push.assert_any_call(
+ key="usage",
+ value={
+ "input_tokens": 5,
+ "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens":
1},
+ "output_tokens": 7,
+ "output_tokens_details": {"reasoning_tokens": 2},
+ "total_tokens": 12,
+ },
+ )
+
+
+def test_openai_response_operator_execute_without_usage():
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID, conn_id=CONN_ID, input_text="Write a haiku.",
model="test_model"
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_response = Mock(
+ spec=Response, output_text="haiku text", id="resp_123",
status="completed", usage=None
+ )
+ mock_hook_instance.create_response.return_value = mock_response
+ operator.hook = mock_hook_instance
+
+ context = _build_execute_context()
+ result = operator.execute(context)
+
+ assert result == "haiku text"
+ context["ti"].xcom_push.assert_any_call(key="usage", value=None)
+
+
+def test_openai_response_operator_execute_skips_xcom_push_when_disabled():
Review Comment:
This test passes unchanged against pre-PR code -- `execute` pushed nothing
there either -- so on its own it binds none of this PR; it is a forward guard
on the `if self.do_xcom_push:` gate. It cannot report that cleanly though: the
mock on line 172 omits `usage`, so if the gate were deleted `execute` would hit
`response.usage` on a `Mock(spec=Response)` and raise `AttributeError` before
`assert_not_called()` runs, and pytest would show the wrong failure. To satisfy
this: parametrize the enabled and disabled cases over one body asserting
`xcom_push.call_count == expected`, and build the response with the existing
`_build_completed_response()` helper so a removed gate fails on the assertion.
That also closes a gap above -- the enabled test uses `assert_any_call` twice
and never pins the count, so it still passes if a third key were pushed.
--
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]
