v1.16.0 - Added PromptRefinerNode with modular AI provider support
- Introduced PromptRefinerNode for post-processing prompts - Added multi-provider support (Ollama, OpenRouter, llama.cpp, etc.) - Implemented external API key loading via api_keys.txt (aligned with PromptCreatorNode) - Added configurable host input (full URL support for remote/local inference) - Integrated output cleaning system to remove thinking/reasoning artifacts - Added robust OpenRouter handling with fallback parsing and debug logging - Added console debug output for final refined prompt - Ensured compatibility with Nemotron and mixed pipelines (local + cloud) Improves prompt quality, modularity, and pipeline flexibility
This commit is contained in:
zeeoale
@@ -0,0 +1,275 @@
|
||||
import requests
|
||||
import os
|
||||
|
||||
|
||||
# =========================
|
||||
# 🧠 CLEAN OUTPUT RULES
|
||||
# =========================
|
||||
CLEAN_OUTPUT_PROMPT = """
|
||||
You must return ONLY the final refined prompt.
|
||||
|
||||
Do NOT include:
|
||||
- explanations
|
||||
- reasoning
|
||||
- thinking
|
||||
- analysis
|
||||
- markdown formatting
|
||||
- headings
|
||||
- bullet points
|
||||
|
||||
Do not produce hidden reasoning or thinking tokens.
|
||||
Do not include <think> or similar tags.
|
||||
|
||||
Output must be a single continuous paragraph.
|
||||
|
||||
No extra text before or after the prompt.
|
||||
|
||||
The output must be ready to use directly.
|
||||
"""
|
||||
|
||||
|
||||
class PromptRefinerNode:
|
||||
|
||||
def __init__(self):
|
||||
self.base_path = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
# =========================
|
||||
# 🔑 API KEYS (COME PCN)
|
||||
# =========================
|
||||
def _read_api_keys(self, base_path):
|
||||
keys_path = os.path.join(base_path, "api_keys.txt")
|
||||
keys = {}
|
||||
|
||||
if os.path.exists(keys_path):
|
||||
with open(keys_path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
if "=" in line:
|
||||
k, v = line.strip().split("=", 1)
|
||||
keys[k.strip()] = v.strip()
|
||||
else:
|
||||
print(f"[PromptRefiner] api_keys.txt NOT FOUND at: {keys_path}")
|
||||
|
||||
return keys
|
||||
|
||||
# =========================
|
||||
# 🧩 INPUTS
|
||||
# =========================
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"prompt_in": ("STRING", {"multiline": True}),
|
||||
"enable_refine": ("BOOLEAN", {"default": True}),
|
||||
"provider": ([
|
||||
"ollama",
|
||||
"openrouter",
|
||||
"openai",
|
||||
"gemini",
|
||||
"cohere",
|
||||
"llama_cpp"
|
||||
],),
|
||||
"refinement_mode": ([
|
||||
"subtle",
|
||||
"balanced",
|
||||
"creative"
|
||||
],),
|
||||
},
|
||||
"optional": {
|
||||
"model": ("STRING", {"default": ""}),
|
||||
"system_prompt": ("STRING", {"multiline": True, "default": ""}),
|
||||
"host": ("STRING", {"default": "http://127.0.0.1:11434"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING")
|
||||
RETURN_NAMES = ("prompt_out", "prompt_raw")
|
||||
FUNCTION = "refine"
|
||||
CATEGORY = "PFN / Prompt"
|
||||
|
||||
# =========================
|
||||
# 🚀 MAIN
|
||||
# =========================
|
||||
def refine(self, prompt_in, enable_refine, provider, refinement_mode, model="", system_prompt="", host="http://127.0.0.1:11434"):
|
||||
|
||||
if not enable_refine:
|
||||
return (prompt_in, prompt_in)
|
||||
|
||||
base_prompt = system_prompt.strip() if system_prompt.strip() else self.get_default_system_prompt(refinement_mode)
|
||||
system_prompt_full = base_prompt + "\n\n" + CLEAN_OUTPUT_PROMPT
|
||||
|
||||
try:
|
||||
raw_response = self.call_provider(
|
||||
provider,
|
||||
model,
|
||||
system_prompt_full,
|
||||
prompt_in,
|
||||
host
|
||||
)
|
||||
|
||||
refined_prompt = self.extract_text(raw_response)
|
||||
|
||||
# 🔥 DEBUG
|
||||
print("\n[PromptRefiner] FINAL PROMPT:\n")
|
||||
print(refined_prompt)
|
||||
print("\n" + "="*60 + "\n")
|
||||
|
||||
except Exception as e:
|
||||
refined_prompt = prompt_in
|
||||
raw_response = f"ERROR: {str(e)}"
|
||||
print(f"[PromptRefiner ERROR] {e}")
|
||||
|
||||
return (refined_prompt, raw_response)
|
||||
|
||||
# =========================
|
||||
# 🧠 SYSTEM MODES
|
||||
# =========================
|
||||
def get_default_system_prompt(self, mode):
|
||||
|
||||
if mode == "subtle":
|
||||
return "Refine the prompt preserving structure and improving clarity."
|
||||
|
||||
elif mode == "balanced":
|
||||
return "Refine the prompt improving clarity, flow and visual richness while preserving identity."
|
||||
|
||||
elif mode == "creative":
|
||||
return "Enhance the prompt with richer detail, atmosphere and depth while keeping the original concept intact."
|
||||
|
||||
return ""
|
||||
|
||||
# =========================
|
||||
# 🔌 ROUTER
|
||||
# =========================
|
||||
def call_provider(self, provider, model, system_prompt, user_prompt, host):
|
||||
|
||||
if provider == "ollama":
|
||||
return self.call_ollama(model, system_prompt, user_prompt, host)
|
||||
|
||||
elif provider == "llama_cpp":
|
||||
return self.call_llama_cpp(system_prompt, user_prompt, host)
|
||||
|
||||
elif provider == "openrouter":
|
||||
return self._enhance_with_openrouter(self.base_path, model, system_prompt, user_prompt)
|
||||
|
||||
elif provider == "openai":
|
||||
return f"[OpenAI NOT IMPLEMENTED]\n{user_prompt}"
|
||||
|
||||
elif provider == "gemini":
|
||||
return f"[Gemini NOT IMPLEMENTED]\n{user_prompt}"
|
||||
|
||||
elif provider == "cohere":
|
||||
return f"[Cohere NOT IMPLEMENTED]\n{user_prompt}"
|
||||
|
||||
else:
|
||||
raise ValueError("Unsupported provider")
|
||||
|
||||
# =========================
|
||||
# 🔌 OLLAMA
|
||||
# =========================
|
||||
def call_ollama(self, model, system_prompt, user_prompt, host):
|
||||
|
||||
url = f"{host}/api/generate"
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"prompt": f"{system_prompt}\n\n{user_prompt}",
|
||||
"stream": False
|
||||
}
|
||||
|
||||
r = requests.post(url, json=payload)
|
||||
return r.json().get("response", "")
|
||||
|
||||
# =========================
|
||||
# 🔌 LLAMA CPP
|
||||
# =========================
|
||||
def call_llama_cpp(self, system_prompt, user_prompt, host):
|
||||
|
||||
url = f"{host}/completion"
|
||||
|
||||
payload = {
|
||||
"prompt": f"{system_prompt}\n\n{user_prompt}",
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 512
|
||||
}
|
||||
|
||||
r = requests.post(url, json=payload)
|
||||
return r.json().get("content", "")
|
||||
|
||||
# =========================
|
||||
# 🌐 OPENROUTER (PCN STYLE + SAFE)
|
||||
# =========================
|
||||
def _enhance_with_openrouter(self, base_path, model, system_prompt, user_prompt):
|
||||
|
||||
keys = self._read_api_keys(base_path)
|
||||
api_key = keys.get("openrouter", "").strip()
|
||||
|
||||
if not api_key:
|
||||
print("[PromptRefiner] ❌ No OpenRouter API key found.")
|
||||
return user_prompt
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt}
|
||||
]
|
||||
}
|
||||
|
||||
try:
|
||||
r = requests.post(
|
||||
"https://openrouter.ai/api/v1/chat/completions",
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=120
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
|
||||
data = r.json()
|
||||
|
||||
print("\n[OpenRouter RAW RESPONSE]")
|
||||
print(data)
|
||||
print("="*60)
|
||||
|
||||
# ✅ standard
|
||||
if "choices" in data:
|
||||
return data["choices"][0]["message"]["content"].strip()
|
||||
|
||||
# ❗ fallback
|
||||
if "output" in data:
|
||||
return data["output"]
|
||||
|
||||
if "content" in data:
|
||||
return data["content"]
|
||||
|
||||
return str(data)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[PromptRefiner OpenRouter ERROR] {e}")
|
||||
return user_prompt
|
||||
|
||||
# =========================
|
||||
# 🧹 CLEAN OUTPUT
|
||||
# =========================
|
||||
def extract_text(self, response):
|
||||
|
||||
if not isinstance(response, str):
|
||||
return str(response)
|
||||
|
||||
text = response.strip()
|
||||
|
||||
if "<think>" in text:
|
||||
text = text.split("</think>")[-1]
|
||||
|
||||
text = text.replace("<think>", "").replace("</think>", "")
|
||||
|
||||
if "Prompt:" in text:
|
||||
text = text.split("Prompt:")[-1]
|
||||
|
||||
text = " ".join(text.split())
|
||||
|
||||
return text.strip()
|
||||
+13
-5
@@ -4,7 +4,7 @@
|
||||
|
||||
MANIFEST = {
|
||||
"name": "Prompt Creator Node",
|
||||
"version": (1, 13, 3),
|
||||
"version": (1, 16, 0),
|
||||
"author": "TK-traumakom",
|
||||
"project": "https://github.com/zeeoale/PromptCreatorNode",
|
||||
"description": (
|
||||
@@ -39,11 +39,11 @@ def _print_promptcreator_banner():
|
||||
██ ██ ██ ██ ██ ██ ██ ██
|
||||
██ ████ ██████ ██████ ███████
|
||||
-----------------------------------------------------
|
||||
1.13.3
|
||||
1.16.0
|
||||
-----------------------------------------------------
|
||||
|
||||
===========================================================
|
||||
PromptCreatorNode v1.13.3
|
||||
PromptCreatorNode v1.16.0
|
||||
A modular cinematic prompt generator for ComfyUI
|
||||
|
||||
Loaded nodes:
|
||||
@@ -52,6 +52,7 @@ Loaded nodes:
|
||||
- Prompt Replay
|
||||
- Prompt Builder
|
||||
- Promt Tags Extractor Node (create tags for galley uploader)
|
||||
- Prompt Refiner Node
|
||||
|
||||
Author: traumakom
|
||||
===========================================================
|
||||
@@ -69,6 +70,9 @@ from .IdentityMixerNode import IdentityMixerNode
|
||||
from .PromptReplayNode import PromptReplayNode
|
||||
from .PromptBuilderNode import PromptBuilderNode
|
||||
from .PromptTagsExtractorNode import PromptTagsExtractorNode
|
||||
from .PromptRefinerNode import PromptRefinerNode
|
||||
|
||||
|
||||
|
||||
# -------------------------
|
||||
# ComfyUI mappings
|
||||
@@ -79,7 +83,9 @@ NODE_CLASS_MAPPINGS = {
|
||||
"IdentityMixerNode": IdentityMixerNode,
|
||||
"PromptReplayNode": PromptReplayNode,
|
||||
"PromptBuilderNode": PromptBuilderNode,
|
||||
"PromptTagsExtractorNode": PromptTagsExtractorNode
|
||||
"PromptTagsExtractorNode": PromptTagsExtractorNode,
|
||||
"PromptRefinerNode": PromptRefinerNode
|
||||
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
@@ -87,6 +93,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"IdentityMixerNode": "Identity Mixer (Parts)",
|
||||
"PromptReplayNode": "Prompt Replay",
|
||||
"PromptBuilderNode": "Prompt Builder",
|
||||
"PFN_PromptTagsExtractor": "PFN Prompt → Tags (Extractor)"
|
||||
"PFN_PromptTagsExtractor": "PFN Prompt → Tags (Extractor)",
|
||||
"PromptRefinerNode": "PFN Prompt Refiner",
|
||||
|
||||
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user