diff --git a/README.md b/README.md index a2d2381..30f4321 100644 --- a/README.md +++ b/README.md @@ -95,8 +95,9 @@ To keep your API key secure, use one of the following methods: - **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). +- **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). +- **request_timeout**: Maximum time, in seconds, to wait for the main OpenRouter completion request before returning an error (default: 120). #### Optional Inputs: diff --git a/node.py b/node.py index 51c91f8..dc5d3e4 100644 --- a/node.py +++ b/node.py @@ -26,9 +26,12 @@ 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 def __init__(self): self.chat_manager = ChatSessionManager() @@ -116,11 +119,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}), } } @@ -139,7 +149,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 @@ -153,17 +163,27 @@ 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 + + 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. @@ -180,8 +200,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() @@ -205,10 +226,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, reasoning_effort="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="none", **kwargs): """ Sends a completion request to the OpenRouter chat completion endpoint. Handles text, optional image, and optional PDF inputs. @@ -236,8 +257,9 @@ 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) # 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 @@ -389,10 +411,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 @@ -487,7 +509,7 @@ class OpenRouterNode: # 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: @@ -625,11 +647,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, **kwargs): """ Check if any input that affects the output has changed. Includes hashing image and PDF data. @@ -675,8 +697,14 @@ 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 # Combine all relevant inputs into a tuple for comparison @@ -684,9 +712,9 @@ 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, 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..7f1231a --- /dev/null +++ b/tests/test_request_timeout.py @@ -0,0 +1,138 @@ +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): + 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, + ) + + 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_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()