fix(chatgpt): avoid recursive image usage parsing

litellm/llms/chatgpt/image_generation/response_parsing.py

@@ -86,12 +86,29 @@ def get_image_generation_usage(response_payload: Any) -> Optional[dict]:
     if not isinstance(response_payload, dict):
         return None
 
-    response = response_payload.get("response")
-    if isinstance(response, dict):
-        image_gen_usage = get_image_generation_usage(response)
+    payloads_to_check: List[dict] = []
+    current_payload = response_payload
+    visited_container_ids = set()
+    while isinstance(current_payload, dict):
+        container_id = id(current_payload)
+        if container_id in visited_container_ids:
+            break
+        visited_container_ids.add(container_id)
+        payloads_to_check.append(current_payload)
+
+        next_payload = current_payload.get("response")
+        if not isinstance(next_payload, dict):
+            break
+        current_payload = next_payload
+
+    for payload in reversed(payloads_to_check):
+        image_gen_usage = _get_image_generation_usage_from_payload(payload)
         if image_gen_usage is not None:
             return image_gen_usage
+    return None
 
+
+def _get_image_generation_usage_from_payload(response_payload: dict) -> Optional[dict]:
     tool_usage = response_payload.get("tool_usage")
     if not isinstance(tool_usage, dict):
         return None

tests/test_litellm/llms/chatgpt/image_generation/test_usage.py

@@ -28,6 +28,54 @@ def test_chatgpt_image_generation_usage_helpers_ignore_invalid_payloads():
     )
 
 
+def test_chatgpt_image_generation_usage_helper_prefers_nested_response_usage():
+    config = ChatGPTImageGenerationConfig()
+
+    usage = config._get_image_generation_usage(
+        {
+            "tool_usage": {
+                "image_gen": {
+                    "input_tokens": 1,
+                    "output_tokens": 2,
+                }
+            },
+            "response": {
+                "tool_usage": {
+                    "image_gen": {
+                        "input_tokens": 10,
+                        "output_tokens": 20,
+                    }
+                }
+            },
+        }
+    )
+
+    assert usage is not None
+    assert usage["input_tokens"] == 10
+    assert usage["output_tokens"] == 20
+    assert usage["total_tokens"] == 30
+
+
+def test_chatgpt_image_generation_usage_helper_handles_cyclic_response_payload():
+    config = ChatGPTImageGenerationConfig()
+    payload = {
+        "tool_usage": {
+            "image_gen": {
+                "input_tokens": 3,
+                "output_tokens": 4,
+            }
+        }
+    }
+    payload["response"] = payload
+
+    usage = config._get_image_generation_usage(payload)
+
+    assert usage is not None
+    assert usage["input_tokens"] == 3
+    assert usage["output_tokens"] == 4
+    assert usage["total_tokens"] == 7
+
+
 def test_chatgpt_image_generation_extracts_tool_usage_from_completed_response():
     config = ChatGPTImageGenerationConfig()
     raw_response = httpx.Response(