From 5fe0479c431f7a3a065f8e6e100a1d9a7a073333 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Fri, 16 Oct 2026 07:15:43 +0200 Subject: [PATCH] Support configurable enterprise safety settings on GenerateContent API --- __init__.py | 9 ++- gemini_nodes.py | 126 ++++++++++++++++++++++++++++++++++++- tests/test_gemini_nodes.py | 71 ++++++++++++++++++++- 3 files changed, 201 insertions(+), 5 deletions(-) diff --git a/__init__.py b/__init__.py index 51728f1..08c5113 100644 --- a/__init__.py +++ b/__init__.py @@ -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, diff --git a/gemini_nodes.py b/gemini_nodes.py index 7954a4c..fe9f46c 100644 --- a/gemini_nodes.py +++ b/gemini_nodes.py @@ -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: diff --git a/tests/test_gemini_nodes.py b/tests/test_gemini_nodes.py index eca1fb0..6d7e947 100644 --- a/tests/test_gemini_nodes.py +++ b/tests/test_gemini_nodes.py @@ -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)