Support caching controls and implicit cache tracking in interactions_nodes

This commit is contained in:
silveroxides
2026-10-04 22:01:18 +02:00
parent 27c4e76abc
commit 639ab6807b
2 changed files with 72 additions and 7 deletions
+33 -7
View File
@@ -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)
+39
View File
@@ -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