diff --git a/__init__.py b/__init__.py index 1f17fbd..9dfef93 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,10 @@ from comfy_api.latest import ComfyExtension, io from .gemini_nodes import GetKeyAPI, SSL_GeminiAPIKeyConfig, SSL_GeminiVideoConfig, SSL_GeminiTextPrompt -from .interactions_nodes import SSL_GeminiInteractionsAPIKeyConfig, SSL_GeminiInteractionsTextPrompt +from .interactions_nodes import ( + SSL_GeminiInteractionsAPIKeyConfig, + SSL_GeminiInteractionsSafetyConfig, + SSL_GeminiInteractionsTextPrompt, +) from .response_schema import SSL_GeminiResponseSchema WEB_DIRECTORY = "./web" @@ -18,6 +22,7 @@ class GeminiExtension(ComfyExtension): SSL_GeminiResponseSchema, SSL_GeminiTextPrompt, SSL_GeminiInteractionsAPIKeyConfig, + SSL_GeminiInteractionsSafetyConfig, SSL_GeminiInteractionsTextPrompt, ] diff --git a/interactions_nodes.py b/interactions_nodes.py index 0593a94..3468573 100644 --- a/interactions_nodes.py +++ b/interactions_nodes.py @@ -50,6 +50,8 @@ class SSL_GeminiInteractionsAPIKeyConfig(IO.ComfyNode): IO.Boolean.Input("use_cache", default=False, optional=True, tooltip="Enable context caching workflow. Interactions API uses implicit server-side caching and session storage."), IO.Int.Input("cache_ttl_minutes", default=60, min=1, step=1, optional=True, tooltip="Context cache TTL minutes (retained for config compatibility)."), 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."), + IO.Combo.Input("safety_level", options=["block_none", "off", "block_only_high", "block_medium_and_above", "block_low_and_above", "default"], default="block_none", optional=True, tooltip="Default safety filter threshold for all harm categories. 'block_none' allows all content; 'off' completely disables filters."), + IO.Combo.Input("safety_method", options=["default", "probability", "severity"], default="default", optional=True, tooltip="Default method for blocking content: probability vs severity score."), ], outputs=[ cls.GemConfig.Output("config") @@ -60,7 +62,8 @@ class SSL_GeminiInteractionsAPIKeyConfig(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: bool = False, - cache_ttl_minutes: int = 60, cache_seed: int = 0) -> IO.NodeOutput: + cache_ttl_minutes: int = 60, cache_seed: int = 0, + safety_level: str = "block_none", safety_method: str = "default") -> IO.NodeOutput: config = { "api_key": api_key, "api_version": api_version, @@ -72,14 +75,74 @@ class SSL_GeminiInteractionsAPIKeyConfig(IO.ComfyNode): "use_cache": use_cache, "cache_ttl_minutes": cache_ttl_minutes, "cache_seed": cache_seed, + "safety_level": safety_level, + "safety_method": safety_method, } return IO.NodeOutput(config) +class SSL_GeminiInteractionsSafetyConfig(IO.ComfyNode): + GeminiSafetyConfig = IO.Custom("GEMINI_SAFETY_CONFIG") + THRESHOLDS = ["block_none", "off", "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_GeminiInteractionsSafetyConfig", + display_name="Configure Gemini Interactions Safety", + category="API/Gemini/Interactions", + inputs=[ + IO.Combo.Input("harassment", options=cls.THRESHOLDS, default="block_none", tooltip="Negative or harmful comments targeting identity."), + IO.Combo.Input("hate_speech", options=cls.THRESHOLDS, default="block_none", tooltip="Content that promotes violence or incites hatred."), + IO.Combo.Input("sexually_explicit", options=cls.THRESHOLDS, default="block_none", tooltip="Contains references to sexual acts or other lewd content."), + IO.Combo.Input("dangerous_content", options=cls.THRESHOLDS, default="block_none", tooltip="Promotes, facilitates, or encourages harmful acts."), + IO.Combo.Input("civic_integrity", options=cls.THRESHOLDS, default="block_none", tooltip="Content that may mislead or influence elections."), + IO.Combo.Input("image_harassment", options=cls.THRESHOLDS, default="block_none", tooltip="Images depicting harassment."), + IO.Combo.Input("image_hate", options=cls.THRESHOLDS, default="block_none", tooltip="Images depicting hate speech."), + IO.Combo.Input("image_sexually_explicit", options=cls.THRESHOLDS, default="block_none", tooltip="Images containing sexually explicit content."), + IO.Combo.Input("image_dangerous_content", options=cls.THRESHOLDS, default="block_none", tooltip="Images depicting dangerous content."), + IO.Combo.Input("jailbreak", options=cls.THRESHOLDS, default="block_none", tooltip="Prompts designed to bypass safety filters."), + IO.Combo.Input("method", options=["default", "probability", "severity"], default="default", tooltip="Method for blocking content: probability vs severity score."), + ], + outputs=[ + cls.GeminiSafetyConfig.Output("safety_config") + ] + ) + + @classmethod + def execute(cls, harassment: str = "block_none", hate_speech: str = "block_none", + sexually_explicit: str = "block_none", dangerous_content: str = "block_none", + civic_integrity: str = "block_none", image_harassment: str = "block_none", + image_hate: str = "block_none", image_sexually_explicit: str = "block_none", + image_dangerous_content: str = "block_none", jailbreak: str = "block_none", + method: str = "default") -> 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_GeminiInteractionsTextPrompt(IO.ComfyNode): GemConfig = IO.Custom("GEMINI_CONFIG") ResponseSchema = IO.Custom("GEMINI_RESPONSE_SCHEMA") GeminiVideoConfig = IO.Custom("GEMINI_VIDEO_CONFIG") + GeminiSafetyConfig = IO.Custom("GEMINI_SAFETY_CONFIG") _cache: dict = {} _seed_map_cache: dict = {} # Maps (input_seed, fingerprint) -> successful_gemini_seed _client_cache: dict = {} # Maps client_key tuple -> genai.Client instance @@ -90,6 +153,19 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): GEMINI_4_FLASH_PREVIEW = "gemini-4-flash-preview" + ALL_HARM_CATEGORIES = [ + "harassment", + "hate_speech", + "sexually_explicit", + "dangerous_content", + "civic_integrity", + "image_harassment", + "image_hate", + "image_sexually_explicit", + "image_dangerous_content", + "jailbreak", + ] + THINKING_MODELS = [ "gemini-1.5-pro-002", "gemini-2.0-flash-thinking-exp", "gemini-2.0-flash-thinking-exp-01-21", "gemini-2.0-flash-thinking-exp-1219", "gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-preview-04-17", "gemini-2.5-pro-exp-03-25", @@ -154,6 +230,9 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): IO.Combo.Input("image_size", options=["None", "512", "1K", "2K", "4K"], default="None", optional=True, tooltip="Generated image resolution. Gemini 3.1 Flash Lite Image supports only 1K; 512 is supported only by Gemini 3.1 Flash Image."), IO.Boolean.Input("store", default=True, optional=True, tooltip="Store interaction on server for multi-turn conversations."), IO.String.Input("previous_interaction_id", default="", optional=True, multiline=False, tooltip="Optional ID of previous interaction to continue multi-turn conversation."), + IO.Combo.Input("safety_level", options=["block_none", "off", "block_only_high", "block_medium_and_above", "block_low_and_above", "default"], default="block_none", optional=True, tooltip="Global safety filter threshold for all 10 text, image, and jailbreak harm categories. 'block_none' allows all content; 'off' completely disables filters."), + IO.Combo.Input("safety_method", options=["default", "probability", "severity"], default="default", optional=True, tooltip="Method for blocking content: probability vs severity score."), + cls.GeminiSafetyConfig.Input("safety_config", optional=True, tooltip="Optional granular safety configuration from Configure Gemini Interactions Safety node."), cls.GeminiVideoConfig.Input("video", optional=True, tooltip="Optional configured Gemini video input with embedded audio and sampling FPS."), cls.ResponseSchema.Input("response_schema", display_name="Response format", optional=True, tooltip="Connect an answer format to request the fields you defined. Leave disconnected for a normal answer."), @@ -352,6 +431,30 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): ) return thinking_level + @classmethod + def _resolve_safety_settings(cls, safety_config, safety_level, safety_method, config): + if safety_config is not None and isinstance(safety_config, list) and safety_config: + return list(safety_config) + + eff_level = safety_level + if eff_level in (None, "default"): + eff_level = config.get("safety_level", "block_none") + + eff_method = safety_method + if eff_method in (None, "default"): + eff_method = config.get("safety_method", "default") + + if eff_level == "default": + return None + + settings = [] + for cat in cls.ALL_HARM_CATEGORIES: + entry: dict[str, Any] = {"type": cat, "threshold": eff_level} + if eff_method != "default": + entry["method"] = eff_method + settings.append(entry) + return settings + @classmethod def _compute_fingerprint_and_check_cache(cls, config, prompt, system_instruction, model, max_output_tokens, include_images, aspect_ratio, bypass_mode, use_seed, seed, @@ -361,7 +464,8 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): include_thoughts=False, thinking_level=None, media_resolution=None, retry_pattern="", max_retries=3, use_cache=False, cache_ttl_minutes=60, cache_seed=0, image_size="None", - store=True, previous_interaction_id="", schema_identity=None): + store=True, previous_interaction_id="", safety_settings=None, + schema_identity=None): def get_tensor_hash(tensor): if tensor is None: @@ -435,6 +539,7 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): int(max_retries), bool(store), str(previous_interaction_id), + str(safety_settings), schema_identity, ) @@ -519,6 +624,8 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): include_thoughts=False, thinking_level=None, media_resolution=None, retry_pattern="", max_retries=3, timeout_fallback_text="", image_size="None", store=True, previous_interaction_id="", + safety_level="block_none", safety_method="default", + safety_config: GeminiSafetyConfig.Type | None = None, video: GeminiVideoConfig.Type | None = None, image_inputs: IO.Autogrow.Type | None = None, response_schema: ResponseSchema.Type | None = None) -> IO.NodeOutput: @@ -528,6 +635,8 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): cache_ttl_minutes = int(config.get("cache_ttl_minutes", 60)) cache_seed = int(config.get("cache_seed", 0)) + resolved_safety_settings = cls._resolve_safety_settings(safety_config, safety_level, safety_method, config) + schema_identity = None schema_snapshot = None if response_schema is not None: @@ -566,6 +675,7 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): cache_ttl_minutes=cache_ttl_minutes, cache_seed=cache_seed, image_size=image_size, store=store, previous_interaction_id=previous_interaction_id, + safety_settings=resolved_safety_settings, schema_identity=schema_identity, ) @@ -815,15 +925,6 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): if include_thoughts: generation_config["thinking_summaries"] = "auto" - # Safety settings - safety_settings = [ - {"type": "harassment", "threshold": "block_none"}, - {"type": "hate_speech", "threshold": "block_none"}, - {"type": "sexually_explicit", "threshold": "block_none"}, - {"type": "dangerous_content", "threshold": "block_none"}, - {"type": "civic_integrity", "threshold": "block_none"}, - ] - # Background worker queue for API call start_time = time.time() result_queue: "queue.Queue[Tuple[str, Any]]" = queue.Queue() @@ -836,8 +937,9 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): kwargs: dict[str, Any] = { "model": model, "input": input_contents, - "safety_settings": safety_settings, } + if resolved_safety_settings is not None: + kwargs["safety_settings"] = resolved_safety_settings if padded_system_instruction: kwargs["system_instruction"] = padded_system_instruction if call_gen_config: @@ -868,6 +970,28 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): return try: + # Check for safety blocks or API errors + api_errors = getattr(api_response, "errors", None) or [] + if api_errors: + err_details = [] + for err in api_errors: + msg = getattr(err, "message", None) or str(err) + code = getattr(err, "code", None) + err_details.append(f"[{code}] {msg}" if code else msg) + raise ValueError(f"Interaction blocked/failed: {'; '.join(err_details)}") + + api_status = getattr(api_response, "status", None) + if api_status in ("failed", "incomplete"): + step_errors = [] + for step in getattr(api_response, "steps", None) or []: + s_err = getattr(step, "error", None) + if s_err: + s_msg = getattr(s_err, "message", None) or str(s_err) + step_errors.append(s_msg) + if step_errors: + raise ValueError(f"Interaction {api_status}: {'; '.join(step_errors)}") + raise ValueError(f"Interaction ended with status: {api_status}") + current_text_output = getattr(api_response, "output_text", "") or "" current_thoughts_output = "" current_image_tensor = None diff --git a/tests/test_interactions_nodes.py b/tests/test_interactions_nodes.py index 857323d..031f48d 100644 --- a/tests/test_interactions_nodes.py +++ b/tests/test_interactions_nodes.py @@ -56,7 +56,8 @@ def mock_interactions_client(monkeypatch): return SimpleNamespace( id=resp.get("id", "int_123"), - status="completed", + status=resp.get("status", "completed"), + errors=resp.get("errors", None), output_text=resp.get("text", ""), output_image=out_img, steps=steps, @@ -183,6 +184,9 @@ def test_basic_text_generation(mock_interactions_client): assert call["system_instruction"] == "test system instruction" assert call["store"] is True assert "safety_settings" in call + assert len(call["safety_settings"]) == 10 + assert {s["type"] for s in call["safety_settings"]} == set(interactions_nodes.SSL_GeminiInteractionsTextPrompt.ALL_HARM_CATEGORIES) + assert all(s["threshold"] == "block_none" for s in call["safety_settings"]) def test_multimodal_image_input(mock_interactions_client): @@ -438,3 +442,69 @@ def test_caching_with_cache_seed_and_usage(mock_interactions_client): assert out2[2] == 777 assert len(mock_interactions_client.calls) == 1 + +def test_safety_config_node(): + schema = interactions_nodes.SSL_GeminiInteractionsSafetyConfig.define_schema() + assert schema.node_id == "SSL_GeminiInteractionsSafetyConfig" + assert schema.category == "API/Gemini/Interactions" + + output = interactions_nodes.SSL_GeminiInteractionsSafetyConfig.execute( + harassment="off", + hate_speech="block_none", + jailbreak="block_only_high", + method="severity", + ) + settings = output[0] + assert isinstance(settings, list) + settings_dict = {s["type"]: s for s in settings} + assert settings_dict["harassment"]["threshold"] == "off" + assert settings_dict["harassment"]["method"] == "severity" + assert settings_dict["hate_speech"]["threshold"] == "block_none" + assert settings_dict["jailbreak"]["threshold"] == "block_only_high" + + +def test_safety_preferences_off_and_method(mock_interactions_client): + mock_interactions_client.responses = [{"text": "Unrestricted response", "id": "int_safe_1"}] + args = _execute_kwargs(safety_level="off", safety_method="severity") + output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) + + assert output[0] == "Unrestricted response" + call = mock_interactions_client.calls[0] + assert "safety_settings" in call + assert len(call["safety_settings"]) == 10 + assert all(s["threshold"] == "off" for s in call["safety_settings"]) + assert all(s["method"] == "severity" for s in call["safety_settings"]) + + +def test_safety_config_input_override(mock_interactions_client): + mock_interactions_client.responses = [{"text": "Granular safety response", "id": "int_safe_2"}] + custom_safety = [ + {"type": "harassment", "threshold": "block_only_high"}, + {"type": "jailbreak", "threshold": "off"}, + ] + args = _execute_kwargs( + safety_level="block_none", # Should be overridden by safety_config + safety_config=custom_safety, + ) + output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) + + assert output[0] == "Granular safety response" + call = mock_interactions_client.calls[0] + assert call["safety_settings"] == custom_safety + + +def test_safety_block_error_reporting(mock_interactions_client): + mock_interactions_client.responses = [{ + "text": "", + "id": "int_blocked", + "errors": [SimpleNamespace(code="SAFETY", message="Content was blocked by safety policy")], + "status": "failed", + }] + + args = _execute_kwargs() + output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) + + assert output[0].startswith("API call/processing error:") + assert "Content was blocked by safety policy" in output[0] + +