chore: Made Operations.response be of GeneratedVideosResponse type

PiperOrigin-RevId: 744340476

Author: IvanLH

google/genai/_replay_api_client.py

@@ -35,6 +35,7 @@
 from ._api_client import HttpResponse
 from ._common import BaseModel
 from .types import HttpOptionsOrDict
+from .types import GenerateVideosOperation
 
 
 def _redact_version_numbers(version_string: str) -> str:
@@ -397,7 +398,12 @@ def _verify_response(self, response_model: BaseModel):
     if isinstance(response_model, list):
       response_model = response_model[0]
     print('response_model: ', response_model.model_dump(exclude_none=True))
-    actual = response_model.model_dump(exclude_none=True, mode='json')
+    if isinstance(response_model, GenerateVideosOperation):
+      actual = response_model.model_dump(
+          exclude={'result'}, exclude_none=True, mode='json'
+      )
+    else:
+      actual = response_model.model_dump(exclude_none=True, mode='json')
     expected = interaction.response.sdk_response_segments[
         self._sdk_response_index
     ]

google/genai/models.py

@@ -3204,8 +3204,16 @@ def _GenerateVideosOperation_from_mldev(
   if getv(from_object, ['error']) is not None:
     setv(to_object, ['error'], getv(from_object, ['error']))
 
-  if getv(from_object, ['response']) is not None:
-    setv(to_object, ['response'], getv(from_object, ['response']))
+  if getv(from_object, ['response', 'generateVideoResponse']) is not None:
+    setv(
+        to_object,
+        ['response'],
+        _GenerateVideosResponse_from_mldev(
+            api_client,
+            getv(from_object, ['response', 'generateVideoResponse']),
+            to_object,
+        ),
+    )
 
   if getv(from_object, ['response', 'generateVideoResponse']) is not None:
     setv(
@@ -3847,7 +3855,13 @@ def _GenerateVideosOperation_from_vertex(
     setv(to_object, ['error'], getv(from_object, ['error']))
 
   if getv(from_object, ['response']) is not None:
-    setv(to_object, ['response'], getv(from_object, ['response']))
+    setv(
+        to_object,
+        ['response'],
+        _GenerateVideosResponse_from_vertex(
+            api_client, getv(from_object, ['response']), to_object
+        ),
+    )
 
   if getv(from_object, ['response']) is not None:
     setv(

google/genai/operations.py

@@ -178,8 +178,16 @@ def _GenerateVideosOperation_from_mldev(
   if getv(from_object, ['error']) is not None:
     setv(to_object, ['error'], getv(from_object, ['error']))
 
-  if getv(from_object, ['response']) is not None:
-    setv(to_object, ['response'], getv(from_object, ['response']))
+  if getv(from_object, ['response', 'generateVideoResponse']) is not None:
+    setv(
+        to_object,
+        ['response'],
+        _GenerateVideosResponse_from_mldev(
+            api_client,
+            getv(from_object, ['response', 'generateVideoResponse']),
+            to_object,
+        ),
+    )
 
   if getv(from_object, ['response', 'generateVideoResponse']) is not None:
     setv(
@@ -285,7 +293,13 @@ def _GenerateVideosOperation_from_vertex(
     setv(to_object, ['error'], getv(from_object, ['error']))
 
   if getv(from_object, ['response']) is not None:
-    setv(to_object, ['response'], getv(from_object, ['response']))
+    setv(
+        to_object,
+        ['response'],
+        _GenerateVideosResponse_from_vertex(
+            api_client, getv(from_object, ['response']), to_object
+        ),
+    )
 
   if getv(from_object, ['response']) is not None:
     setv(

google/genai/tests/models/test_generate_videos.py

@@ -154,9 +154,9 @@ def test_text_and_image_to_video_poll(client):
 def test_create_operation_to_poll(client):
   if client.vertexai:
     # Fill in project and location for record mode
-    operation_name = "projects/<project>/locations/<location>/publishers/google/models/veo-2.0-generate-001/operations/ce4324d0-1a9f-4fd0-98de-0c54e2ca5798"
+    operation_name = "projects/<project>/locations/<location>/publishers/google/models/veo-2.0-generate-001/operations/4a040b94-2343-4748-9322-c284371bdbd5"
   else:
-    operation_name = "models/veo-2.0-generate-001/operations/s421ckxsogje"
+    operation_name = "models/veo-2.0-generate-001/operations/xtc75gxjir7d"
 
   operation = types.GenerateVideosOperation(
       name=operation_name,
@@ -185,6 +185,8 @@ async def test_text_to_video_poll_async(client):
     # Skip the sleep when in replay mode.
     if not isinstance(client._api_client, _replay_api_client.ReplayApiClient):
       time.sleep(20)
-    operation = await client.aio.operations.get(operation=operation)
+    operation = await client.aio.operations.get(
+        operation=operation
+    )
 
   assert operation.result.generated_videos[0].video.uri

google/genai/tunings.py

@@ -545,9 +545,6 @@ def _Operation_from_mldev(
   if getv(from_object, ['error']) is not None:
     setv(to_object, ['error'], getv(from_object, ['error']))
 
-  if getv(from_object, ['response']) is not None:
-    setv(to_object, ['response'], getv(from_object, ['response']))
-
   return to_object
 
 

google/genai/types.py

@@ -5469,9 +5469,8 @@ class GenerateVideosOperation(_common.BaseModel):
       default=None,
       description="""The error result of the operation in case of failure or cancellation.""",
   )
-  response: Optional[dict[str, Any]] = Field(
-      default=None,
-      description="""The normal response of the operation in case of success.""",
+  response: Optional[GenerateVideosResponse] = Field(
+      default=None, description="""The generated videos."""
   )
   result: Optional[GenerateVideosResponse] = Field(
       default=None, description="""The generated videos."""
@@ -5493,8 +5492,8 @@ class GenerateVideosOperationDict(TypedDict, total=False):
   error: Optional[dict[str, Any]]
   """The error result of the operation in case of failure or cancellation."""
 
-  response: Optional[dict[str, Any]]
-  """The normal response of the operation in case of success."""
+  response: Optional[GenerateVideosResponseDict]
+  """The generated videos."""
 
   result: Optional[GenerateVideosResponseDict]
   """The generated videos."""
@@ -6691,10 +6690,6 @@ class Operation(_common.BaseModel):
       default=None,
       description="""The error result of the operation in case of failure or cancellation.""",
   )
-  response: Optional[dict[str, Any]] = Field(
-      default=None,
-      description="""The normal response of the operation in case of success.""",
-  )
 
 
 class OperationDict(TypedDict, total=False):
@@ -6712,9 +6707,6 @@ class OperationDict(TypedDict, total=False):
   error: Optional[dict[str, Any]]
   """The error result of the operation in case of failure or cancellation."""
 
-  response: Optional[dict[str, Any]]
-  """The normal response of the operation in case of success."""
-
 
 OperationOrDict = Union[Operation, OperationDict]