diff --git a/README.md b/README.md index 383b71a..d8a0fd5 100644 --- a/README.md +++ b/README.md @@ -56,11 +56,27 @@ pip install -r requirements.txt The OpenRouter node provides a simple interface to interact with various LLMs through the OpenRouter API. +### API Key Security + +> [!WARNING] +> Entering your API key directly into the node's input field in ComfyUI will embed it in the **workflow metadata** of every image you save. This is a security risk if you share your images. + +To keep your API key secure, use one of the following methods: + +1. **JSON Config File (Recommended)**: + - Create a file named `openrouter_api_key.json` in the node directory. + - Add your key like this: `{"api_key": "your_key_here"}`. + - Leave the **api_key** field blank in the ComfyUI interface. + +2. **Environment Variable**: + - Set an environment variable named `LLM_KEY` on your system. + - The node will automatically pick it up if the UI field is empty. + ### Inputs #### Required Inputs: -- **api_key**: Your OpenRouter API key. You can get one from [OpenRouter](https://openrouter.ai/). +- **api_key**: Your OpenRouter API key. Can be left blank if provided via `openrouter_api_key.json` or `LLM_KEY` environment variable. - **system_prompt**: The system prompt that sets the behavior of the LLM. - **user_message_box**: The user message to send to the LLM. - **model**: The model to use for generation. The node automatically fetches the list of available models from OpenRouter. diff --git a/node.py b/node.py index 75b4d56..c7f4a09 100644 --- a/node.py +++ b/node.py @@ -8,6 +8,7 @@ import torch import tiktoken from PIL import Image import hashlib # Added for hashing PDF bytes in IS_CHANGED +import os from .chat_manager import ChatSessionManager # Define a placeholder type name for PDF data. @@ -32,6 +33,35 @@ class OpenRouterNode: def __init__(self): self.chat_manager = ChatSessionManager() + def get_api_key(self, 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): """ @@ -136,6 +166,7 @@ class OpenRouterNode: 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." @@ -188,8 +219,12 @@ class OpenRouterNode: """ # 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.", placeholder_image, "Stats N/A", "Credits N/A") + 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 = { @@ -643,6 +678,8 @@ class OpenRouterNode: # Combine all relevant inputs into a tuple for comparison # Use primitive types where possible for reliable hashing/comparison + # Note: We don't hash the resolved API key from file/env to avoid re-running + # if only the hidden key changes (which is rare), but we still pass the UI value. 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) diff --git a/openrouter_api_key.json.example b/openrouter_api_key.json.example new file mode 100644 index 0000000..3124938 --- /dev/null +++ b/openrouter_api_key.json.example @@ -0,0 +1,3 @@ +{ + "api_key": "your_openrouter_api_key_here" +}