scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

fix(chatgpt): avoid recursive image extraction

Browse source
This commit is contained in:
CrystalVibe28 2026-05-14 21:43:19 +08:00
parent ffd11fdfa3
commit 49c53492e6
2 changed files with 42 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
litellm/llms/chatgpt/image_generation/transformation.py
View file

@ -509,25 +509,38 @@ class ChatGPTImageGenerationConfig(BaseImageGenerationConfig):
response_payload = payload.get("response")
if isinstance(response_payload, dict):
candidates.extend(self._extract_images_recursive(response_payload))
candidates.extend(self._extract_images_from_nested_value(response_payload))
candidates.extend(self._extract_images_recursive(payload))
candidates.extend(self._extract_images_from_nested_value(payload))
return self._dedupe(candidates), self._dedupe(partial_images)
def _extract_images_recursive(self, value: Any) -> List[str]:
def _extract_images_from_nested_value(self, value: Any) -> List[str]:
images: List[str] = []
if isinstance(value, dict):
value_type = value.get("type")
if value_type in ("image_generation_call", "image_generation"):
images.extend(self._get_image_strings_from_dict(value))
elif isinstance(value.get("b64_json"), str):
images.append(value["b64_json"])
values_to_visit = [value]
visited_container_ids = set()
for child_value in value.values():
images.extend(self._extract_images_recursive(child_value))
elif isinstance(value, list):
for item in value:
images.extend(self._extract_images_recursive(item))
while values_to_visit:
current_value = values_to_visit.pop()
if isinstance(current_value, dict):
container_id = id(current_value)
if container_id in visited_container_ids:
continue
visited_container_ids.add(container_id)
value_type = current_value.get("type")
if value_type in ("image_generation_call", "image_generation"):
images.extend(self._get_image_strings_from_dict(current_value))
elif isinstance(current_value.get("b64_json"), str):
images.append(current_value["b64_json"])
values_to_visit.extend(reversed(list(current_value.values())))
elif isinstance(current_value, list):
container_id = id(current_value)
if container_id in visited_container_ids:
continue
visited_container_ids.add(container_id)
values_to_visit.extend(reversed(current_value))
return self._dedupe(images)
@staticmethod

15
tests/image_gen_tests/test_chatgpt_image_generation.py
View file

@ -190,6 +190,21 @@ def test_chatgpt_image_generation_extracts_b64_from_sse_completed_response(
assert response._hidden_params["model"] == "gpt-image-2"
def test_chatgpt_image_generation_extracts_b64_from_deep_nested_payload(
monkeypatch, tmp_path
):
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
config = ChatGPTImageGenerationConfig()
nested_payload = {"type": "image_generation_call", "result": "b64-image-data"}
for _ in range(1200):
nested_payload = {"nested": [nested_payload]}
images, partial_images = config._extract_images_from_payload(nested_payload)
assert images == ["b64-image-data"]
assert partial_images == []
def test_chatgpt_image_generation_extracts_tool_usage_from_completed_response(
monkeypatch, tmp_path
):