From 92529f388e98d1b3953ba645e9c459d6f6145ef6 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Fri, 28 Aug 2026 15:08:28 +0200 Subject: [PATCH] feat(cache): add Gemini context caching --- README.md | 4 + gemini_nodes.py | 177 ++++++++++++++++++++++++++-- tests/test_gemini_nodes.py | 236 +++++++++++++++++++++++++++++++++++++ 3 files changed, 407 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index a7a9b43..4e0f5bb 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,7 @@ Use the `Expanded Gemini Text/Image` node for text generation: - Enter your prompt text in `prompt`. - Adjust generation parameters (temperature, top_p, top_k, etc.). - If needed, connect an image to the `input_image` input for image understanding. +- Enable `use_cache` on the configuration node to reuse unchanged prompt, system instruction, and image context through Gemini after the local result cache misses. ### 3. Proxy Settings @@ -71,6 +72,9 @@ If you are in China or other regions requiring a proxy: ### API Key Configuration Node - `api_key`: Google Gemini API key +- `use_cache`: Enable Gemini API context caching. Disabled by default. +- `cache_ttl_minutes`: Number of minutes Gemini retains cached context. +- `cache_seed`: Generation seed used while `use_cache` is enabled. ### Text Generation Node diff --git a/gemini_nodes.py b/gemini_nodes.py index adb732f..a1e8f88 100644 --- a/gemini_nodes.py +++ b/gemini_nodes.py @@ -10,7 +10,7 @@ from io import BytesIO import folder_paths # type: ignore[reportMissingImports] from comfy_api.latest import ComfyExtension, UI, IO # type: ignore[reportMissingImports] from google import genai -from google.genai import types +from google.genai import errors, types import time import traceback import threading @@ -115,6 +115,9 @@ class SSL_GeminiAPIKeyConfig(IO.ComfyNode): IO.String.Input("vertexai_project", optional=True), IO.String.Input("vertexai_location", optional=True), 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="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, optional=True, tooltip="Gemini generation seed used when context caching is enabled."), ], outputs=[ cls.GemConfig.Output("config") @@ -124,7 +127,8 @@ class SSL_GeminiAPIKeyConfig(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=False, + cache_ttl_minutes=60, cache_seed=0) -> IO.NodeOutput: config = { "api_key": api_key, "api_version": api_version, @@ -133,6 +137,9 @@ class SSL_GeminiAPIKeyConfig(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) @@ -143,6 +150,8 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): _cache: dict = {} _seed_map_cache: dict = {} # Maps (input_seed, fingerprint_without_seed) -> successful_gemini_seed _client_cache: dict = {} # Maps client_key tuple -> genai.Client instance + _context_cache: dict = {} + _context_cache_lock = threading.Lock() GEMINI_3_7_FLASH = "gemini-3.7-flash" # Placeholder only. Keep this out of the model combo until Google's hosted-model @@ -333,7 +342,8 @@ class SSL_GeminiTextPrompt(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): + retry_pattern="", max_retries=3, use_cache=False, + cache_ttl_minutes=60, cache_seed=0): # 1. Hashing Images def get_tensor_hash(tensor): @@ -384,6 +394,8 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): # If none of the above, all effective thinking/image params remain default/ignored 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()}" @@ -403,6 +415,9 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): eff_thinking_budget, # EFFECTIVE thinking_budget use_seed, int(seed) if use_seed else 0, # Only use seed in cache if use_seed is True + bool(use_cache), + int(cache_ttl_minutes) if use_cache else 0, + int(cache_seed) if use_cache and use_seed else 0, image_hashes, bool(use_proxy), str(proxy_host), @@ -488,6 +503,87 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): client_args = {"proxy": proxy_url} if proxy_url else None return types.HttpOptions(api_version=api_version, client_args=client_args) + @staticmethod + def _build_context_cache_key(client_key, model, system_instruction, prompt, + image_hashes, media_resolution): + return ( + client_key, + model, + system_instruction, + prompt, + tuple(image_hashes), + str(media_resolution), + ) + + @staticmethod + def _context_cache_expiry(cache, fallback_ttl_minutes): + expire_time = getattr(cache, "expire_time", None) + if expire_time is not None: + return expire_time.timestamp() + return time.time() + (int(fallback_ttl_minutes) * 60) + + @classmethod + def _get_or_create_context_cache(cls, client, context_key, model, contents, + system_instruction, cache_ttl_minutes): + ttl = f"{int(cache_ttl_minutes) * 60}s" + with cls._context_cache_lock: + entry = cls._context_cache.get(context_key) + if entry is not None and entry["expires_at"] > time.time() + 1: + if entry["ttl_minutes"] == int(cache_ttl_minutes): + print(f"[INFO] Reusing Gemini context cache {entry['name']}") + return entry["name"] + try: + updated = client.caches.update( + name=entry["name"], + config=types.UpdateCachedContentConfig(ttl=ttl), + ) + entry = { + "name": updated.name, + "expires_at": cls._context_cache_expiry(updated, cache_ttl_minutes), + "ttl_minutes": int(cache_ttl_minutes), + } + cls._context_cache[context_key] = entry + print(f"[INFO] Updated Gemini context cache TTL to {cache_ttl_minutes} minutes") + return entry["name"] + except Exception: + cls._context_cache.pop(context_key, None) + elif entry is not None: + cls._context_cache.pop(context_key, None) + + display_hash = hashlib.sha256(repr(context_key).encode("utf-8")).hexdigest()[:16] + created = client.caches.create( + model=model, + config=types.CreateCachedContentConfig( + contents=contents, + system_instruction=system_instruction, + display_name=f"comfyui-{display_hash}", + ttl=ttl, + ), + ) + if not created.name: + raise ValueError("Gemini did not return a context cache name.") + cls._context_cache[context_key] = { + "name": created.name, + "expires_at": cls._context_cache_expiry(created, cache_ttl_minutes), + "ttl_minutes": int(cache_ttl_minutes), + } + print(f"[INFO] Created Gemini context cache {created.name}") + return created.name + + @classmethod + def _invalidate_context_cache(cls, context_key): + with cls._context_cache_lock: + cls._context_cache.pop(context_key, None) + + @staticmethod + def _is_missing_context_cache_error(error): + if isinstance(error, errors.APIError) and error.status == 404: + return True + message = str(error).lower() + cache_reference = any(term in message for term in ("cached content", "cached_content", "cachedcontents")) + missing_reference = any(term in message for term in ("expired", "invalid", "not found")) + return cache_reference and missing_reference + @classmethod 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, @@ -566,19 +662,22 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): @classmethod def execute(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens, include_images, aspect_ratio, bypass_mode, thinking_budget, - image_inputs: IO.Autogrow.Type | None = None, use_proxy=False, proxy_host="127.0.0.1", proxy_port=7890, use_seed=False, seed=0, timeout=30, include_thoughts=False, thinking_level=None, media_resolution=None, - retry_pattern="", max_retries=3, timeout_fallback_text="") -> IO.NodeOutput: + retry_pattern="", max_retries=3, timeout_fallback_text="", + image_inputs: IO.Autogrow.Type | None = None) -> IO.NodeOutput: print(f"[INFO] SSL_GeminiTextPrompt 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)) fingerprint, cached = cls._compute_fingerprint_and_check_cache( config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens, include_images, aspect_ratio, bypass_mode, thinking_budget, use_seed, seed, image_inputs, use_proxy, proxy_host, proxy_port, timeout, include_thoughts, thinking_level, media_resolution, - retry_pattern, max_retries + retry_pattern, max_retries, use_cache, cache_ttl_minutes, cache_seed ) if cached is not None: @@ -598,8 +697,9 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): padded_system_instruction = cls._pad_text_with_joiners(system_instruction) print(padded_system_instruction) - actual_seed = cls._handle_seed(use_seed, seed) - input_seed = seed # Store the original input seed for cache key + api_seed = cache_seed if use_cache else seed + actual_seed = cls._handle_seed(use_seed, api_seed) + input_seed = api_seed # Store the original API seed for retry cache key # Check if we have a cached successful gemini seed for this input seed # Build a cache key that excludes the seed itself (to match different gemini seeds for same input) @@ -739,6 +839,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): # Prepare contents (images + prompt) image_frames, image_batch_counts = cls._ordered_image_frames(image_inputs) + context_image_hashes = [] if image_frames: try: @@ -750,6 +851,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): img_byte_arr = BytesIO() pil_img.save(img_byte_arr, format='PNG') img_bytes = img_byte_arr.getvalue() + context_image_hashes.append(hashlib.sha256(img_bytes).hexdigest()) if model in cls.MEDIA_RES_MODELS and media_resolution is not None and media_resolution != "unspecified": img_part = types.Part.from_bytes( data=img_bytes, @@ -784,6 +886,14 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): # Logic Update: Only request IMAGE modality if include_images is TRUE AND it's the correct model # Otherwise we stick to TEXT modality response_modalities = ["IMAGE", "TEXT"] if (include_images and model in cls.IMAGE_MODELS) else ["TEXT"] + context_key = cls._build_context_cache_key( + client_key, + model, + padded_system_instruction, + padded_prompt, + context_image_hashes, + media_resolution, + ) generate_content_config = cls._build_generate_content_config( model=model, @@ -791,7 +901,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): top_p=top_p, top_k=top_k, max_output_tokens=max_output_tokens, - seed=seed, + seed=api_seed, include_images=include_images, response_modalities=response_modalities, aspect_ratio=aspect_ratio, @@ -807,11 +917,58 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): generate_content_config.seed = actual_seed except Exception: pass + elif use_cache: + try: + generate_content_config.seed = None + except Exception: + pass # API call in background thread start_time = time.time() result_queue: "queue.Queue[Tuple[str, Any]]" = queue.Queue() + def generate_once(): + if not use_cache: + return client.models.generate_content( + model=model, + contents=contents, + config=generate_content_config, + ) + + for cache_attempt in range(2): + try: + cache_name = cls._get_or_create_context_cache( + client, + context_key, + model, + contents, + padded_system_instruction, + cache_ttl_minutes, + ) + except Exception as cache_error: + print(f"[WARNING] Gemini context cache unavailable; using full request: {cache_error}") + return client.models.generate_content( + model=model, + contents=contents, + config=generate_content_config, + ) + + cached_config = generate_content_config.model_copy(deep=True) + cached_config.cached_content = cache_name + cached_config.system_instruction = None + try: + return client.models.generate_content( + model=model, + contents=" ", + config=cached_config, + ) + except Exception as cache_error: + if cache_attempt == 0 and cls._is_missing_context_cache_error(cache_error): + cls._invalidate_context_cache(context_key) + print("[INFO] Recreating expired Gemini context cache") + continue + raise + def api_call(): last_api_exception = None api_response = None @@ -823,7 +980,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): generate_content_config.seed = actual_seed + attempt except Exception: pass - response = client.models.generate_content(model=model, contents=contents, config=generate_content_config) + response = generate_once() if not (response.candidates and getattr(response.candidates[0].content, 'parts', None)): finish_reason = "UNKNOWN" if response.candidates: diff --git a/tests/test_gemini_nodes.py b/tests/test_gemini_nodes.py index a54391c..3e1160a 100644 --- a/tests/test_gemini_nodes.py +++ b/tests/test_gemini_nodes.py @@ -66,6 +66,238 @@ def test_image_autogrow_is_last_input(): assert schema.inputs[-1].id == "image_inputs" +def test_context_cache_reuses_full_input_after_local_fingerprint_miss(monkeypatch): + captured = {"creates": [], "generations": []} + + class FakeCaches: + @staticmethod + def create(**kwargs): + captured["creates"].append(kwargs) + return SimpleNamespace(name="cachedContents/test", expire_time=None) + + @staticmethod + def update(**kwargs): + raise AssertionError("TTL should not update") + + class FakeModels: + @staticmethod + def generate_content(**kwargs): + captured["generations"].append(kwargs) + part = SimpleNamespace(text="ok", inline_data=None) + content = SimpleNamespace(parts=[part]) + return SimpleNamespace(candidates=[SimpleNamespace(content=content)]) + + class FakeClient: + def __init__(self, **kwargs): + self.caches = FakeCaches() + self.models = FakeModels() + + monkeypatch.setattr(gemini_nodes.genai, "Client", FakeClient) + gemini_nodes.SSL_GeminiTextPrompt._cache.clear() + gemini_nodes.SSL_GeminiTextPrompt._client_cache.clear() + gemini_nodes.SSL_GeminiTextPrompt._context_cache.clear() + config = { + "api_key": "test-key", + "api_version": "v1", + "use_vertexai_env": False, + "vertexai_express": False, + "vertexai_project": "", + "vertexai_location": "", + "google_application_credentials": "", + "use_cache": True, + "cache_ttl_minutes": 15, + "cache_seed": 123, + } + kwargs = _execute_kwargs(config) + kwargs.update( + use_seed=True, + seed=1, + image_inputs={"image_1": torch.full((1, 2, 2, 3), 0.25)}, + ) + + first = gemini_nodes.SSL_GeminiTextPrompt.execute(**kwargs) + kwargs["seed"] = 2 + second = gemini_nodes.SSL_GeminiTextPrompt.execute(**kwargs) + third = gemini_nodes.SSL_GeminiTextPrompt.execute(**kwargs) + + assert first[0] == second[0] == third[0] == "ok" + assert len(captured["creates"]) == 1 + create_config = captured["creates"][0]["config"] + assert create_config.system_instruction == "test system instruction" + assert create_config.ttl == "900s" + assert create_config.contents[0].parts[-1].text == "test prompt" + assert len(create_config.contents[0].parts) == 2 + assert len(captured["generations"]) == 2 + assert all(call["contents"] == " " for call in captured["generations"]) + assert all(call["config"].cached_content == "cachedContents/test" for call in captured["generations"]) + assert all(call["config"].system_instruction is None for call in captured["generations"]) + assert all(call["config"].seed == 123 for call in captured["generations"]) + + +def test_context_cache_creation_failure_uses_original_request(monkeypatch): + captured = {} + + class FakeCaches: + @staticmethod + def create(**kwargs): + raise RuntimeError("unsupported cache") + + class FakeModels: + @staticmethod + def generate_content(**kwargs): + captured.update(kwargs) + part = SimpleNamespace(text="ok", inline_data=None) + content = SimpleNamespace(parts=[part]) + return SimpleNamespace(candidates=[SimpleNamespace(content=content)]) + + class FakeClient: + def __init__(self, **kwargs): + self.caches = FakeCaches() + self.models = FakeModels() + + monkeypatch.setattr(gemini_nodes.genai, "Client", FakeClient) + gemini_nodes.SSL_GeminiTextPrompt._cache.clear() + gemini_nodes.SSL_GeminiTextPrompt._client_cache.clear() + gemini_nodes.SSL_GeminiTextPrompt._context_cache.clear() + config = { + "api_key": "test-key", + "api_version": "v1", + "use_vertexai_env": False, + "vertexai_express": False, + "vertexai_project": "", + "vertexai_location": "", + "google_application_credentials": "", + "use_cache": True, + "cache_ttl_minutes": 60, + "cache_seed": 456, + } + kwargs = _execute_kwargs(config) + kwargs.update( + model="gemini-3-pro-image", + include_images=True, + use_seed=True, + seed=1, + ) + + output = gemini_nodes.SSL_GeminiTextPrompt.execute(**kwargs) + + assert output[0] == "ok" + assert captured["contents"] == "test prompt" + assert captured["config"].cached_content is None + assert captured["config"].system_instruction[0].text == "test system instruction" + assert captured["config"].seed == 456 + + kwargs.update(use_seed=False, seed=2) + output = gemini_nodes.SSL_GeminiTextPrompt.execute(**kwargs) + + assert output[0] == "ok" + assert captured["config"].seed is None + + kwargs["config"] = config | {"use_cache": False} + output = gemini_nodes.SSL_GeminiTextPrompt.execute(**kwargs) + + assert output[0] == "ok" + assert captured["config"].seed == 2 + + +def test_missing_context_cache_is_recreated_once(monkeypatch): + captured = {"creates": 0, "generations": 0} + + class FakeCaches: + @staticmethod + def create(**kwargs): + captured["creates"] += 1 + return SimpleNamespace( + name=f"cachedContents/test-{captured['creates']}", + expire_time=None, + ) + + class FakeModels: + @staticmethod + def generate_content(**kwargs): + captured["generations"] += 1 + if captured["generations"] == 1: + raise RuntimeError("cached content not found") + part = SimpleNamespace(text="ok", inline_data=None) + content = SimpleNamespace(parts=[part]) + return SimpleNamespace(candidates=[SimpleNamespace(content=content)]) + + class FakeClient: + def __init__(self, **kwargs): + self.caches = FakeCaches() + self.models = FakeModels() + + monkeypatch.setattr(gemini_nodes.genai, "Client", FakeClient) + gemini_nodes.SSL_GeminiTextPrompt._cache.clear() + gemini_nodes.SSL_GeminiTextPrompt._client_cache.clear() + gemini_nodes.SSL_GeminiTextPrompt._context_cache.clear() + config = { + "api_key": "test-key", + "api_version": "v1", + "use_vertexai_env": False, + "vertexai_express": False, + "vertexai_project": "", + "vertexai_location": "", + "google_application_credentials": "", + "use_cache": True, + "cache_ttl_minutes": 60, + "cache_seed": 456, + } + kwargs = _execute_kwargs(config) + kwargs.update(use_seed=True, seed=1) + + output = gemini_nodes.SSL_GeminiTextPrompt.execute(**kwargs) + + assert output[0] == "ok" + assert captured == {"creates": 2, "generations": 2} + + +def test_context_cache_ttl_change_updates_existing_resource(): + captured = {"creates": 0, "updates": []} + + class FakeCaches: + @staticmethod + def create(**kwargs): + captured["creates"] += 1 + return SimpleNamespace(name="cachedContents/test", expire_time=None) + + @staticmethod + def update(**kwargs): + captured["updates"].append(kwargs) + return SimpleNamespace(name="cachedContents/test", expire_time=None) + + client = SimpleNamespace(caches=FakeCaches()) + gemini_nodes.SSL_GeminiTextPrompt._context_cache.clear() + + first = gemini_nodes.SSL_GeminiTextPrompt._get_or_create_context_cache( + client, ("context",), "gemini-2.0-flash", "prompt", "system", 15 + ) + second = gemini_nodes.SSL_GeminiTextPrompt._get_or_create_context_cache( + client, ("context",), "gemini-2.0-flash", "prompt", "system", 30 + ) + + assert first == second == "cachedContents/test" + assert captured["creates"] == 1 + assert len(captured["updates"]) == 1 + assert captured["updates"][0]["config"].ttl == "1800s" + + +def test_context_cache_identity_covers_every_cached_input(): + build_key = gemini_nodes.SSL_GeminiTextPrompt._build_context_cache_key + base = build_key(("standard", "key-a"), "model-a", "system-a", "prompt-a", ["image-a"], "low") + variants = { + build_key(("standard", "key-b"), "model-a", "system-a", "prompt-a", ["image-a"], "low"), + build_key(("standard", "key-a"), "model-b", "system-a", "prompt-a", ["image-a"], "low"), + build_key(("standard", "key-a"), "model-a", "system-b", "prompt-a", ["image-a"], "low"), + build_key(("standard", "key-a"), "model-a", "system-a", "prompt-b", ["image-a"], "low"), + build_key(("standard", "key-a"), "model-a", "system-a", "prompt-a", ["image-b"], "low"), + build_key(("standard", "key-a"), "model-a", "system-a", "prompt-a", ["image-a"], "high"), + } + + assert base not in variants + assert len(variants) == 6 + + def test_vertex_config_socket_ids_remain_stable(): schema = gemini_nodes.SSL_GeminiAPIKeyConfig.define_schema() @@ -77,7 +309,11 @@ def test_vertex_config_socket_ids_remain_stable(): "vertexai_project", "vertexai_location", "google_application_credentials", + "use_cache", + "cache_ttl_minutes", + "cache_seed", ] + assert all(input_.optional for input_ in schema.inputs[-3:]) def test_image_inputs_use_clean_autogrow_socket_ids():