df18925e8d
Preserve existing workflows while adding capability-based catalogs, dedicated image requests, resumable native video, validated audio and explicit chat tiers. Cover compatibility, billing-sensitive retries, credentials and cancellation. Co-authored-by: Arthur Reboul Salze <22814805+ArthurReboulSalze@users.noreply.github.com> Co-authored-by: Dookie <me@djdookie.de>
981 lines
48 KiB
Python
981 lines
48 KiB
Python
import requests
|
|
import json
|
|
import time
|
|
import base64
|
|
import io
|
|
import numpy as np
|
|
import torch
|
|
import tiktoken
|
|
from PIL import Image
|
|
import hashlib # Added for hashing PDF bytes in IS_CHANGED
|
|
import os
|
|
import re
|
|
import logging
|
|
from .chat_manager import ChatSessionManager
|
|
|
|
# Define a placeholder type name for PDF data.
|
|
# The actual input connection will accept '*' but we check the structure.
|
|
# Expecting a dictionary: {"filename": str, "bytes": bytes}
|
|
PDF_DATA_TYPE = "*" # Use '*' to accept any type, check structure later
|
|
|
|
class OpenRouterNode:
|
|
"""
|
|
OpenRouter chat, image, and video requests with optional media inputs.
|
|
Preserves Output, image, Stats, Credits and appends native VIDEO.
|
|
"""
|
|
|
|
models_cache = None
|
|
last_fetch_time = 0
|
|
cache_duration = 3600 # Cache duration in seconds (1 hour)
|
|
default_request_timeout = 120
|
|
min_request_timeout = 1
|
|
max_request_timeout = 3600
|
|
reasoning_effort_options = ("auto", "none", "minimal", "low", "medium", "high", "xhigh")
|
|
default_reasoning_effort = "auto"
|
|
|
|
def __init__(self):
|
|
self.chat_manager = ChatSessionManager()
|
|
self.last_video_job_id = ""
|
|
|
|
@staticmethod
|
|
def get_api_key(api_key_ui):
|
|
"""
|
|
Resolves the API key from:
|
|
1. UI input field (if not empty)
|
|
2. Environment variable 'LLM_KEY'
|
|
3. config file 'openrouter_api_key.json' in node directory
|
|
"""
|
|
if api_key_ui and api_key_ui.strip():
|
|
return api_key_ui.strip()
|
|
|
|
# Check environment variable
|
|
env_key = os.environ.get("LLM_KEY")
|
|
if env_key and env_key.strip():
|
|
return env_key.strip()
|
|
|
|
# Check JSON file
|
|
config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "openrouter_api_key.json")
|
|
if os.path.exists(config_path):
|
|
try:
|
|
with open(config_path, 'r') as f:
|
|
config = json.load(f)
|
|
file_key = config.get("api_key")
|
|
if file_key and file_key.strip():
|
|
return file_key.strip()
|
|
except Exception as e:
|
|
print(f"Error reading openrouter_api_key.json: {e}")
|
|
|
|
return ""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""
|
|
Defines the input specification for this node.
|
|
Includes optional inputs for image and PDF data.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"api_key": ("STRING", {
|
|
"multiline": False,
|
|
"default": ""
|
|
}),
|
|
"system_prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": "You are a helpful assistant."
|
|
}),
|
|
"user_message_box": ("STRING", {
|
|
"multiline": True,
|
|
"default": "Hello, how are you?"
|
|
}),
|
|
"model": (cls.fetch_openrouter_models(),),
|
|
"web_search": ("BOOLEAN", {"default": False}),
|
|
"cheapest": ("BOOLEAN", {"default": True}),
|
|
"fastest": ("BOOLEAN", {"default": False}),
|
|
"aspect_ratio": ([
|
|
"auto",
|
|
"1:1 (1024x1024)",
|
|
"2:3 (832x1248)",
|
|
"3:2 (1248x832)",
|
|
"3:4 (864x1184)",
|
|
"4:3 (1184x864)",
|
|
"4:5 (896x1152)",
|
|
"5:4 (1152x896)",
|
|
"9:16 (768x1344)",
|
|
"16:9 (1344x768)",
|
|
"21:9 (1536x672)",
|
|
"1:4 (google/gemini-3.1-flash-image-preview (Nano Banana 2) only)",
|
|
"4:1 (google/gemini-3.1-flash-image-preview (Nano Banana 2) only)",
|
|
"1:8 (google/gemini-3.1-flash-image-preview (Nano Banana 2) only)",
|
|
"8:1 (google/gemini-3.1-flash-image-preview (Nano Banana 2) only)",
|
|
], {"default": "auto"}),
|
|
"image_resolution": (["auto", "0.5K", "1K", "2K", "4K"], {"default": "auto"}),
|
|
"reasoning_effort": (list(cls.reasoning_effort_options), {"default": cls.default_reasoning_effort}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": "fixed"}),
|
|
"temperature": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 2.0,
|
|
"step": 0.01,
|
|
"display": "slider",
|
|
"round": 0.01,
|
|
}),
|
|
"pdf_engine": (["auto", "mistral-ocr", "pdf-text"], {"default": "auto"}),
|
|
"chat_mode": ("BOOLEAN", {"default": False}),
|
|
"request_timeout": ("INT", {
|
|
"default": cls.default_request_timeout,
|
|
"min": cls.min_request_timeout,
|
|
"max": cls.max_request_timeout,
|
|
"step": 1,
|
|
"display": "number",
|
|
}),
|
|
},
|
|
"optional": {
|
|
"pdf_data": (PDF_DATA_TYPE,), # Use '*' and check structure in generate_response
|
|
"user_message_input": ("STRING", {"forceInput": True}),
|
|
"audio_data": ("AUDIO",),
|
|
"request_type": (["chat", "image", "video"], {"default": "chat"}),
|
|
"service_tier": (["auto", "default", "flex", "priority", "ultrafast"], {"default": "auto"}),
|
|
"image_quality": (["auto", "low", "medium", "high"], {"default": "auto"}),
|
|
"image_background": (["auto", "opaque", "transparent"], {"default": "auto"}),
|
|
"video_mode": (["text_to_video", "first_frame", "first_last_frame", "reference_images"], {"default": "text_to_video"}),
|
|
"video_duration": ("STRING", {"default": "auto"}),
|
|
"video_resolution": ("STRING", {"default": "auto"}),
|
|
"video_generate_audio": ("BOOLEAN", {"default": False}),
|
|
"video_wait_timeout": ("INT", {"default": 900, "min": 1, "max": 3600}),
|
|
"video_job_id": ("STRING", {"default": ""}),
|
|
},
|
|
"hidden": {"unique_id": "UNIQUE_ID"},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "IMAGE", "STRING", "STRING", "VIDEO")
|
|
RETURN_NAMES = ("Output", "image", "Stats", "Credits", "video")
|
|
|
|
FUNCTION = "generate_response"
|
|
CATEGORY = "LLM"
|
|
|
|
@classmethod
|
|
def fetch_openrouter_models(cls):
|
|
"""Read the cached union without blocking ComfyUI's schema/event loop."""
|
|
from . import openrouter_catalog
|
|
snapshot = openrouter_catalog.get_catalog()
|
|
ids = {m["id"] for kind in ("chat", "image", "video")
|
|
for m in snapshot.get(kind, []) if isinstance(m.get("id"), str)}
|
|
return sorted(ids) or ["openai/gpt-4o"]
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, model, aspect_ratio="auto", image_resolution="auto",
|
|
image_quality="auto", image_background="auto"):
|
|
# Discovery is asynchronous. Do not reject a saved/manual model merely
|
|
# because the server's widget snapshot predates a catalog refresh.
|
|
return True
|
|
|
|
def validate_temperature(self, temperature):
|
|
"""
|
|
Validates and converts temperature value to float within acceptable range.
|
|
"""
|
|
try:
|
|
temp = float(temperature)
|
|
return max(0.0, min(2.0, temp)) # Clamp between 0.0 and 2.0
|
|
except (ValueError, TypeError):
|
|
return 1.0 # Return default if conversion fails
|
|
|
|
def validate_request_timeout(self, request_timeout):
|
|
"""
|
|
Validates and converts request timeout to seconds within an acceptable range.
|
|
"""
|
|
try:
|
|
timeout = int(request_timeout)
|
|
return max(self.min_request_timeout, min(self.max_request_timeout, timeout))
|
|
except (ValueError, TypeError):
|
|
return self.default_request_timeout
|
|
|
|
@classmethod
|
|
def validate_reasoning_effort(cls, reasoning_effort):
|
|
"""
|
|
Validates OpenRouter reasoning effort. "auto" means do not send a
|
|
reasoning override and let OpenRouter/model defaults apply.
|
|
"""
|
|
if isinstance(reasoning_effort, str):
|
|
normalized_effort = reasoning_effort.strip().lower()
|
|
if normalized_effort in cls.reasoning_effort_options:
|
|
return normalized_effort
|
|
return cls.default_reasoning_effort
|
|
|
|
def fetch_credits(self, api_key, timeout=None):
|
|
"""
|
|
Fetches the user's credits information from the OpenRouter API.
|
|
Returns a formatted string with remaining credits.
|
|
"""
|
|
api_key = self.get_api_key(api_key)
|
|
if not api_key:
|
|
return "API Key not provided."
|
|
|
|
url = "https://openrouter.ai/api/v1/credits"
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
"HTTP-Referer": "https://github.com/gabe-init/ComfyUI-Openrouter_node",
|
|
"X-Title": "ComfyUI OpenRouter LLM Node",
|
|
}
|
|
|
|
try:
|
|
validated_timeout = self.validate_request_timeout(timeout)
|
|
response = requests.get(url, headers=headers, timeout=validated_timeout)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
# Check if 'data' and expected keys exist
|
|
if "data" in result and "total_credits" in result["data"] and "total_usage" in result["data"]:
|
|
total_credits = result["data"]["total_credits"]
|
|
total_usage = result["data"]["total_usage"]
|
|
remaining = total_credits - total_usage
|
|
credits_text = f"Remaining: ${remaining:.3f}"
|
|
else:
|
|
credits_text = "Could not parse credit data from response."
|
|
|
|
return credits_text
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
# Provide more context about the error
|
|
error_message = f"Error fetching credits: {str(e)}"
|
|
if hasattr(e, 'response') and e.response is not None:
|
|
error_message += f" | Status Code: {e.response.status_code}"
|
|
return error_message
|
|
except json.JSONDecodeError:
|
|
return "Error fetching credits: Could not decode JSON response."
|
|
|
|
@staticmethod
|
|
def _safe_error(error, api_key=""):
|
|
message = str(error)
|
|
if api_key:
|
|
message = message.replace(api_key, "[redacted]")
|
|
return re.sub(r"data:[^\s\"']+", "[media omitted]", message)[:600]
|
|
|
|
def _remember_video_job(self, job_id, unique_id=None):
|
|
"""Expose recovery before polling, including when Comfy hides cancellation errors."""
|
|
if not isinstance(job_id, str) or not re.fullmatch(r"[A-Za-z0-9_-]{1,256}", job_id):
|
|
raise ValueError("OpenRouter returned an invalid video job ID.")
|
|
self.last_video_job_id = job_id
|
|
logger = logging.getLogger(__name__)
|
|
logger.info("OpenRouter video job: %s. Resume with video_job_id=%s.", job_id, job_id)
|
|
if unique_id is None:
|
|
return
|
|
try:
|
|
from server import PromptServer
|
|
server = PromptServer.instance
|
|
server.send_sync("openrouter.video_job", {"node_id": str(unique_id), "job_id": job_id},
|
|
sid=getattr(server, "client_id", None))
|
|
except Exception:
|
|
# UI delivery must never turn an accepted paid job into a failure.
|
|
logger.warning("OpenRouter video job: %s. Resume with video_job_id=%s; frontend update unavailable.",
|
|
job_id, job_id)
|
|
|
|
@staticmethod
|
|
def _chat_capabilities(model):
|
|
"""Wait for discovery once; retain catalog variants such as :free."""
|
|
from .openrouter_catalog import get_model, require_model
|
|
try:
|
|
return require_model("chat", model)
|
|
except ValueError as error:
|
|
base_model = model
|
|
while ":" in base_model and base_model.rsplit(":", 1)[1] in {"floor", "nitro", "online"}:
|
|
base_model = base_model.rsplit(":", 1)[0]
|
|
if base_model != model:
|
|
# The blocking lookup has finished; use its refreshed snapshot.
|
|
try:
|
|
return get_model("chat", base_model)
|
|
except ValueError:
|
|
pass
|
|
raise error
|
|
|
|
@staticmethod
|
|
def _chat_image_config(model, aspect_ratio, image_resolution):
|
|
"""Validate chat image controls using the matching image catalog record."""
|
|
from .openrouter_catalog import get_model
|
|
ratio = str(aspect_ratio).split(" ", 1)[0]
|
|
resolution = "512" if image_resolution == "0.5K" else image_resolution
|
|
try:
|
|
record = get_model("image", model)
|
|
parameters = record.get("supported_parameters") or {}
|
|
except ValueError:
|
|
parameters = {}
|
|
|
|
def supports(name, value):
|
|
descriptor = parameters.get(name)
|
|
return isinstance(descriptor, dict) and descriptor.get("type") == "enum" and value in descriptor.get("values", [])
|
|
|
|
config = {}
|
|
if ratio != "auto":
|
|
if not supports("aspect_ratio", ratio):
|
|
raise ValueError(f"Cannot verify aspect ratio {ratio} for this chat image model; choose auto or use the Image API.")
|
|
config["aspect_ratio"] = ratio
|
|
# The old node always stored 1K, including for GPT image models. It is
|
|
# neutral when unavailable; explicit non-default choices must validate.
|
|
if resolution == "1K" and not supports("resolution", resolution):
|
|
resolution = "auto"
|
|
if resolution != "auto":
|
|
if not model.startswith("google/gemini-") or not supports("resolution", resolution):
|
|
raise ValueError(f"Resolution {resolution} is not supported for this chat image model; choose auto or use the Image API.")
|
|
config["image_size"] = resolution
|
|
return config
|
|
|
|
def generate_response(self, api_key, system_prompt, user_message_box, model,
|
|
web_search, cheapest, fastest, temperature, pdf_engine, chat_mode,
|
|
request_timeout=120, aspect_ratio="auto", image_resolution="1K", seed=0,
|
|
pdf_data=None, user_message_input=None, reasoning_effort="auto",
|
|
audio_data=None, request_type="chat", service_tier="auto",
|
|
image_quality="auto", image_background="auto", video_mode="text_to_video",
|
|
video_duration="auto", video_resolution="auto", video_generate_audio=False,
|
|
video_wait_timeout=900, video_job_id="", unique_id=None, **kwargs):
|
|
"""Keep the original four output positions; append native video output."""
|
|
placeholder = torch.zeros((1, 1, 1, 3), dtype=torch.float32)
|
|
key = self.get_api_key(api_key)
|
|
self.last_video_job_id = video_job_id.strip() if request_type == "video" else ""
|
|
try:
|
|
if not key:
|
|
raise ValueError("API Key not provided. Set LLM_KEY or openrouter_api_key.json.")
|
|
if request_type not in ("chat", "image", "video"):
|
|
raise ValueError("request_type must be chat, image, or video")
|
|
if service_tier not in ("auto", "default", "flex", "priority", "ultrafast"):
|
|
raise ValueError("Unsupported service tier")
|
|
timeout = self.validate_request_timeout(request_timeout)
|
|
prompt = user_message_input if user_message_input is not None else user_message_box
|
|
ratio = str(aspect_ratio).split(" ", 1)[0]
|
|
if request_type == "chat":
|
|
if pdf_data is not None and (not isinstance(pdf_data, dict) or
|
|
not isinstance(pdf_data.get("bytes"), bytes) or not pdf_data["bytes"]):
|
|
raise ValueError("pdf_data must contain nonempty PDF bytes")
|
|
has_images = any(value is not None for name, value in kwargs.items() if re.fullmatch(r"image_\d+", name))
|
|
# 1K was an unconditional saved value in older text workflows.
|
|
image_controls = ratio != "auto" or image_resolution not in ("auto", "1K")
|
|
try:
|
|
capabilities = self._chat_capabilities(model)
|
|
except ValueError as error:
|
|
if audio_data is not None or has_images or image_controls:
|
|
raise ValueError("Cannot verify the selected chat model's media capabilities. Refresh models before submitting audio, images, or image settings.") from error
|
|
capabilities = {} # Plain custom chat IDs may be absent from discovery.
|
|
architecture = capabilities.get("architecture") or {}
|
|
inputs = architecture.get("input_modalities")
|
|
for modality, present in (("audio", audio_data is not None), ("image", has_images)):
|
|
if present and not isinstance(inputs, list):
|
|
raise ValueError(f"Cannot verify whether the selected chat model accepts {modality} input.")
|
|
if present and modality not in inputs:
|
|
raise ValueError(f"The selected chat model does not support {modality} input")
|
|
if image_controls and "image" not in architecture.get("output_modalities", []):
|
|
raise ValueError("The selected chat model does not advertise image output for these image settings.")
|
|
audio_content = None
|
|
if audio_data is not None:
|
|
from .openrouter_audio import prepare_audio
|
|
audio_content = prepare_audio(audio_data)["block"]
|
|
result = self._generate_chat(
|
|
key, system_prompt, user_message_box, model, web_search, cheapest, fastest,
|
|
temperature, pdf_engine, chat_mode, request_timeout=timeout,
|
|
aspect_ratio=aspect_ratio, image_resolution=image_resolution, seed=seed,
|
|
pdf_data=pdf_data, user_message_input=user_message_input,
|
|
reasoning_effort=reasoning_effort, service_tier=service_tier,
|
|
audio_content=audio_content, capabilities=capabilities, **kwargs)
|
|
return (*result, None)
|
|
|
|
# Recovery must not process or validate inputs from the original job.
|
|
resuming = request_type == "video" and bool(video_job_id.strip())
|
|
references = []
|
|
if not resuming:
|
|
if audio_data is not None or pdf_data is not None:
|
|
raise ValueError("Audio and PDF inputs are supported in chat mode only.")
|
|
for name in sorted((k for k in kwargs if re.fullmatch(r"image_\d+", k)),
|
|
key=lambda k: int(k.split("_")[1])):
|
|
if kwargs[name] is not None:
|
|
references.append("data:image/png;base64," + self.image_to_base64(kwargs[name]))
|
|
if request_type == "image":
|
|
from .openrouter_images import generate_images
|
|
from .openrouter_catalog import require_model
|
|
record = require_model("image", model)
|
|
resolution = image_resolution
|
|
# Older workflows always stored 1K even when the model did not
|
|
# support resolution. Do not turn that inherited default into a
|
|
# new unsupported /images parameter.
|
|
if resolution == "1K" and "resolution" not in record.get("supported_parameters", {}):
|
|
resolution = "auto"
|
|
result = generate_images(
|
|
key, model, prompt, reference_urls=references, aspect_ratio=ratio,
|
|
resolution=resolution, quality=image_quality, background=image_background,
|
|
seed=seed, timeout=timeout)
|
|
tensors = []
|
|
for raw in result["images"]:
|
|
with Image.open(io.BytesIO(raw)) as img:
|
|
mode = "RGBA" if "A" in img.getbands() or "transparency" in img.info else "RGB"
|
|
array = np.asarray(img.convert(mode), dtype=np.float32) / 255.0
|
|
tensors.append(torch.from_numpy(array.copy()).unsqueeze(0))
|
|
if not tensors:
|
|
raise ValueError("OpenRouter returned no raster images")
|
|
image = torch.cat(tensors, dim=0)
|
|
video = None
|
|
else:
|
|
from .openrouter_video import generate_video
|
|
self.last_video_job_id = video_job_id.strip()
|
|
result = generate_video(
|
|
key, model, prompt, reference_urls=references, mode=video_mode,
|
|
duration=video_duration, resolution=video_resolution, aspect_ratio=ratio,
|
|
generate_audio=video_generate_audio, seed=seed, request_timeout=timeout,
|
|
wait_timeout=video_wait_timeout, job_id=video_job_id,
|
|
on_job=lambda job: self._remember_video_job(job, unique_id))
|
|
image, video = placeholder, result["video"]
|
|
usage = result.get("usage") or {}
|
|
stats = f"Model: {model}"
|
|
cost = result.get("cost", usage.get("cost"))
|
|
if isinstance(cost, (int, float)):
|
|
stats += f", Cost: ${cost:.6f}"
|
|
if result.get("job_id"):
|
|
stats += f", Job: {result['job_id']}"
|
|
return (result.get("text", ""), image, stats, self.fetch_credits(key, timeout=timeout), video)
|
|
except Exception as error:
|
|
message = self._safe_error(error, key)
|
|
if request_type == "video" and self.last_video_job_id:
|
|
if f"video_job_id={self.last_video_job_id}" not in message:
|
|
message += f" | resume with video_job_id={self.last_video_job_id}"
|
|
if request_type == "video":
|
|
# A None VIDEO causes a misleading error in downstream SaveVideo
|
|
# and caches the failed generation as a successful node result.
|
|
failure = RuntimeError(f"OpenRouter video error: {message}")
|
|
failure.openrouter_job_id = self.last_video_job_id
|
|
raise failure from None
|
|
return (f"Error: {message}", placeholder, "Stats N/A due to error", "Credits N/A due to error", None)
|
|
|
|
def _generate_chat(self, api_key, system_prompt, user_message_box, model,
|
|
web_search, cheapest, fastest, temperature, pdf_engine, chat_mode,
|
|
request_timeout=120, aspect_ratio="auto", image_resolution="1K", seed=0,
|
|
pdf_data=None, user_message_input=None, reasoning_effort="auto",
|
|
service_tier="auto", audio_content=None, capabilities=None, **kwargs):
|
|
"""
|
|
Sends a completion request to the OpenRouter chat completion endpoint.
|
|
Handles text, optional image, and optional PDF inputs.
|
|
|
|
Returns four outputs:
|
|
(1) Output: the LLM's text response
|
|
(2) image: an image tensor if the response contains an image, else empty tensor
|
|
(3) Stats: a string with tokens per second, prompt tokens, completion tokens
|
|
(4) Credits: a string with the user's credit information
|
|
"""
|
|
# Create empty placeholder image
|
|
placeholder_image = torch.zeros((1, 1, 1, 3), dtype=torch.float32)
|
|
|
|
# Resolve API key
|
|
api_key = self.get_api_key(api_key)
|
|
|
|
if not api_key:
|
|
return ("Error: API Key not provided. Set LLM_KEY env var or use openrouter_api_key.json", placeholder_image, "Stats N/A", "Credits N/A")
|
|
|
|
url = "https://openrouter.ai/api/v1/chat/completions"
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
"HTTP-Referer": "https://github.com/gabe-init/ComfyUI-Openrouter_node",
|
|
"X-Title": "ComfyUI OpenRouter LLM Node",
|
|
}
|
|
|
|
# Validate and convert temperature
|
|
validated_temp = self.validate_temperature(temperature)
|
|
validated_timeout = self.validate_request_timeout(request_timeout)
|
|
validated_reasoning_effort = self.validate_reasoning_effort(reasoning_effort)
|
|
|
|
# Decide whether to use user_message_input or user_message_box
|
|
user_text = user_message_input if user_message_input is not None else user_message_box
|
|
|
|
# Initialize session_path
|
|
session_path = None
|
|
|
|
# Handle chat mode
|
|
if chat_mode:
|
|
# Get or create a chat session
|
|
session_path, messages = self.chat_manager.get_or_create_session(user_text, system_prompt)
|
|
|
|
# Check if we need to update the system prompt (for existing sessions)
|
|
if messages and messages[0]["role"] == "system" and messages[0]["content"] != system_prompt:
|
|
# Update system prompt if it has changed
|
|
messages[0]["content"] = system_prompt
|
|
else:
|
|
# Non-chat mode: Build the messages array, starting with a system prompt.
|
|
messages = [
|
|
{"role": "system", "content": system_prompt},
|
|
]
|
|
|
|
# --- Build the user message content ---
|
|
user_content_blocks = []
|
|
|
|
# 1. Add Text part (always present)
|
|
user_content_blocks.append({
|
|
"type": "text",
|
|
"text": user_text
|
|
})
|
|
if audio_content is not None:
|
|
user_content_blocks.append(audio_content)
|
|
|
|
# 2. Add Image parts (optional) - support multiple images from kwargs
|
|
# Process all image_N inputs from kwargs
|
|
image_keys = sorted([k for k in kwargs if re.fullmatch(r'image_\d+', k)],
|
|
key=lambda x: int(x.split('_')[1]))
|
|
|
|
for image_key in image_keys:
|
|
if kwargs[image_key] is not None:
|
|
try:
|
|
img_str = self.image_to_base64(kwargs[image_key])
|
|
user_content_blocks.append({
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/png;base64,{img_str}"
|
|
}
|
|
})
|
|
except Exception as e:
|
|
print(f"Error processing {image_key}: {e}")
|
|
return (f"Error processing {image_key}: {e}", placeholder_image, "Stats N/A", "Credits N/A")
|
|
|
|
# 3. Add PDF part (optional)
|
|
pdf_filename = "document.pdf" # Default filename if not provided
|
|
if pdf_data is not None:
|
|
# Validate pdf_data structure (expecting dict with 'filename' and 'bytes')
|
|
if isinstance(pdf_data, dict) and "bytes" in pdf_data and isinstance(pdf_data["bytes"], bytes):
|
|
pdf_bytes = pdf_data["bytes"]
|
|
# Use provided filename if available and valid, otherwise use default
|
|
if "filename" in pdf_data and isinstance(pdf_data["filename"], str) and pdf_data["filename"].strip():
|
|
pdf_filename = pdf_data["filename"]
|
|
|
|
try:
|
|
base64_pdf = base64.b64encode(pdf_bytes).decode('utf-8')
|
|
data_url = f"data:application/pdf;base64,{base64_pdf}"
|
|
user_content_blocks.append({
|
|
"type": "file",
|
|
"file": {
|
|
"filename": pdf_filename,
|
|
"file_data": data_url
|
|
}
|
|
})
|
|
except Exception as e:
|
|
print(f"Error encoding PDF: {e}")
|
|
return (f"Error encoding PDF: {e}", placeholder_image, "Stats N/A", "Credits N/A")
|
|
else:
|
|
# Handle case where pdf_data is not in the expected format
|
|
print(f"Warning: pdf_data input is not in the expected format (dict with 'filename' and 'bytes'). PDF not included.")
|
|
# Optionally return an error or just proceed without the PDF
|
|
# return ("Error: Invalid PDF data format.", "Stats N/A", "Credits N/A")
|
|
|
|
|
|
# Determine message format based on content type
|
|
# Use simple string format for text-only requests to ensure compatibility
|
|
# Use structured format only when we have multimodal content
|
|
has_multimodal_content = len(user_content_blocks) > 1 or any(block.get("type") != "text" for block in user_content_blocks)
|
|
|
|
if has_multimodal_content:
|
|
# Use structured format for multimodal content
|
|
new_user_message = {
|
|
"role": "user",
|
|
"content": user_content_blocks
|
|
}
|
|
else:
|
|
# Use simple string format for text-only requests
|
|
new_user_message = {
|
|
"role": "user",
|
|
"content": user_text
|
|
}
|
|
|
|
if chat_mode:
|
|
# In chat mode, append to existing conversation (but don't save yet - wait for response)
|
|
messages.append(new_user_message)
|
|
else:
|
|
# In non-chat mode, messages array already has system prompt, just append user message
|
|
messages.append(new_user_message)
|
|
|
|
# --- Apply model modifiers ---
|
|
modified_model = model
|
|
# Check if model already has modifiers to avoid duplication
|
|
if web_search and ":online" not in modified_model:
|
|
modified_model = f"{modified_model}:online"
|
|
if ":online" not in modified_model:
|
|
if cheapest and ":floor" not in modified_model:
|
|
modified_model = f"{modified_model}:floor"
|
|
elif fastest and not cheapest and ":nitro" not in modified_model:
|
|
modified_model = f"{modified_model}:nitro"
|
|
|
|
|
|
# --- Construct the final payload ---
|
|
data = {
|
|
"model": modified_model,
|
|
"messages": messages,
|
|
"temperature": validated_temp,
|
|
"seed": seed
|
|
}
|
|
if validated_reasoning_effort != "auto":
|
|
data["reasoning"] = {"effort": validated_reasoning_effort}
|
|
if service_tier != "auto":
|
|
data["service_tier"] = service_tier
|
|
|
|
capabilities = capabilities or {}
|
|
output_modalities = capabilities.get("architecture", {}).get("output_modalities", [])
|
|
if "image" in output_modalities:
|
|
data["modalities"] = [m for m in ("text", "image") if m in output_modalities]
|
|
config = self._chat_image_config(capabilities.get("id", model), aspect_ratio, image_resolution)
|
|
if config:
|
|
data["image_config"] = config
|
|
|
|
print(f"Payload: model={modified_model}")
|
|
|
|
# Add plugins if a specific PDF engine is selected
|
|
if pdf_engine != "auto":
|
|
data["plugins"] = [
|
|
{
|
|
"id": "file-parser",
|
|
"pdf": {
|
|
"engine": pdf_engine
|
|
}
|
|
}
|
|
]
|
|
|
|
# --- Pre-calculate text input tokens (rough estimate) ---
|
|
# Note: Actual token count depends on the model and includes parsed PDF/image data.
|
|
# Rely on the API response for accurate usage stats.
|
|
text_token_estimate = 0
|
|
try:
|
|
text_token_estimate = self.count_tokens(system_prompt, model) + self.count_tokens(user_text, model)
|
|
except Exception as e:
|
|
print(f"Warning: Token counting failed - {e}")
|
|
|
|
|
|
# --- Make API Call and Process Response ---
|
|
try:
|
|
start_time = time.time()
|
|
response = requests.post(url, headers=headers, json=data, timeout=validated_timeout)
|
|
response.raise_for_status() # Raises HTTPError for bad responses (4xx or 5xx)
|
|
end_time = time.time()
|
|
|
|
result = response.json()
|
|
|
|
# --- Extract results and calculate stats ---
|
|
if not result.get("choices") or not result["choices"][0].get("message"):
|
|
raise ValueError("Invalid response format from API: 'choices' or 'message' missing.")
|
|
|
|
# Parse response for text and image content
|
|
message = result["choices"][0]["message"]
|
|
text_output = message.get("content") or ""
|
|
image_tensor = placeholder_image
|
|
|
|
# Check for images in the separate images field (OpenRouter format)
|
|
if message.get("images"):
|
|
print(f"Found {len(message['images'])} image(s) in API response")
|
|
try:
|
|
# Get the first image from the images array
|
|
first_image = message["images"][0]
|
|
image_url = first_image["image_url"]["url"]
|
|
|
|
if image_url.startswith("data:image"):
|
|
base64_str = image_url.split(",", 1)[1]
|
|
try:
|
|
# Convert base64 to image tensor
|
|
image_tensor = self.base64_to_image(base64_str)
|
|
print(f"Successfully decoded image from API response")
|
|
except Exception as e:
|
|
print(f"Error decoding image: {e}")
|
|
else:
|
|
raise ValueError("Chat image response must contain encoded image data.")
|
|
except Exception as e:
|
|
print(f"Error processing images from response: {e}")
|
|
else:
|
|
print("No images found in API response - this may be normal if the model doesn't support image generation or the prompt didn't request an image")
|
|
|
|
# Also handle legacy multimodal content format as fallback
|
|
if isinstance(text_output, list):
|
|
text_parts = []
|
|
for content in text_output:
|
|
if isinstance(content, dict):
|
|
if content.get("type") == "text":
|
|
text_parts.append(content.get("text", ""))
|
|
elif content.get("type") == "image_url":
|
|
# Extract base64 image data
|
|
image_url = content["image_url"]["url"]
|
|
if image_url.startswith("data:image"):
|
|
base64_str = image_url.split(",", 1)[1]
|
|
try:
|
|
# Convert base64 to image tensor
|
|
image_tensor = self.base64_to_image(base64_str)
|
|
except Exception as e:
|
|
print(f"Error decoding image: {e}")
|
|
text_output = "\n".join(text_parts)
|
|
|
|
response_ms = result.get("response_ms", None)
|
|
api_usage = result.get("usage", {})
|
|
prompt_tokens = api_usage.get("prompt_tokens", text_token_estimate) # Use API value if available
|
|
completion_tokens = api_usage.get("completion_tokens", 0)
|
|
if completion_tokens == 0 and text_output: # Estimate completion tokens if API doesn't provide them
|
|
try:
|
|
completion_tokens = self.count_tokens(text_output, model)
|
|
except Exception as e:
|
|
print(f"Warning: Completion token counting failed - {e}")
|
|
|
|
|
|
# Calculate tokens per second (TPS)
|
|
tps = 0
|
|
elapsed_time = end_time - start_time
|
|
if response_ms is not None:
|
|
server_elapsed_time = response_ms / 1000.0
|
|
if server_elapsed_time > 0:
|
|
tps = completion_tokens / server_elapsed_time
|
|
elif elapsed_time > 0:
|
|
# Use client-side timing as fallback, less accurate due to network latency
|
|
tps = completion_tokens / elapsed_time
|
|
# Optional: apply a heuristic correction factor if needed, but server time is better
|
|
# correction_factor = 1.28 # Example factor, might need tuning
|
|
# tps *= correction_factor
|
|
|
|
stats_text = (
|
|
f"TPS: {tps:.2f}, "
|
|
f"Prompt Tokens: {prompt_tokens}, "
|
|
f"Completion Tokens: {completion_tokens}, "
|
|
f"Temp: {validated_temp:.1f}, "
|
|
f"Model: {modified_model}" # Display the actual model used
|
|
)
|
|
if pdf_engine != "auto":
|
|
stats_text += f", PDF Engine: {pdf_engine}"
|
|
if validated_reasoning_effort != "auto":
|
|
stats_text += f", Reasoning: {validated_reasoning_effort}"
|
|
if result.get("service_tier"):
|
|
stats_text += f", Service tier: {result['service_tier']}"
|
|
if isinstance(api_usage.get("cost"), (int, float)):
|
|
stats_text += f", Cost: ${api_usage['cost']:.6f}"
|
|
|
|
|
|
# Fetch credits information AFTER the main request
|
|
credits_text = self.fetch_credits(api_key, timeout=validated_timeout)
|
|
|
|
# Save conversation in chat mode
|
|
if chat_mode and session_path:
|
|
# Append assistant's response to the conversation
|
|
assistant_message = {
|
|
"role": "assistant",
|
|
"content": text_output
|
|
}
|
|
messages.append(assistant_message)
|
|
|
|
# Save the updated conversation
|
|
self.chat_manager.save_conversation(session_path, messages)
|
|
|
|
return (text_output, image_tensor, stats_text, credits_text)
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
error_message = f"API Request Error: {self._safe_error(e, api_key)}"
|
|
if hasattr(e, 'response') and e.response is not None:
|
|
error_message += f" | Status: {e.response.status_code}"
|
|
else:
|
|
error_message += " (Network or connection issue)"
|
|
print(f"ERROR: {error_message}")
|
|
return (error_message, placeholder_image, "Stats N/A due to error", "Credits N/A due to error")
|
|
except Exception as e:
|
|
print(f"ERROR: Node Error: {str(e)}")
|
|
return (f"Node Error: {str(e)}", placeholder_image, "Stats N/A due to error", "Credits N/A due to error")
|
|
|
|
@staticmethod
|
|
def image_to_base64(image):
|
|
"""
|
|
Converts a ComfyUI IMAGE (torch.Tensor, BHWC, float 0-1)
|
|
into a base64-encoded PNG string.
|
|
"""
|
|
if not isinstance(image, torch.Tensor):
|
|
raise TypeError("Input 'image' is not a torch.Tensor")
|
|
|
|
# Remove batch dimension if present
|
|
if image.ndim == 4:
|
|
if image.shape[0] != 1:
|
|
raise ValueError("Connect one image per numbered input; image batches are not supported.")
|
|
image = image.squeeze(0) # Shape HWC
|
|
|
|
if image.ndim != 3:
|
|
raise ValueError(f"Unexpected image dimensions: {image.shape}. Expected HWC.")
|
|
|
|
# Convert float tensor (0-1) to numpy array (0-255, uint8)
|
|
image_np = image.detach().cpu().float().numpy()
|
|
if not np.isfinite(image_np).all():
|
|
raise ValueError("Image contains non-finite values")
|
|
if image_np.dtype != np.uint8:
|
|
if image_np.min() < 0 or image_np.max() > 1:
|
|
print("Warning: Image tensor values outside [0, 1] range. Clamping.")
|
|
image_np = np.clip(image_np, 0, 1)
|
|
image_np = (image_np * 255).astype(np.uint8)
|
|
|
|
# Convert numpy array to PIL Image
|
|
if image_np.shape[-1] not in (3, 4):
|
|
raise ValueError("Image must contain RGB or RGBA channels")
|
|
pil_image = Image.fromarray(image_np)
|
|
|
|
# Save PIL Image to a bytes buffer as PNG
|
|
buffered = io.BytesIO()
|
|
pil_image.save(buffered, format="PNG")
|
|
|
|
# Encode the bytes buffer to base64 string
|
|
return base64.b64encode(buffered.getvalue()).decode('utf-8')
|
|
|
|
@staticmethod
|
|
def base64_to_image(base64_str: str) -> torch.Tensor:
|
|
"""
|
|
Converts a base64 image string to a ComfyUI image tensor
|
|
Returns tensor in [1, H, W, 3] format with values in [0, 1]
|
|
"""
|
|
try:
|
|
# Decode base64 string to image
|
|
img_data = base64.b64decode(base64_str)
|
|
img = Image.open(io.BytesIO(img_data))
|
|
img = img.convert("RGB")
|
|
|
|
# Convert to numpy array and normalize to [0, 1]
|
|
img_array = np.array(img).astype(np.float32) / 255.0
|
|
|
|
# Add batch dimension: [1, H, W, 3]
|
|
img_tensor = torch.from_numpy(img_array).unsqueeze(0)
|
|
|
|
print(f"Successfully converted base64 to image tensor: {img_tensor.shape}")
|
|
return img_tensor
|
|
|
|
except Exception as e:
|
|
print(f"Error in base64_to_image: {e}")
|
|
# Return a small placeholder image instead of failing
|
|
return torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
|
|
|
@staticmethod
|
|
def count_tokens(text, model):
|
|
"""
|
|
Count tokens for a given text using tiktoken.
|
|
Uses model-specific encodings where possible, falls back to cl100k_base.
|
|
Handles potential errors during encoding.
|
|
"""
|
|
if not text or not isinstance(text, str):
|
|
return 0
|
|
|
|
# Strip any model modifiers like :floor, :nitro, :online
|
|
base_model = model.split(':')[0] if ':' in model else model
|
|
|
|
# Simplified mapping, cl100k_base is common for many recent models
|
|
encoding_name = "cl100k_base"
|
|
try:
|
|
# List known models/prefixes that definitely use cl100k_base
|
|
# Add others if known, but cl100k_base is a safe default for many
|
|
cl100k_models = [
|
|
"openai/gpt-4", "openai/gpt-3.5", "openai/gpt-4o",
|
|
"anthropic/claude",
|
|
"google/gemini",
|
|
"meta-llama/llama-2", "meta-llama/llama-3",
|
|
"mistralai/mistral", "mistralai/mixtral",
|
|
]
|
|
# Check if the base_model or its prefix matches known cl100k models
|
|
is_cl100k = any(base_model.startswith(prefix) for prefix in cl100k_models)
|
|
|
|
if is_cl100k:
|
|
encoding_name = "cl100k_base"
|
|
# else: # Add logic for other encodings if needed, e.g., p50k_base for older models
|
|
# pass # Stick with cl100k_base as default for now
|
|
|
|
encoding = tiktoken.get_encoding(encoding_name)
|
|
token_count = len(encoding.encode(text, disallowed_special=())) # Allow special tokens
|
|
return token_count
|
|
|
|
except Exception as e:
|
|
print(f"Warning: Tiktoken error for model '{model}' (base: '{base_model}', encoding: '{encoding_name}'): {e}. Falling back to estimation.")
|
|
# Fallback: Estimate tokens based on characters (rough approximation)
|
|
# Average ~4 chars per token is a common heuristic
|
|
return max(1, round(len(text) / 4))
|
|
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, api_key, system_prompt, user_message_box, model,
|
|
web_search, cheapest, fastest, temperature, pdf_engine, chat_mode,
|
|
request_timeout=120, aspect_ratio="auto", image_resolution="1K", seed=0,
|
|
pdf_data=None, user_message_input=None, reasoning_effort="auto",
|
|
audio_data=None, request_type="chat", service_tier="auto",
|
|
image_quality="auto", image_background="auto", video_mode="text_to_video",
|
|
video_duration="auto", video_resolution="auto", video_generate_audio=False,
|
|
video_wait_timeout=900, video_job_id="", **kwargs):
|
|
"""
|
|
Check if any input that affects the output has changed.
|
|
Includes hashing image and PDF data.
|
|
"""
|
|
if request_type == "video" and video_job_id.strip():
|
|
# Recovery must work even if original media disappeared or changed.
|
|
return float("nan") # Poll again after a previous local timeout.
|
|
|
|
# Hash image data if present - handle multiple images from kwargs
|
|
image_hashes = []
|
|
image_keys = sorted([k for k in kwargs if re.fullmatch(r'image_\d+', k)],
|
|
key=lambda x: int(x.split('_')[1]))
|
|
|
|
for image_key in image_keys:
|
|
if kwargs[image_key] is not None:
|
|
image = kwargs[image_key]
|
|
if isinstance(image, torch.Tensor):
|
|
try:
|
|
hasher = hashlib.sha256()
|
|
normalized = image.detach().cpu().float().contiguous()
|
|
hasher.update(str(tuple(normalized.shape)).encode())
|
|
hasher.update(normalized.numpy().tobytes())
|
|
image_hashes.append(hasher.hexdigest())
|
|
except Exception as e:
|
|
print(f"Warning: Could not hash {image_key} data for IS_CHANGED: {e}")
|
|
raise ValueError(f"Could not hash {image_key}") from e
|
|
else:
|
|
image_hashes.append(None)
|
|
|
|
|
|
# Hash PDF data if present and valid
|
|
pdf_hash = None
|
|
if pdf_data is not None and isinstance(pdf_data, dict) and "bytes" in pdf_data and isinstance(pdf_data["bytes"], bytes):
|
|
try:
|
|
hasher = hashlib.sha256()
|
|
hasher.update(pdf_data["bytes"])
|
|
pdf_hash = hasher.hexdigest()
|
|
# Optionally include filename in hash if it affects processing?
|
|
# if "filename" in pdf_data: hasher.update(pdf_data["filename"].encode())
|
|
except Exception as e:
|
|
print(f"Warning: Could not hash pdf data for IS_CHANGED: {e}")
|
|
pdf_hash = "pdf_hashing_error" # Use a placeholder on error
|
|
elif pdf_data is not None:
|
|
# Handle cases where pdf_data is present but not in the expected format
|
|
pdf_hash = "invalid_pdf_data_format"
|
|
|
|
|
|
# Ensure temperature is consistently represented (e.g., as float)
|
|
try:
|
|
temp_float = float(temperature) if isinstance(temperature, (str, int, float)) else 1.0
|
|
temp_float = max(0.0, min(2.0, temp_float))
|
|
except (ValueError, TypeError):
|
|
temp_float = 1.0
|
|
|
|
try:
|
|
timeout_int = int(request_timeout)
|
|
timeout_int = max(cls.min_request_timeout, min(cls.max_request_timeout, timeout_int))
|
|
except (ValueError, TypeError):
|
|
timeout_int = cls.default_request_timeout
|
|
|
|
validated_reasoning_effort = cls.validate_reasoning_effort(reasoning_effort)
|
|
|
|
|
|
# Combine all relevant inputs into a tuple for comparison
|
|
# Use primitive types where possible for reliable hashing/comparison
|
|
# Note: api_key here is the UI value only. Keys resolved from the LLM_KEY
|
|
# env var or openrouter_api_key.json are intentionally NOT part of the
|
|
# cache key — they're treated as user environment, not workflow inputs.
|
|
audio_hash = None
|
|
if audio_data is not None:
|
|
from .openrouter_audio import audio_fingerprint
|
|
audio_hash = audio_fingerprint(audio_data)
|
|
return (hashlib.sha256(api_key.encode()).hexdigest(), system_prompt, user_message_box, model,
|
|
web_search, cheapest, fastest, temp_float, pdf_engine, chat_mode,
|
|
timeout_int, aspect_ratio, image_resolution, seed, validated_reasoning_effort,
|
|
tuple(image_hashes), pdf_hash, user_message_input, audio_hash, request_type, service_tier,
|
|
image_quality, image_background, video_mode, video_duration, video_resolution,
|
|
video_generate_audio, video_wait_timeout, video_job_id)
|
|
|
|
# Node class mappings
|
|
NODE_CLASS_MAPPINGS = {
|
|
"OpenRouterNode": OpenRouterNode,
|
|
"openrouter_node": OpenRouterNode
|
|
}
|
|
|
|
# Node display name mappings
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"OpenRouterNode": "OpenRouter LLM Node (Text/Multi-Image/PDF/Chat)" # Updated name
|
|
}
|