From 639ab6807b00a4bbad0199843fb188df8c85cd83 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sun, 4 Oct 2026 22:01:18 +0200 Subject: [PATCH] Support caching controls and implicit cache tracking in interactions_nodes --- interactions_nodes.py | 40 ++++++++++++++++++++++++++------ tests/test_interactions_nodes.py | 39 +++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 7 deletions(-) diff --git a/interactions_nodes.py b/interactions_nodes.py index 2a78cc9..0593a94 100644 --- a/interactions_nodes.py +++ b/interactions_nodes.py @@ -47,6 +47,9 @@ class SSL_GeminiInteractionsAPIKeyConfig(IO.ComfyNode): IO.String.Input("vertexai_project", optional=True, tooltip="Google Cloud Project ID."), IO.String.Input("vertexai_location", optional=True, tooltip="Google Cloud Location/Region (e.g. us-central1)."), IO.String.Input("google_application_credentials", default="", optional=True, multiline=False, tooltip="Optional absolute path to an application default credentials JSON file for Vertex AI."), + 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."), ], outputs=[ cls.GemConfig.Output("config") @@ -56,7 +59,8 @@ class SSL_GeminiInteractionsAPIKeyConfig(IO.ComfyNode): @classmethod 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 = "") -> IO.NodeOutput: + google_application_credentials: str | None = "", use_cache: bool = False, + cache_ttl_minutes: int = 60, cache_seed: int = 0) -> IO.NodeOutput: config = { "api_key": api_key, "api_version": api_version, @@ -65,6 +69,9 @@ class SSL_GeminiInteractionsAPIKeyConfig(IO.ComfyNode): "vertexai_project": vertexai_project, "vertexai_location": vertexai_location, "google_application_credentials": google_application_credentials, + "use_cache": use_cache, + "cache_ttl_minutes": cache_ttl_minutes, + "cache_seed": cache_seed, } return IO.NodeOutput(config) @@ -352,7 +359,8 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): image_inputs: IO.Autogrow.Type | None = None, use_proxy=False, proxy_host="127.0.0.1", proxy_port=7890, timeout=30, include_thoughts=False, thinking_level=None, media_resolution=None, - retry_pattern="", max_retries=3, image_size="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): def get_tensor_hash(tensor): @@ -391,6 +399,8 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): eff_include_thoughts = include_thoughts sanitized_config = dict(config) + for cache_setting in ("use_cache", "cache_ttl_minutes", "cache_seed"): + sanitized_config.pop(cache_setting, None) api_key = sanitized_config.get("api_key") if api_key: sanitized_config["api_key"] = f"sha256:{hashlib.sha256(str(api_key).encode('utf-8')).hexdigest()}" @@ -407,6 +417,9 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): str(bypass_mode), use_seed, int(seed) if use_seed else 0, + bool(use_cache), + int(cache_ttl_minutes) if use_cache else 0, + int(cache_seed) if use_cache and use_seed else 0, video_hash, video_mime_type, video_fps, @@ -511,6 +524,10 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): response_schema: ResponseSchema.Type | None = None) -> IO.NodeOutput: print(f"[INFO] SSL_GeminiInteractionsTextPrompt execute called, model: {model}") + use_cache = bool(config.get("use_cache", False)) + cache_ttl_minutes = int(config.get("cache_ttl_minutes", 60)) + cache_seed = int(config.get("cache_seed", 0)) + schema_identity = None schema_snapshot = None if response_schema is not None: @@ -534,7 +551,7 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): ) except Exception as e: print(f"[ERROR] Error processing input video: {e}") - output_seed = seed if use_seed else 0 + output_seed = cache_seed if use_cache and use_seed else seed if use_seed else 0 return IO.NodeOutput(f"Error processing input video: {e}", cls.generate_empty_image(), output_seed, "", "", "") video_hash = hashlib.sha256(video_bytes).hexdigest() if video_bytes is not None else None @@ -545,8 +562,10 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): image_inputs, use_proxy, proxy_host, proxy_port, timeout, include_thoughts, thinking_level, media_resolution, - retry_pattern, max_retries, image_size=image_size, - store=store, previous_interaction_id=previous_interaction_id, + retry_pattern, max_retries, use_cache=use_cache, + cache_ttl_minutes=cache_ttl_minutes, cache_seed=cache_seed, + image_size=image_size, store=store, + previous_interaction_id=previous_interaction_id, schema_identity=schema_identity, ) @@ -564,8 +583,9 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): if bypass_mode == "system_instruction" or bypass_mode == "both": padded_system_instruction = cls._pad_text_with_joiners(system_instruction) - actual_seed = cls._handle_seed(use_seed, seed) - input_seed = seed + api_seed = cache_seed if use_cache else seed + actual_seed = cls._handle_seed(use_seed, api_seed) + input_seed = api_seed seed_cache_key = (input_seed, fingerprint) if use_seed and retry_pattern and max_retries > 0: @@ -893,6 +913,12 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode): if not current_text_output and model_output_texts: current_text_output = "".join(model_output_texts) + usage = getattr(api_response, "usage", None) + if usage: + cached_tokens = getattr(usage, "total_cached_tokens", 0) or getattr(usage, "cached_content_token_count", 0) or 0 + if cached_tokens: + print(f"[INFO] Gemini Interactions implicit cache hit: {cached_tokens} cached tokens") + # Also check output_image convenience property if current_image_tensor is None and hasattr(api_response, "output_image"): out_img = getattr(api_response, "output_image", None) diff --git a/tests/test_interactions_nodes.py b/tests/test_interactions_nodes.py index 3fa8d3f..857323d 100644 --- a/tests/test_interactions_nodes.py +++ b/tests/test_interactions_nodes.py @@ -60,6 +60,7 @@ def mock_interactions_client(monkeypatch): output_text=resp.get("text", ""), output_image=out_img, steps=steps, + usage=SimpleNamespace(total_cached_tokens=resp.get("cached_tokens", 0)) if "cached_tokens" in resp else None, ) class FakeClient: @@ -399,3 +400,41 @@ def test_timeout_fallback(mock_interactions_client, monkeypatch): output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert output[0] == "Custom timeout message" + + +def test_interactions_api_key_config_with_cache(): + output = interactions_nodes.SSL_GeminiInteractionsAPIKeyConfig.execute( + api_key="secret123", + api_version="v1beta", + use_vertexai_env=False, + vertexai_express=False, + use_cache=True, + cache_ttl_minutes=120, + cache_seed=999, + ) + config = output[0] + assert config["use_cache"] is True + assert config["cache_ttl_minutes"] == 120 + assert config["cache_seed"] == 999 + + +def test_caching_with_cache_seed_and_usage(mock_interactions_client): + mock_interactions_client.responses = [ + {"text": "Cached response", "id": "int_cache_hit_1", "cached_tokens": 4096}, + ] + + config = _config(use_cache=True, cache_seed=777) + args = _execute_kwargs(config=config, use_seed=True, seed=123) + out1 = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) + + assert out1[0] == "Cached response" + assert out1[2] == 777 # actual_seed matches cache_seed when use_cache=True + assert len(mock_interactions_client.calls) == 1 + assert mock_interactions_client.calls[0]["generation_config"]["seed"] == 777 + + # Second call with identical params hits local result cache + out2 = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) + assert out2[0] == "Cached response" + assert out2[2] == 777 + assert len(mock_interactions_client.calls) == 1 +