Skip to content

Commit c543ad8

Browse files
authored
fix(api-nodes-gemini): raise exception when no candidates due to safety block (Comfy-Org#11848)
1 parent 5ac1372 commit c543ad8

File tree

1 file changed

+19
-9
lines changed

1 file changed

+19
-9
lines changed

comfy_api_nodes/nodes_gemini.py

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,7 @@ def get_parts_by_type(response: GeminiGenerateContentResponse, part_type: Litera
130130
Returns:
131131
List of response parts matching the requested type.
132132
"""
133-
if response.candidates is None:
133+
if not response.candidates:
134134
if response.promptFeedback and response.promptFeedback.blockReason:
135135
feedback = response.promptFeedback
136136
raise ValueError(
@@ -141,14 +141,24 @@ def get_parts_by_type(response: GeminiGenerateContentResponse, part_type: Litera
141141
"try changing it to `IMAGE+TEXT` to view the model's reasoning and understand why image generation failed."
142142
)
143143
parts = []
144-
for part in response.candidates[0].content.parts:
145-
if part_type == "text" and part.text:
146-
parts.append(part)
147-
elif part.inlineData and part.inlineData.mimeType == part_type:
148-
parts.append(part)
149-
elif part.fileData and part.fileData.mimeType == part_type:
150-
parts.append(part)
151-
# Skip parts that don't match the requested type
144+
blocked_reasons = []
145+
for candidate in response.candidates:
146+
if candidate.finishReason and candidate.finishReason.upper() == "IMAGE_PROHIBITED_CONTENT":
147+
blocked_reasons.append(candidate.finishReason)
148+
continue
149+
if candidate.content is None or candidate.content.parts is None:
150+
continue
151+
for part in candidate.content.parts:
152+
if part_type == "text" and part.text:
153+
parts.append(part)
154+
elif part.inlineData and part.inlineData.mimeType == part_type:
155+
parts.append(part)
156+
elif part.fileData and part.fileData.mimeType == part_type:
157+
parts.append(part)
158+
159+
if not parts and blocked_reasons:
160+
raise ValueError(f"Gemini API blocked the request. Reasons: {blocked_reasons}")
161+
152162
return parts
153163

154164

0 commit comments

Comments
 (0)