feat: add secure API key handling via environment variables and JSON config
This commit is contained in:
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"api_key": "your_openrouter_api_key_here"
|
||||
}
|
||||
Reference in New Issue
Block a user