diff --git a/README.md b/README.md index a2d2381..0da565b 100644 --- a/README.md +++ b/README.md @@ -94,9 +94,14 @@ To keep your API key secure, use one of the following methods: - **model**: The model to use for generation. The node automatically fetches the list of available models from OpenRouter. - **web_search**: Enable web search capability by appending `:online` to the model ID. This costs $4 per 1000 queries and automatically uses your openrouter balance. - **cheapest**: Route to the cheapest provider by appending `:floor` to the model ID (enabled by default). -- **fastest**: Route to the fastest provider by appending `:nitro` to the model ID (disabled by default). -- **temperature**: Controls the randomness of the model's output (0.0 to 2.0). -- **chat_mode**: Enable conversation mode to maintain context across messages (disabled by default). +- **fastest**: Route to the fastest provider by appending `:nitro` to the model ID (disabled by default). +- **temperature**: Controls the randomness of the model's output (0.0 to 2.0). +- **reasoning_effort**: Controls OpenRouter's unified reasoning parameter. `auto` sends no override and lets OpenRouter/model defaults apply. `none` explicitly disables reasoning. `minimal`, `low`, `medium`, `high`, and `xhigh` request increasing reasoning effort where the selected model supports it. +- **chat_mode**: Enable conversation mode to maintain context across messages (disabled by default). +- **request_timeout**: Maximum time, in seconds, to wait for the main OpenRouter completion request before returning an error (default: 120). + +> [!NOTE] +> OpenRouter normalizes reasoning across providers, but behavior still depends on the selected model. Some providers map effort levels to token budgets, Gemini 3 maps effort to Google's thinking levels, `xhigh` may be mapped down when unsupported, and some reasoning models do not return visible reasoning tokens. #### Optional Inputs: diff --git a/node.py b/node.py index cfa51fa..19494ff 100644 --- a/node.py +++ b/node.py @@ -26,9 +26,14 @@ class OpenRouterNode: 3) "Credits": a string showing your remaining OpenRouter account balance """ - models_cache = None - last_fetch_time = 0 - cache_duration = 3600 # Cache duration in seconds (1 hour) + 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() @@ -105,6 +110,7 @@ class OpenRouterNode: "8:1 (google/gemini-3.1-flash-image-preview (Nano Banana 2) only)", ], {"default": "auto"}), "image_resolution": (["1K", "2K", "4K"], {"default": "1K"}), + "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, @@ -115,11 +121,18 @@ class OpenRouterNode: "round": 0.01, }), "pdf_engine": (["auto", "mistral-ocr", "pdf-text"], {"default": "auto"}), - "chat_mode": ("BOOLEAN", {"default": False}), - }, - "optional": { - "pdf_data": (PDF_DATA_TYPE,), # Use '*' and check structure in generate_response - "user_message_input": ("STRING", {"forceInput": True}), + "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}), } } @@ -138,7 +151,7 @@ class OpenRouterNode: if (cls.models_cache is None) or (current_time - cls.last_fetch_time > cls.cache_duration): url = "https://openrouter.ai/api/v1/models" try: - response = requests.get(url) + response = requests.get(url, timeout=cls.default_request_timeout) response.raise_for_status() models = response.json()["data"] # Filter for models that support chat completions if needed, but API handles this @@ -152,17 +165,39 @@ class OpenRouterNode: cls.models_cache = ["error_fetching_models", "google/gemma-3-27b-it", "openai/gpt-4o"] # Example fallbacks return cls.models_cache if cls.models_cache else ["error_fetching_models"] # Ensure it's never empty - def validate_temperature(self, temperature): - """ - Validates and converts temperature value to float within acceptable range. - """ - try: + 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 fetch_credits(self, api_key): + 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. @@ -179,8 +214,9 @@ class OpenRouterNode: "X-Title": "ComfyUI OpenRouter LLM Node", } - try: - response = requests.get(url, headers=headers) + try: + validated_timeout = self.validate_request_timeout(timeout) + response = requests.get(url, headers=headers, timeout=validated_timeout) response.raise_for_status() result = response.json() @@ -204,10 +240,10 @@ class OpenRouterNode: except json.JSONDecodeError: return "Error fetching credits: Could not decode JSON response." - def generate_response(self, api_key, system_prompt, user_message_box, model, - web_search, cheapest, fastest, temperature, pdf_engine, chat_mode, - aspect_ratio="auto", image_resolution="1K", seed=0, - pdf_data=None, user_message_input=None, **kwargs): + 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", **kwargs): """ Sends a completion request to the OpenRouter chat completion endpoint. Handles text, optional image, and optional PDF inputs. @@ -235,8 +271,10 @@ class OpenRouterNode: "X-Title": "ComfyUI OpenRouter LLM Node", } - # Validate and convert temperature - validated_temp = self.validate_temperature(temperature) + # 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 and user_message_input.strip() else user_message_box @@ -355,12 +393,14 @@ class OpenRouterNode: # --- Construct the final payload --- - data = { - "model": modified_model, - "messages": messages, - "temperature": validated_temp, - "seed": seed - } + data = { + "model": modified_model, + "messages": messages, + "temperature": validated_temp, + "seed": seed + } + if validated_reasoning_effort != "auto": + data["reasoning"] = {"effort": validated_reasoning_effort} print(f"Payload: model={modified_model}") @@ -387,10 +427,10 @@ class OpenRouterNode: # --- Make API Call and Process Response --- try: - start_time = time.time() - response = requests.post(url, headers=headers, json=data) - response.raise_for_status() # Raises HTTPError for bad responses (4xx or 5xx) - end_time = time.time() + 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() # Debug: print truncated response to see what OpenRouter returned @@ -480,12 +520,14 @@ class OpenRouterNode: 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 pdf_engine != "auto": + stats_text += f", PDF Engine: {pdf_engine}" + if validated_reasoning_effort != "auto": + stats_text += f", Reasoning: {validated_reasoning_effort}" # Fetch credits information AFTER the main request - credits_text = self.fetch_credits(api_key) + credits_text = self.fetch_credits(api_key, timeout=validated_timeout) # Save conversation in chat mode if chat_mode and session_path: @@ -623,11 +665,11 @@ class OpenRouterNode: 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, - aspect_ratio="auto", image_resolution="1K", seed=0, - pdf_data=None, user_message_input=None, **kwargs): + @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", **kwargs): """ Check if any input that affects the output has changed. Includes hashing image and PDF data. @@ -673,8 +715,16 @@ class OpenRouterNode: 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 + 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 @@ -682,9 +732,10 @@ class OpenRouterNode: # 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. - return (api_key, system_prompt, user_message_box, model, - web_search, cheapest, fastest, temp_float, pdf_engine, chat_mode, - aspect_ratio, image_resolution, seed, tuple(image_hashes), pdf_hash, user_message_input) + return (api_key, 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) # Node class mappings NODE_CLASS_MAPPINGS = { diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_request_timeout.py b/tests/test_request_timeout.py new file mode 100644 index 0000000..f0160f0 --- /dev/null +++ b/tests/test_request_timeout.py @@ -0,0 +1,204 @@ +import importlib.util +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + + +class FakeRequestException(Exception): + def __init__(self, message="", response=None): + super().__init__(message) + self.response = response + + +class FakeTimeout(FakeRequestException): + pass + + +def install_dependency_stubs(): + requests_module = types.ModuleType("requests") + exceptions = types.SimpleNamespace( + RequestException=FakeRequestException, + Timeout=FakeTimeout, + ) + requests_module.exceptions = exceptions + requests_module.get = Mock() + requests_module.post = Mock() + sys.modules["requests"] = requests_module + + torch_module = types.ModuleType("torch") + torch_module.float32 = "float32" + torch_module.Tensor = type("Tensor", (), {}) + torch_module.zeros = Mock(return_value="placeholder-image") + sys.modules["torch"] = torch_module + + tiktoken_module = types.ModuleType("tiktoken") + tiktoken_module.get_encoding = Mock() + sys.modules["tiktoken"] = tiktoken_module + + pil_module = types.ModuleType("PIL") + image_module = types.ModuleType("PIL.Image") + image_module.open = Mock() + image_module.fromarray = Mock() + pil_module.Image = image_module + sys.modules["PIL"] = pil_module + sys.modules["PIL.Image"] = image_module + + sys.modules.setdefault("numpy", types.ModuleType("numpy")) + + +def load_node_module(): + install_dependency_stubs() + root = Path(__file__).resolve().parents[1] + package_name = "comfy_openrouter_node_test" + + package = types.ModuleType(package_name) + package.__path__ = [str(root)] + sys.modules[package_name] = package + + spec = importlib.util.spec_from_file_location( + f"{package_name}.node", + root / "node.py", + submodule_search_locations=[str(root)], + ) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +class RequestTimeoutTests(unittest.TestCase): + def setUp(self): + self.node_module = load_node_module() + self.node = self.node_module.OpenRouterNode() + self.node.fetch_credits = Mock(return_value="Remaining: $1.000") + self.node.count_tokens = Mock(return_value=1) + + def call_generate_response(self, request_timeout=120, reasoning_effort="auto"): + return self.node.generate_response( + api_key="test-key", + system_prompt="system", + user_message_box="hello", + model="openai/gpt-4o", + web_search=False, + cheapest=False, + fastest=False, + temperature=1.0, + pdf_engine="auto", + chat_mode=False, + request_timeout=request_timeout, + reasoning_effort=reasoning_effort, + ) + + def test_main_openrouter_request_uses_configured_timeout(self): + response = Mock() + response.raise_for_status = Mock() + response.json.return_value = { + "choices": [{"message": {"content": "done"}}], + "usage": {"prompt_tokens": 2, "completion_tokens": 3}, + } + self.node_module.requests.post.return_value = response + + result = self.call_generate_response(request_timeout=45) + + self.assertEqual(result[0], "done") + self.node_module.requests.post.assert_called_once() + self.assertEqual(self.node_module.requests.post.call_args.kwargs["timeout"], 45) + + def test_default_reasoning_effort_omits_reasoning_override(self): + response = Mock() + response.raise_for_status = Mock() + response.json.return_value = { + "choices": [{"message": {"content": "done"}}], + "usage": {"prompt_tokens": 2, "completion_tokens": 3}, + } + self.node_module.requests.post.return_value = response + + result = self.call_generate_response() + + self.assertEqual(result[0], "done") + payload = self.node_module.requests.post.call_args.kwargs["json"] + self.assertNotIn("reasoning", payload) + + def test_explicit_reasoning_effort_is_sent(self): + response = Mock() + response.raise_for_status = Mock() + response.json.return_value = { + "choices": [{"message": {"content": "done"}}], + "usage": {"prompt_tokens": 2, "completion_tokens": 3}, + } + self.node_module.requests.post.return_value = response + + result = self.call_generate_response(reasoning_effort="high") + + self.assertEqual(result[0], "done") + payload = self.node_module.requests.post.call_args.kwargs["json"] + self.assertEqual(payload["reasoning"], {"effort": "high"}) + + def test_invalid_reasoning_effort_falls_back_to_auto(self): + response = Mock() + response.raise_for_status = Mock() + response.json.return_value = { + "choices": [{"message": {"content": "done"}}], + "usage": {"prompt_tokens": 2, "completion_tokens": 3}, + } + self.node_module.requests.post.return_value = response + + result = self.call_generate_response(reasoning_effort="unsupported") + + self.assertEqual(result[0], "done") + payload = self.node_module.requests.post.call_args.kwargs["json"] + self.assertNotIn("reasoning", payload) + + def test_is_changed_includes_reasoning_effort(self): + base_args = dict( + api_key="test-key", + system_prompt="system", + user_message_box="hello", + model="openai/gpt-4o", + web_search=False, + cheapest=False, + fastest=False, + temperature=1.0, + pdf_engine="auto", + chat_mode=False, + request_timeout=120, + ) + + low_key = self.node_module.OpenRouterNode.IS_CHANGED(**base_args, reasoning_effort="low") + high_key = self.node_module.OpenRouterNode.IS_CHANGED(**base_args, reasoning_effort="high") + + self.assertNotEqual(low_key, high_key) + + def test_timeout_exception_returns_clear_error(self): + self.node_module.requests.post.side_effect = FakeTimeout("request timed out") + + result = self.call_generate_response(request_timeout=2) + + self.assertIn("API Request Error", result[0]) + self.assertIn("request timed out", result[0]) + self.assertEqual(result[2], "Stats N/A due to error") + self.assertEqual(result[3], "Credits N/A due to error") + + def test_fetch_credits_uses_configured_timeout(self): + self.node.fetch_credits = self.node_module.OpenRouterNode().fetch_credits + response = Mock() + response.raise_for_status = Mock() + response.json.return_value = { + "data": { + "total_credits": 10.0, + "total_usage": 3.25, + } + } + self.node_module.requests.get.return_value = response + + credits = self.node.fetch_credits("test-key", timeout=45) + + self.assertEqual(credits, "Remaining: $6.750") + self.node_module.requests.get.assert_called_once() + self.assertEqual(self.node_module.requests.get.call_args.kwargs["timeout"], 45) + + +if __name__ == "__main__": + unittest.main()