kaxil commented on code in PR #70132:
URL: https://github.com/apache/airflow/pull/70132#discussion_r3625195081
##########
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:
Dropping the `json.loads` fallback here regresses the exact case this PR
targets. `rehydrate_pydantic_output` returns `raw` unchanged when `output_type`
isn't a `BaseModel` subclass (its docstring notes the caller is expected to
apply its own `json.loads`), so with `output_type=list[str]` this returns the
JSON string `'["tag-a","tag-b"]'` instead of the list. The new
`test_execute_with_hitl_rehydrates_non_base_model_output` fails on exactly this
(`assert '["tag-a", "tag-b"]' == ['tag-a', 'tag-b']`). Keeping the previous
`json.loads`/`except` fallback after the rehydrate call fixes it while
preserving the `BaseModel` handling.
--
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]
