From 08c47aa1a4237daed90d86d38ccb4590c25c8ab8 Mon Sep 17 00:00:00 2001
From: Kaituo Huang <kaituo@google.com>
Date: Fri, 30 Jan 2026 10:50:05 -0800
Subject: [PATCH] fix: Add metadata in batch inlined response

PiperOrigin-RevId: 863306844
---
 google/genai/batches.py                       |  3 ++
 .../test_create_with_inlined_requests.py      | 46 +++++++++++++------
 google/genai/tests/batches/test_get.py        |  2 +-
 google/genai/types.py                         |  7 +++
 4 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/google/genai/batches.py b/google/genai/batches.py
index 5223a4df0..21dce5ce9 100644
--- a/google/genai/batches.py
+++ b/google/genai/batches.py
@@ -1208,6 +1208,9 @@ def _InlinedResponse_from_mldev(
         ),
     )
 
+  if getv(from_object, ['metadata']) is not None:
+    setv(to_object, ['metadata'], getv(from_object, ['metadata']))
+
   if getv(from_object, ['error']) is not None:
     setv(to_object, ['error'], getv(from_object, ['error']))
 
diff --git a/google/genai/tests/batches/test_create_with_inlined_requests.py b/google/genai/tests/batches/test_create_with_inlined_requests.py
index ad312dd5c..f528162cd 100644
--- a/google/genai/tests/batches/test_create_with_inlined_requests.py
+++ b/google/genai/tests/batches/test_create_with_inlined_requests.py
@@ -42,20 +42,36 @@
     },
 ]
 
-_INLINED_REQUEST = {
-    'contents': [{
-        'parts': [{
-            'text': 'Hello!',
+_INLINED_REQUESTS = [
+    {
+        'contents': [{
+            'parts': [{
+                'text': 'what is the number after 1? return just the number.',
+            }],
+            'role': 'user',
         }],
-        'role': 'user',
-    }],
-    'metadata': {
-        'key': 'request-1',
+        'metadata': {
+            'key': 'request-1',
+        },
+        'config': {
+            'safety_settings': _SAFETY_SETTINGS,
+        },
     },
-    'config': {
-        'safety_settings': _SAFETY_SETTINGS,
+    {
+        'contents': [{
+            'parts': [{
+                'text': 'what is the number after 2? return just the number.',
+            }],
+            'role': 'user',
+        }],
+        'metadata': {
+            'key': 'request-2',
+        },
+        'config': {
+            'safety_settings': _SAFETY_SETTINGS,
+        },
     },
-}
+]
 _INLINED_TEXT_REQUEST_UNION = {
     'contents': [{
         'parts': [{
@@ -154,7 +170,7 @@
         name='test_union_with_inlined_request',
         parameters=types._CreateBatchJobParameters(
             model=_MLDEV_GEMINI_MODEL,
-            src=[_INLINED_REQUEST],
+            src=_INLINED_REQUESTS,
             config={
                 'display_name': _DISPLAY_NAME,
             },
@@ -166,7 +182,7 @@
         name='test_with_inlined_request',
         parameters=types._CreateBatchJobParameters(
             model=_MLDEV_GEMINI_MODEL,
-            src={'inlined_requests': [_INLINED_REQUEST]},
+            src={'inlined_requests': _INLINED_REQUESTS},
             config={
                 'display_name': _DISPLAY_NAME,
             },
@@ -177,7 +193,7 @@
         name='test_with_inlined_request_config',
         parameters=types._CreateBatchJobParameters(
             model=_MLDEV_GEMINI_MODEL,
-            src={'inlined_requests': [_INLINED_TEXT_REQUEST]},
+            src={'inlined_requests': _INLINED_REQUESTS},
             config={
                 'display_name': _DISPLAY_NAME,
             },
@@ -247,7 +263,7 @@ async def test_async_create(client):
   with pytest_helper.exception_if_vertex(client, ValueError):
     batch_job = await client.aio.batches.create(
         model=_GEMINI_MODEL,
-        src=[_INLINED_REQUEST],
+        src=_INLINED_REQUESTS,
     )
     assert batch_job.name.startswith('batches/')
     assert (
diff --git a/google/genai/tests/batches/test_get.py b/google/genai/tests/batches/test_get.py
index b26740e16..5f010ab18 100644
--- a/google/genai/tests/batches/test_get.py
+++ b/google/genai/tests/batches/test_get.py
@@ -29,7 +29,7 @@
     f'batchPredictionJobs/{_BATCH_JOB_NAME}'
 )
 # MLDev batch operation name.
-_MLDEV_BATCH_OPERATION_NAME = 'batches/0yew7plxupyybd7appsrq5vw7w0lp3l79lab'
+_MLDEV_BATCH_OPERATION_NAME = 'batches/z2p8ksus4lyxt25rntl3fpd67p2niw4hfij5'
 _INVALID_BATCH_JOB_NAME = 'invalid_name'
 
 
diff --git a/google/genai/types.py b/google/genai/types.py
index a6ed2f4d6..fca0fe101 100644
--- a/google/genai/types.py
+++ b/google/genai/types.py
@@ -14382,6 +14382,10 @@ class InlinedResponse(_common.BaseModel):
       description="""The response to the request.
       """,
   )
+  metadata: Optional[dict[str, str]] = Field(
+      default=None,
+      description="""The metadata to be associated with the request.""",
+  )
   error: Optional[JobError] = Field(
       default=None,
       description="""The error encountered while processing the request.
@@ -14396,6 +14400,9 @@ class InlinedResponseDict(TypedDict, total=False):
   """The response to the request.
       """
 
+  metadata: Optional[dict[str, str]]
+  """The metadata to be associated with the request."""
+
   error: Optional[JobErrorDict]
   """The error encountered while processing the request.
       """