gopidesupavan commented on code in PR #67644:
URL: https://github.com/apache/airflow/pull/67644#discussion_r3321080481
##########
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:
can we move this to small helper may be to output_types file? its same thing
in the agentoperator?
--
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]
