kaxil commented on code in PR #67644:
URL: https://github.com/apache/airflow/pull/67644#discussion_r3321298662
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -141,7 +169,25 @@ def execute(self, context: Context) -> Any:
if self.require_approval:
self.defer_for_approval(context, output) # type: ignore[misc]
- if isinstance(output, BaseModel):
+ if self._serialize_model_output and isinstance(output, BaseModel):
+ # ``serialize_output=True`` was set explicitly, or this is an
+ # older Airflow version without ``airflow.sdk.serde.allow_class``.
+ # Either way, dump to dict so XCom carries a plain JSON payload.
output = output.model_dump()
return output
+
+ def execute_complete(self, context: Context, generated_output: str, event:
dict[str, Any]) -> Any:
+ """Resume after human review and restore the Pydantic model for XCom
consumers."""
+ output = super().execute_complete(context, generated_output, event)
+ if isinstance(self.output_type, type) and issubclass(self.output_type,
BaseModel):
Review Comment:
Done in `0840270de4` -- pulled the validate-json/try/except/serialize block
into `airflow.providers.common.ai.utils.output_type.rehydrate_pydantic_output`
and call it from both `LLMOperator.execute_complete` and the `AgentOperator`
HITL branch. Thanks for the catch.
--
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]
