nodes/silveroxides-ComfyUI_Gemini_Expanded_API
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Support configurable enterprise safety settings on GenerateContent API

Browse Source
This commit is contained in:
silveroxides
2026-10-16 07:15:43 +02:00
parent 66ae9a5e88
commit 5fe0479c43
3 changed files with 201 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+8 -1
View File
@@ -1,5 +1,11 @@
from comfy_api.latest import ComfyExtension, io
from .gemini_nodes import GetKeyAPI, SSL_GeminiAPIKeyConfig, SSL_GeminiVideoConfig, SSL_GeminiTextPrompt
from .gemini_nodes import (
GetKeyAPI,
SSL_GeminiAPIKeyConfig,
SSL_GeminiSafetyConfig,
SSL_GeminiVideoConfig,
SSL_GeminiTextPrompt,
)
from .interactions_nodes import (
SSL_GeminiInteractionsAPIKeyConfig,
SSL_GeminiInteractionsSafetyConfig,
@@ -19,6 +25,7 @@ class GeminiExtension(ComfyExtension):
return [
GetKeyAPI,
SSL_GeminiAPIKeyConfig,
SSL_GeminiSafetyConfig,
SSL_GeminiVideoConfig,
SSL_GeminiResponseSchema,
SSL_GeminiTextPrompt,
gemini_nodes.py
+123 -3
View File
@@ -98,8 +98,67 @@ class GetKeyAPI(IO.ComfyNode):
class SSL_GeminiSafetyConfig(IO.ComfyNode):
GeminiSafetyConfig = IO.Custom("GEMINI_SAFETY_CONFIG")
THRESHOLDS = ["off", "block_none", "block_only_high", "block_medium_and_above", "block_low_and_above", "default"]
@classmethod
def define_schema(cls) -> IO.Schema:
return IO.Schema(
node_id="SSL_GeminiSafetyConfig",
display_name="Configure Gemini Safety",
category="API/Gemini",
inputs=[
IO.Combo.Input("harassment", options=cls.THRESHOLDS, default="off", tooltip="Negative or harmful comments targeting identity."),
IO.Combo.Input("hate_speech", options=cls.THRESHOLDS, default="off", tooltip="Content that promotes violence or incites hatred."),
IO.Combo.Input("sexually_explicit", options=cls.THRESHOLDS, default="off", tooltip="Contains references to sexual acts or other lewd content."),
IO.Combo.Input("dangerous_content", options=cls.THRESHOLDS, default="off", tooltip="Promotes, facilitates, or encourages harmful acts."),
IO.Combo.Input("civic_integrity", options=cls.THRESHOLDS, default="off", tooltip="Content that may mislead or influence elections."),
IO.Combo.Input("image_harassment", options=cls.THRESHOLDS, default="off", tooltip="Images depicting harassment."),
IO.Combo.Input("image_hate", options=cls.THRESHOLDS, default="off", tooltip="Images depicting hate speech."),
IO.Combo.Input("image_sexually_explicit", options=cls.THRESHOLDS, default="off", tooltip="Images containing sexually explicit content."),
IO.Combo.Input("image_dangerous_content", options=cls.THRESHOLDS, default="off", tooltip="Images depicting dangerous content."),
IO.Combo.Input("jailbreak", options=cls.THRESHOLDS, default="default", tooltip="Prompts designed to bypass safety filters. 'default' keeps filter disabled."),
IO.Combo.Input("method", options=["default", "severity", "probability"], default="severity", tooltip="Method for blocking content: severity vs probability score."),
],
outputs=[
cls.GeminiSafetyConfig.Output("safety_config")
]
)
@classmethod
def execute(cls, harassment: str = "off", hate_speech: str = "off",
sexually_explicit: str = "off", dangerous_content: str = "off",
civic_integrity: str = "off", image_harassment: str = "off",
image_hate: str = "off", image_sexually_explicit: str = "off",
image_dangerous_content: str = "off", jailbreak: str = "default",
method: str = "severity") -> IO.NodeOutput:
categories = {
"harassment": harassment,
"hate_speech": hate_speech,
"sexually_explicit": sexually_explicit,
"dangerous_content": dangerous_content,
"civic_integrity": civic_integrity,
"image_harassment": image_harassment,
"image_hate": image_hate,
"image_sexually_explicit": image_sexually_explicit,
"image_dangerous_content": image_dangerous_content,
"jailbreak": jailbreak,
}
settings = []
for cat, threshold in categories.items():
if threshold != "default":
entry: dict[str, Any] = {"type": cat, "threshold": threshold}
if method != "default":
entry["method"] = method
settings.append(entry)
return IO.NodeOutput(settings)
class SSL_GeminiAPIKeyConfig(IO.ComfyNode):
GemConfig = IO.Custom("GEMINI_CONFIG")
GeminiSafetyConfig = IO.Custom("GEMINI_SAFETY_CONFIG")
@classmethod
def define_schema(cls) -> IO.Schema:
@@ -118,6 +177,7 @@ class SSL_GeminiAPIKeyConfig(IO.ComfyNode):
IO.Boolean.Input("use_cache", default=False, optional=True, tooltip="Reuse unchanged Gemini input context through the API cache after a local result-cache miss."),
IO.Int.Input("cache_ttl_minutes", default=60, min=1, step=1, optional=True, tooltip="How long Gemini keeps explicitly cached context."),
IO.Int.Input("cache_seed", default=0, min=0, max=2147483647, control_after_generate=True, optional=True, tooltip="Gemini generation seed used when context caching is enabled."),
cls.GeminiSafetyConfig.Input("safety_override", optional=True, tooltip="Optional safety override settings from Configure Gemini Safety node."),
],
outputs=[
cls.GemConfig.Output("config")
@@ -128,7 +188,8 @@ class SSL_GeminiAPIKeyConfig(IO.ComfyNode):
def execute(cls, api_key: str, api_version: str, use_vertexai_env: bool, vertexai_express: bool,
vertexai_project: str | None = "", vertexai_location: str | None = "",
google_application_credentials: str | None = "", use_cache=False,
cache_ttl_minutes=60, cache_seed=0) -> IO.NodeOutput:
cache_ttl_minutes=60, cache_seed=0,
safety_override: GeminiSafetyConfig.Type | None = None) -> IO.NodeOutput:
config = {
"api_key": api_key,
"api_version": api_version,
@@ -140,6 +201,7 @@ class SSL_GeminiAPIKeyConfig(IO.ComfyNode):
"use_cache": use_cache,
"cache_ttl_minutes": cache_ttl_minutes,
"cache_seed": cache_seed,
"safety_override": safety_override,
}
return IO.NodeOutput(config)
@@ -460,6 +522,61 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
)
return thinking_level
CORE_HARM_CATEGORIES = [
"harassment",
"hate_speech",
"sexually_explicit",
"dangerous_content",
"civic_integrity",
"image_harassment",
"image_hate",
"image_sexually_explicit",
"image_dangerous_content",
]
@classmethod
def _resolve_safety_settings(cls, config: dict[str, Any]) -> list[types.SafetySetting]:
override = config.get("safety_override")
if override and isinstance(override, list) and override:
settings = []
for entry in override:
if isinstance(entry, types.SafetySetting):
settings.append(entry)
elif isinstance(entry, dict):
cat_name = entry.get("type") or entry.get("category", "")
cat_str = str(cat_name)
cat_enum = cat_str if cat_str.startswith("HARM_CATEGORY_") else f"HARM_CATEGORY_{cat_str.upper()}"
thresh = entry.get("threshold", "BLOCK_NONE")
thresh_enum = str(thresh).upper()
method = entry.get("method")
kwargs: dict[str, Any] = {
"category": cat_enum,
"threshold": thresh_enum,
}
if method and method != "default":
kwargs["method"] = str(method).upper()
settings.append(types.SafetySetting(**kwargs))
if settings:
return settings
is_enterprise = bool(config.get("use_vertexai_env", False) or config.get("vertexai_express", False))
if is_enterprise:
return [
types.SafetySetting(
category=f"HARM_CATEGORY_{cat.upper()}",
threshold="OFF",
method="SEVERITY",
)
for cat in cls.CORE_HARM_CATEGORIES
]
return [
types.SafetySetting(
category=f"HARM_CATEGORY_{cat.upper()}",
threshold="BLOCK_NONE",
)
for cat in cls.CORE_HARM_CATEGORIES
]
@classmethod
def _compute_fingerprint_and_check_cache(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
include_images, aspect_ratio, bypass_mode, thinking_budget, use_seed, seed,
@@ -720,9 +837,9 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
def _build_generate_content_config(cls, model, temperature, top_p, top_k, max_output_tokens, seed,
include_images, response_modalities, aspect_ratio, padded_system_instruction,
thinking_level, thinking_budget, include_thoughts, media_resolution,
image_size="None", allow_all_people=False):
image_size="None", allow_all_people=False, safety_settings=None):
# Centralized builder for GenerateContentConfig used by different model/feature branches
safety = [
safety = safety_settings if safety_settings is not None else [
types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="BLOCK_NONE"),
types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="BLOCK_NONE"),
types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="BLOCK_NONE"),
@@ -1116,6 +1233,8 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
else:
print("[WARNING] Vertex context cache skipped because inline cached media exceeds 10 MB; using full request")
resolved_safety_settings = cls._resolve_safety_settings(config)
generate_content_config = cls._build_generate_content_config(
model=model,
temperature=temperature,
@@ -1133,6 +1252,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
media_resolution=media_resolution,
image_size=image_size,
allow_all_people=bool(config.get("use_vertexai_env", False) or config.get("vertexai_express", False)),
safety_settings=resolved_safety_settings,
)
if schema_snapshot is not None:
tests/test_gemini_nodes.py
+70 -1
View File
@@ -824,8 +824,9 @@ def test_vertex_config_socket_ids_remain_stable():
"use_cache",
"cache_ttl_minutes",
"cache_seed",
"safety_override",
]
assert all(input_.optional for input_ in schema.inputs[-3:])
assert all(input_.optional for input_ in schema.inputs[-4:])
def test_image_inputs_use_clean_autogrow_socket_ids():
@@ -1626,3 +1627,71 @@ def test_timeout_returns_custom_fallback_without_caching(monkeypatch):
assert output[0] == "configured fallback"
assert output[3] == output[4] == ""
assert gemini_nodes.SSL_GeminiTextPrompt._cache == {}
def test_gemini_safety_config_node():
schema = gemini_nodes.SSL_GeminiSafetyConfig.define_schema()
assert schema.node_id == "SSL_GeminiSafetyConfig"
assert schema.category == "API/Gemini"
output = gemini_nodes.SSL_GeminiSafetyConfig.execute()
custom_safety = output[0]
assert isinstance(custom_safety, list)
assert len(custom_safety) == 9 # default jailbreak is 'default' so omitted
assert all(s["threshold"] == "off" for s in custom_safety)
assert all(s["method"] == "severity" for s in custom_safety)
def test_gemini_safety_override_reached_generate_content(structured_client):
safety_out = gemini_nodes.SSL_GeminiSafetyConfig.execute(
harassment="block_low_and_above",
hate_speech="block_medium_and_above",
method="probability",
)
custom_safety = safety_out[0]
config = _config(
api_key="test-key",
safety_override=custom_safety,
)
args = _execute_kwargs(config)
gemini_nodes.SSL_GeminiTextPrompt.execute(**args)
call = structured_client.calls[0]
sent_safety = call["config"].safety_settings
assert sent_safety is not None
assert len(sent_safety) == 9
harass = next(s for s in sent_safety if "HARASSMENT" in str(s.category) and "IMAGE" not in str(s.category))
assert "BLOCK_LOW_AND_ABOVE" in str(harass.threshold)
assert "PROBABILITY" in str(harass.method)
hate = next(s for s in sent_safety if "HATE_SPEECH" in str(s.category))
assert "BLOCK_MEDIUM_AND_ABOVE" in str(hate.threshold)
assert "PROBABILITY" in str(hate.method)
def test_gemini_enterprise_auto_safety_resolution(structured_client):
config = _config(api_key="test-key", vertexai_express=True)
args = _execute_kwargs(config)
gemini_nodes.SSL_GeminiTextPrompt.execute(**args)
call = structured_client.calls[0]
sent_safety = call["config"].safety_settings
assert sent_safety is not None
assert len(sent_safety) == 9
assert all("OFF" in str(s.threshold) for s in sent_safety)
assert all("SEVERITY" in str(s.method) for s in sent_safety)
def test_gemini_standard_auto_safety_resolution(structured_client):
config = _config(api_key="test-key")
args = _execute_kwargs(config)
gemini_nodes.SSL_GeminiTextPrompt.execute(**args)
call = structured_client.calls[0]
sent_safety = call["config"].safety_settings
assert sent_safety is not None
assert len(sent_safety) == 9
assert all("BLOCK_NONE" in str(s.threshold) for s in sent_safety)
assert all(s.method is None for s in sent_safety)