From a93c069dd59fd85ef4c40e33400e0f4685f9915b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 12 Mar 2026 13:22:56 -0700 Subject: [PATCH] [Fix] Add max_depth guard to BFL _read_image_bytes recursive function Use the standard depth/max_depth pattern with DEFAULT_MAX_RECURSE_DEPTH to guard the recursive list-unwrapping in _read_image_bytes, matching the existing pattern used by _read_all_bytes in vertex_imagen. Co-Authored-By: Claude Opus 4.6 --- .../image_edit/transformation.py | 22 ++++++++++++------- .../code_coverage_tests/recursive_detector.py | 1 + 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 172b3b1485..c88ee5010d 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import httpx from httpx._types import RequestFiles +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -189,17 +190,22 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): endpoint = self._get_model_endpoint(model) return f"{base_url}{endpoint}" - def _read_image_bytes(self, image: Any) -> bytes: + def _read_image_bytes( + self, + image: Any, + depth: int = 0, + max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, + ) -> bytes: """Read image bytes from various input types.""" - # Unwrap nested lists iteratively to avoid recursion - for _ in range(10): - if isinstance(image, list): - image = image[0] - else: - break - + if depth > max_depth: + raise ValueError( + f"Max recursion depth {max_depth} reached while reading image bytes for Black Forest Labs image edit." + ) if isinstance(image, bytes): return image + elif isinstance(image, list): + # If it's a list, take the first image + return self._read_image_bytes(image[0], depth=depth + 1, max_depth=max_depth) elif isinstance(image, str): if image.startswith(("http://", "https://")): # Download image from URL diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 3710971229..99b5125b19 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -44,6 +44,7 @@ IGNORE_FUNCTIONS = [ "extract_text_from_a2a_message", # max depth set (default 10) to prevent infinite recursion in A2A message parsing. "_convert_to_json_serializable_dict", # max depth set (default 20) and circular reference protection to prevent infinite recursion. "dict", # max depth set. _LiteLLMParamsDictView.dict() calls builtin dict(), not itself. + "_read_image_bytes", # max depth set. ]