57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import json
|
|
import uuid
|
|
|
|
from custom_nodes.comfyui_prompt_bench.api import _fetch_ollama_models, post_ollama_text
|
|
from custom_nodes.comfyui_prompt_bench.prompt import build_prompt
|
|
|
|
|
|
class PBTextModel:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"system_prompt": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"default": "",
|
|
"dynamicPrompts": True,
|
|
"tooltip": "System prompt applied to the Ollama model.",
|
|
},
|
|
),
|
|
"rating": ("RATING",),
|
|
"model": (
|
|
"COMBO",
|
|
{
|
|
"options": _fetch_ollama_models(),
|
|
"remote": {
|
|
"route": "/prompt_bench/ollama/models",
|
|
"refresh_button": True,
|
|
},
|
|
"tooltip": "Ollama model name. Refresh to fetch the current local model list.",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"text": ("STRING",),
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return str(uuid.uuid4())
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING",)
|
|
RETURN_NAMES = ("enhanced_prompt", "model", "text",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Prompt Bench"
|
|
|
|
def process(self, system_prompt: str, rating: dict, model: str, text: str = ""):
|
|
rating_context = json.dumps(rating, ensure_ascii=True, sort_keys=True)
|
|
prompt = build_prompt("", system_prompt, rating_context, text)
|
|
|
|
# Override here since we get feedback loop text
|
|
text = post_ollama_text(model, prompt)
|
|
|
|
return text, model, text,
|