Change serialization inside GenAIGeminiCreateBatchJobOperator (#61253)

Co-authored-by: Anton Nitochkin <nitochkin@google.com>

Author: Crowiant

providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py

@@ -578,7 +578,7 @@ def execute(self, context: Context):
                 job_results = self._prepare_results_for_xcom(job)
                 context["ti"].xcom_push(key="job_results", value=job_results)
 
-        return dict(job)
+        return job.model_dump(mode="json")
 
     def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any]:
         if event["status"] == "error":
@@ -588,7 +588,7 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str,
         if self.retrieve_result and job.error is None:
             job_results = self._prepare_results_for_xcom(job)
             context["ti"].xcom_push(key="job_results", value=job_results)
-        return dict(job)
+        return job.model_dump(mode="json")
 
 
 class GenAIGeminiGetBatchJobOperator(GoogleCloudBaseOperator):

providers/google/tests/unit/google/cloud/operators/test_gen_ai.py

@@ -107,6 +107,20 @@
 TEST_BATCH_JOB_NAME = "test-name"
 TEST_FILE_NAME = "test-file"
 TEST_FILE_PATH = "test/path/to/file"
+TEST_CREATE_BATCH_JOB_RESPONSE = {
+    "src": None,
+    "dest": "test-batch-job-destination",
+    "name": "test-name",
+    "error": None,
+    "model": "test-model",
+    "state": "JOB_STATE_SUCCEEDED",
+    "end_time": "test-end-datetime",
+    "start_time": None,
+    "create_time": "test-create-datetime",
+    "update_time": "test-update-datetime",
+    "display_name": "test-display-name",
+    "completion_stats": None,
+}
 
 
 def assert_warning(msg: str, warnings):
@@ -301,6 +315,59 @@ def test_execute(self, mock_hook):
             create_batch_job_config=None,
         )
 
+    @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+    def test_execute_return_value(self, mock_hook):
+        expected_return = TEST_CREATE_BATCH_JOB_RESPONSE
+
+        mock_job = mock.MagicMock()
+        mock_job.model_dump.return_value = expected_return
+        mock_hook.return_value.create_batch_job.return_value = mock_job
+
+        op = GenAIGeminiCreateBatchJobOperator(
+            task_id=TASK_ID,
+            project_id=GCP_PROJECT,
+            location=GCP_LOCATION,
+            model=TEST_GEMINI_MODEL,
+            gcp_conn_id=GCP_CONN_ID,
+            impersonation_chain=IMPERSONATION_CHAIN,
+            input_source=TEST_BATCH_JOB_INLINED_REQUESTS,
+            gemini_api_key=TEST_GEMINI_API_KEY,
+            deferrable=False,
+            wait_until_complete=False,
+        )
+
+        result = op.execute(context={"ti": mock.MagicMock()})
+
+        assert result == expected_return
+        mock_job.model_dump.assert_called_once_with(mode="json")
+
+    @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+    def test_execute_complete_return_value(self, mock_hook):
+        expected_return = TEST_CREATE_BATCH_JOB_RESPONSE
+
+        event = {"status": "success", "job_name": "test-name"}
+
+        mock_job = mock.MagicMock()
+        mock_job.model_dump.return_value = expected_return
+        mock_hook.return_value.get_batch_job.return_value = mock_job
+
+        op = GenAIGeminiCreateBatchJobOperator(
+            task_id=TASK_ID,
+            project_id=GCP_PROJECT,
+            location=GCP_LOCATION,
+            model=TEST_GEMINI_MODEL,
+            gcp_conn_id=GCP_CONN_ID,
+            impersonation_chain=IMPERSONATION_CHAIN,
+            input_source=TEST_BATCH_JOB_INLINED_REQUESTS,
+            gemini_api_key=TEST_GEMINI_API_KEY,
+        )
+
+        result = op.execute_complete(context={"ti": mock.MagicMock()}, event=event)
+
+        assert result == expected_return
+        mock_hook.return_value.get_batch_job.assert_called_once_with("test-name")
+        mock_job.model_dump.assert_called_once_with(mode="json")
+
 
 class TestGenAIGeminiGetBatchJobOperator:
     @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))