Validate model-returned product and image ids

actions/generate_storyboard.py

@@ -285,19 +285,25 @@ def execute(
 
   # Group images by product ID to build context
   products: dict[str, Product] = {}
+  valid_product_image_combinations = {}
 
   for img_obj in images:
     product_id = str(img_obj.get(Dimension.PRODUCT_ID.value, "1"))
-    product_description = str(img_obj.get("product_description", None))
+    desc = img_obj.get("product_description")
+    product_description = str(desc) if desc is not None else ""
     if product_id not in products:
       products[product_id] = Product(product_id, product_description)
+    if product_id not in valid_product_image_combinations:
+      valid_product_image_combinations[product_id] = set()
 
+    image_id = str(img_obj.get(Dimension.IMAGE_ID.value, "1"))
     product = products[product_id]
     image = Image(
-        id=str(img_obj.get(Dimension.IMAGE_ID.value, "1")),
+        id=image_id,
         uri=gcs.get_uri(img_obj[Key.FILE.value]),
     )
     product.images.append(image)
+    valid_product_image_combinations[product_id].add(image_id)
 
   # Add products and images
   prompt_parts.append(genai.types.Part.from_text(text="### Products & Images:\n\n"))
@@ -373,41 +379,67 @@ def execute(
       ],
   )
 
-  start_time = time.time()
-  response = client.models.generate_content(
-      model=gemini_model,
-      contents=prompt_parts,
-      config=config,
-  )
-  end_time = time.time()
-  logger.info(
-      "Gemini API request completed in %.2f seconds.", end_time - start_time
-  )
-
-  segments = []
-  if (
-      response.candidates
-      and response.candidates[0].content
-      and response.candidates[0].content.parts
-  ):
-    parts = list(response.candidates[0].content.parts)
-    for part in parts:
-      if part.text:
-        segments.append(part.text)
-  storyboard_text = "".join(segments)
-  if not storyboard_text:
-    raise ValueError(
-        "The model returned no storyboard text. This can happen when the"
-        " model reaches the output token limit before emitting a response."
+  result = None
+  for _ in range(4):  # 1 initial call + 3 retries
+    start_time = time.time()
+    response = client.models.generate_content(
+        model=gemini_model,
+        contents=prompt_parts,
+        config=config,
+    )
+    end_time = time.time()
+    logger.info(
+        "Gemini API request completed in %.2f seconds.", end_time - start_time
     )
-  json_result = json.loads(storyboard_text)
 
-  if not json_result.get("storyboard"):
-    raise ValueError("No storyboard found in the response.")
+    segments = []
+    if (
+        response.candidates
+        and response.candidates[0].content
+        and response.candidates[0].content.parts
+    ):
+      parts = list(response.candidates[0].content.parts)
+      for part in parts:
+        if part.text:
+          segments.append(part.text)
+    storyboard_text = "".join(segments)
+    if not storyboard_text:
+      raise ValueError(
+          "The model returned no storyboard text. This can happen when the"
+          " model reaches the output token limit before emitting a response."
+      )
+    candidate_result = json.loads(storyboard_text)
+
+    if not candidate_result.get("storyboard"):
+      raise ValueError("No storyboard found in the response.")
+
+    is_valid = True
+    for scene in candidate_result["storyboard"]:
+      product_id = scene.get(Dimension.PRODUCT_ID.value)
+      image_id = scene.get(Dimension.IMAGE_ID.value)
+      if (
+          product_id not in valid_product_image_combinations
+          or image_id
+          not in valid_product_image_combinations.get(product_id, set())
+      ):
+        is_valid = False
+        logger.warning(
+            "Invalid scene from Gemini: product_id=%s, image_id=%s",
+            product_id,
+            image_id,
+        )
+        break  # This scene is invalid, so the whole result is.
+    if is_valid:
+      result = candidate_result
+      break
+  if not result:
+    raise RuntimeError(
+        "Gemini repeatedly scripted invalid image/product combinations"
+    )
 
   valid_scenes = []
 
-  for scene in json_result["storyboard"]:
+  for scene in result["storyboard"]:
     image_id = scene.get(Dimension.IMAGE_ID.value)
     product_id = scene.get(Dimension.PRODUCT_ID.value)
     video_prompt = scene.get("video_prompt")

actions/test/test_generate_storyboard.py

@@ -27,6 +27,198 @@
 
 class TestGenerateStoryboard(unittest.TestCase):
 
+  def _make_mock_response(self, storyboard_dict):
+    mock_response = MagicMock()
+    mock_candidate = MagicMock()
+    mock_part = MagicMock()
+    mock_part.text = json.dumps(storyboard_dict)
+    mock_candidate.content.parts = [mock_part]
+    mock_response.candidates = [mock_candidate]
+    return mock_response
+
+  @patch("actions.generate_storyboard.genai.Client")
+  def test_execute_hallucinated_product_id_triggers_retry(
+      self, mock_genai_client_class
+  ):
+    mock_client = MagicMock()
+    mock_genai_client_class.return_value = mock_client
+    invalid_resp = self._make_mock_response({
+        "storyboard": [{
+            "image_id": "i1",
+            "product_id": "p_invalid",
+            "scene_name": "Scene 1",
+            "video_prompt": "Prompt 1",
+        }]
+    })
+    valid_resp = self._make_mock_response({
+        "storyboard": [{
+            "image_id": "i1",
+            "product_id": "p1",
+            "scene_name": "Scene 1",
+            "video_prompt": "Prompt 1",
+        }]
+    })
+    mock_client.models.generate_content.side_effect = [invalid_resp, valid_resp]
+
+    result = generate_storyboard.execute(
+        self.mock_gcs,
+        self.mock_params,
+        self.mock_images,
+        self.mock_user_prompt,
+        self.gemini_model,
+        self.gemini_model_location,
+    )
+
+    self.assertEqual(mock_client.models.generate_content.call_count, 2)
+    self.assertEqual(
+        result,
+        {"storyboard": [{Key.FILE.value: "gs://bucket/storyboard.json"}]},
+    )
+
+  @patch("actions.generate_storyboard.genai.Client")
+  def test_execute_hallucinated_image_id_triggers_retry(
+      self, mock_genai_client_class
+  ):
+    mock_client = MagicMock()
+    mock_genai_client_class.return_value = mock_client
+    invalid_resp = self._make_mock_response({
+        "storyboard": [{
+            "image_id": "i_invalid",
+            "product_id": "p1",
+            "scene_name": "Scene 1",
+            "video_prompt": "Prompt 1",
+        }]
+    })
+    valid_resp = self._make_mock_response({
+        "storyboard": [{
+            "image_id": "i1",
+            "product_id": "p1",
+            "scene_name": "Scene 1",
+            "video_prompt": "Prompt 1",
+        }]
+    })
+    mock_client.models.generate_content.side_effect = [invalid_resp, valid_resp]
+
+    result = generate_storyboard.execute(
+        self.mock_gcs,
+        self.mock_params,
+        self.mock_images,
+        self.mock_user_prompt,
+        self.gemini_model,
+        self.gemini_model_location,
+    )
+
+    self.assertEqual(mock_client.models.generate_content.call_count, 2)
+    self.assertEqual(
+        result,
+        {"storyboard": [{Key.FILE.value: "gs://bucket/storyboard.json"}]},
+    )
+
+  @patch("actions.generate_storyboard.genai.Client")
+  def test_execute_exhausting_retries_raises_runtime_error(
+      self, mock_genai_client_class
+  ):
+    mock_client = MagicMock()
+    mock_genai_client_class.return_value = mock_client
+    invalid_resp = self._make_mock_response({
+        "storyboard": [{
+            "image_id": "i_invalid",
+            "product_id": "p1",
+            "scene_name": "Scene 1",
+            "video_prompt": "Prompt 1",
+        }]
+    })
+    mock_client.models.generate_content.return_value = invalid_resp
+
+    with self.assertRaises(RuntimeError) as cm:
+      generate_storyboard.execute(
+          self.mock_gcs,
+          self.mock_params,
+          self.mock_images,
+          self.mock_user_prompt,
+          self.gemini_model,
+          self.gemini_model_location,
+      )
+
+    self.assertEqual(mock_client.models.generate_content.call_count, 4)
+    self.assertIn("Gemini repeatedly scripted invalid image/product combinations", str(cm.exception))
+
+  @patch("actions.generate_storyboard.genai.Client")
+  def test_execute_valid_ids_pass_through_unchanged(
+      self, mock_genai_client_class
+  ):
+    mock_client = MagicMock()
+    mock_genai_client_class.return_value = mock_client
+    valid_resp = self._make_mock_response({
+        "storyboard": [{
+            "image_id": "i1",
+            "product_id": "p1",
+            "scene_name": "Scene 1",
+            "video_prompt": "Prompt 1",
+        }]
+    })
+    mock_client.models.generate_content.return_value = valid_resp
+
+    generate_storyboard.execute(
+        self.mock_gcs,
+        self.mock_params,
+        self.mock_images,
+        self.mock_user_prompt,
+        self.gemini_model,
+        self.gemini_model_location,
+    )
+
+    self.mock_gcs.store.assert_called_once()
+    stored_json = json.loads(self.mock_gcs.store.call_args[0][0])
+    self.assertEqual(stored_json["storyboard"][0]["product_id"], "p1")
+    self.assertEqual(stored_json["storyboard"][0]["image_id"], "i1")
+
+  @patch("actions.generate_storyboard.genai.Client")
+  def test_execute_missing_product_description_yields_empty_string_not_none(
+      self, mock_genai_client_class
+  ):
+    mock_client = MagicMock()
+    mock_genai_client_class.return_value = mock_client
+    mock_client.models.generate_content.return_value = self._make_mock_response({
+        "storyboard": [{
+            "image_id": "i1",
+            "product_id": "p1",
+            "scene_name": "Scene 1",
+            "video_prompt": "Prompt 1",
+        }]
+    })
+
+    # Test with both omitted and explicitly None product_description
+    images = [
+        {
+            Dimension.PRODUCT_ID.value: "p1",
+            Dimension.IMAGE_ID.value: "i1",
+            Key.FILE.value: "path/to/img1.jpg",
+            # product_description omitted
+        },
+        {
+            Dimension.PRODUCT_ID.value: "p2",
+            Dimension.IMAGE_ID.value: "i2",
+            Key.FILE.value: "path/to/img2.jpg",
+            "product_description": None,
+        },
+    ]
+
+    generate_storyboard.execute(
+        self.mock_gcs,
+        self.mock_params,
+        images,
+        self.mock_user_prompt,
+        self.gemini_model,
+        self.gemini_model_location,
+    )
+
+    _, kwargs = mock_client.models.generate_content.call_args
+    prompt_parts = kwargs["contents"]
+    all_text = "".join(p.text for p in prompt_parts if getattr(p, "text", None))
+    self.assertNotIn("'None'", all_text)
+    self.assertNotIn("Product description:", all_text)
+
   def setUp(self):
     self.mock_gcs = MagicMock()
     self.mock_params = {Key.GCP_PROJECT.value: "test-project"}

actions_lib/gemini.py

@@ -18,7 +18,7 @@
 
 import json
 import logging
-import mimetypes
+from types import MappingProxyType
 
 from google import genai
 from google.genai import types
@@ -29,6 +29,21 @@
 
 logger = logging.getLogger(__name__)
 
+ALLOWED_MIME_TYPES = MappingProxyType({
+    "png": "image/png",
+    "jpeg": "image/jpeg",
+    "jpg": "image/jpeg",
+    "webp": "image/webp",
+    "avif": "image/avif",
+    "heic": "image/heic",
+    "heif": "image/heif",
+    "mp4": "video/mp4",
+    "mov": "video/quicktime",
+    "avi": "video/x-msvideo",
+    "webm": "video/webm",
+    "mkv": "video/x-matroska",
+})
+
 
 def get_mime_type(uri: str) -> str:
     """Checks the file and returns its MIME type if supported.
@@ -40,21 +55,8 @@ def get_mime_type(uri: str) -> str:
       the mimetype string.
     """
     file_extension = uri.split(".")[-1].lower()
-    allowed_types = [
-        "png",
-        "jpeg",
-        "jpg",
-        "webp",
-        "heic",
-        "heif",
-        "mp4",
-        "mov",
-        "avi",
-        "webm",
-        "mkv",
-    ]
-    if file_extension in allowed_types:
-        return mimetypes.types_map["." + file_extension]
+    if file_extension in ALLOWED_MIME_TYPES:
+        return ALLOWED_MIME_TYPES[file_extension]
 
     raise ValueError(f"File extension '{file_extension}' ({uri}) is not allowed.")