kaxil commented on code in PR #70132:
URL: https://github.com/apache/airflow/pull/70132#discussion_r3680669058
##########
providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py:
##########
@@ -271,5 +271,5 @@ def _to_string(output: Any) -> str:
if isinstance(output, BaseModel):
return output.model_dump_json()
if not isinstance(output, str):
- return str(output)
+ return TypeAdapter(type(output)).dump_json(output).decode()
Review Comment:
Same unguarded `TypeAdapter` on the serialize side: `str(output)` could
never raise, `TypeAdapter(type(output))` can. It's reachable through
pydantic-ai output functions, which only warn ("Falling back to unconstrained
schema") for a non-schema-able return type and then run fine --
`Agent(TestModel(), output_type=make_widget)` where `make_widget` returns a
plain class makes this raise `PydanticSchemaGenerationError` where it
previously produced `'Widget(0)'`.
This one fires at line 111, before the session XCom push, so the task dies
before the reviewer ever sees a session. `except Exception: return str(output)`
restores the old degradation, and pairs with the schema-build guard on the
rehydrate side.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -556,4 +550,6 @@ def regenerate_with_feedback(self, *, feedback: str,
message_history: Any) -> tu
output = result.output
if isinstance(output, BaseModel):
output = output.model_dump_json()
- return str(output), result.all_messages()
+ elif not isinstance(output, str):
+ output = TypeAdapter(type(output)).dump_json(output).decode()
+ return output, result.all_messages()
Review Comment:
`AgentOperator` already inherits `_to_string` from `HITLReviewMixin`, and
these lines are now exactly what it does, so `return
self._to_string(result.output), result.all_messages()` collapses it.
Worth noting the same block also lives in
`LLMApprovalMixin.defer_for_approval` (approval.py:108-112), so that's three
copies. The split is how the original bug arose in the first place:
`approval.py` was fixed in #70075 while these two weren't. If you add the guard
from the other comments, one shared helper next to `rehydrate_pydantic_output`
fixes all three at once and gives the round-trip a matching pair.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -493,16 +492,11 @@ def execute(self, context: Context) -> Any:
output,
message_history=result.all_messages(),
)
- if isinstance(self.output_type, type) and
issubclass(self.output_type, BaseModel):
- return rehydrate_pydantic_output(
- self.output_type,
- result_str,
- serialize_output=self._serialize_model_output,
- )
- try:
- return json.loads(result_str)
- except (ValueError, TypeError):
- return result_str
+ return rehydrate_pydantic_output(
Review Comment:
This also changes the default `output_type=str`, and neither the description
nor a test mentions it: an approved string of `'42'` used to come back as `42`,
`'{"text": "hi"}'` as a dict, `'true'` as `True` and `'null'` as `None`; all
four now stay strings. That's the right behaviour -- it makes the HITL path
consistent with the non-HITL path, which returns `result.output` unparsed --
but `str` is the default `output_type` and `enable_hitl_review` has shipped in
every release from 0.1.0 through 0.6.0, so it's worth a line in the description
rather than arriving silently.
Running the six new tests against the pre-fix source gives 3 failed, 3
passed. `test_execute_with_hitl_rehydrates_non_base_model_output` is in the
passing group: it patches `run_hitl_review` to return a hand-written
`'["tag-a", "tag-b"]'`, so `_to_string` never runs and the `json.loads` removal
is never exercised. Nothing currently covers the `execute()` half of this diff.
A case with `output_type=str` and a JSON-parseable approved string would pin it.
--
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]
