142 lines
5.6 KiB
Python
142 lines
5.6 KiB
Python
# api_key_input.py
|
|
import os
|
|
import sys
|
|
import logging
|
|
from typing import Any, Dict, Optional, Tuple
|
|
|
|
# Ensure parent directory is in path
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
parent_dir = os.path.dirname(current_dir)
|
|
if parent_dir not in sys.path:
|
|
sys.path.insert(0, parent_dir)
|
|
|
|
from context_payload import extract_context
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class APIKeyInput:
|
|
"""
|
|
A node to input API keys for various providers through context.
|
|
This allows dynamic API key management without hardcoding keys in provider nodes.
|
|
The API key is passed through context and can override provider-specific keys.
|
|
"""
|
|
|
|
SUPPORTED_PROVIDERS = [
|
|
"openai", "gemini", "google", "anthropic", "groq", "huggingface",
|
|
"bfl", "wavespeed", "ollama", "lmstudio", "textgen", "kobold",
|
|
"llamacpp", "vllm", "transformers", "custom"
|
|
]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"provider": (cls.SUPPORTED_PROVIDERS, {
|
|
"default": "gemini",
|
|
"tooltip": "Select the provider this API key is for"
|
|
}),
|
|
"api_key": ("STRING", {
|
|
"multiline": False,
|
|
"default": "",
|
|
"tooltip": "Enter the API key for the selected provider"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"context": ("*", {"tooltip": "Optional context to merge with"}),
|
|
"override_existing": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Whether to override existing API keys in context/provider config"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("*",)
|
|
RETURN_NAMES = ("context",)
|
|
FUNCTION = "set_api_key"
|
|
CATEGORY = "llm_toolkit/config"
|
|
|
|
def set_api_key(
|
|
self,
|
|
provider: str,
|
|
api_key: str,
|
|
context: Optional[Dict[str, Any]] = None,
|
|
override_existing: bool = False
|
|
) -> Tuple[Dict[str, Any]]:
|
|
"""
|
|
Sets an API key for a specific provider in the context.
|
|
|
|
Args:
|
|
provider: The provider name (e.g., 'gemini', 'openai')
|
|
api_key: The API key string
|
|
context: Existing context to merge with
|
|
override_existing: Whether to override existing API keys
|
|
|
|
Returns:
|
|
Updated context with API key information
|
|
"""
|
|
logger.info(f"APIKeyInput: Setting API key for provider '{provider}'")
|
|
|
|
# Validate inputs
|
|
if not api_key or not api_key.strip():
|
|
logger.warning(f"APIKeyInput: Empty API key provided for provider '{provider}'")
|
|
|
|
api_key = api_key.strip()
|
|
|
|
# Initialize or copy context
|
|
if context is None:
|
|
output_context = {}
|
|
elif isinstance(context, dict):
|
|
output_context = context.copy()
|
|
else:
|
|
# Handle ContextPayload or other formats
|
|
unwrapped = extract_context(context)
|
|
if isinstance(unwrapped, dict):
|
|
output_context = unwrapped.copy()
|
|
output_context.setdefault("passthrough_data", context)
|
|
else:
|
|
output_context = {"passthrough_data": context}
|
|
|
|
# Initialize api_keys section in context
|
|
if "api_keys" not in output_context:
|
|
output_context["api_keys"] = {}
|
|
|
|
# Check if we should set the key
|
|
existing_key = output_context["api_keys"].get(provider)
|
|
|
|
if existing_key and not override_existing:
|
|
# Secure logging: only show first 5 characters
|
|
masked_existing = existing_key[:5] + "..." if len(existing_key) > 5 else "..."
|
|
masked_new = api_key[:5] + "..." if len(api_key) > 5 else "..."
|
|
logger.info(f"APIKeyInput: Provider '{provider}' already has API key ({masked_existing}), keeping existing (override=False)")
|
|
else:
|
|
# Set the new key
|
|
output_context["api_keys"][provider] = api_key
|
|
|
|
# Secure logging: only show first 5 characters
|
|
masked_key = api_key[:5] + "..." if len(api_key) > 5 else "..."
|
|
action = "overriding" if existing_key else "setting"
|
|
logger.info(f"APIKeyInput: {action} API key for provider '{provider}' ({masked_key})")
|
|
|
|
# For backward compatibility, also update provider_config if it exists and matches
|
|
provider_config = output_context.get("provider_config")
|
|
if provider_config and isinstance(provider_config, dict):
|
|
config_provider = provider_config.get("provider_name", "").lower()
|
|
|
|
# Check if this API key is for the current provider config
|
|
if config_provider == provider.lower() or (provider == "google" and config_provider == "gemini"):
|
|
if override_existing or not provider_config.get("api_key") or provider_config.get("api_key") == "1234":
|
|
provider_config["api_key"] = api_key
|
|
logger.info(f"APIKeyInput: Updated provider_config API key for '{config_provider}'")
|
|
|
|
logger.info(f"APIKeyInput: Context now contains API keys for: {list(output_context['api_keys'].keys())}")
|
|
|
|
return (output_context,)
|
|
|
|
|
|
# --- Node Mappings ---
|
|
NODE_CLASS_MAPPINGS = {
|
|
"APIKeyInput": APIKeyInput
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"APIKeyInput": "API Key Input (LLMToolkit)"
|
|
} |