50 lines
1.3 KiB
Python
50 lines
1.3 KiB
Python
import uuid
|
|
|
|
from custom_nodes.comfyui_prompt_bench.api import OllamaChatMessage
|
|
|
|
|
|
class PBListChatView:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"chat": ("CHAT",),
|
|
"text": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "Latest rendered chat text.",
|
|
},
|
|
),
|
|
},
|
|
"optional": {}
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return str(uuid.uuid4())
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Prompt Bench"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, chat: list[OllamaChatMessage], text: str = ""):
|
|
formatted_parts = []
|
|
|
|
for msg in chat:
|
|
role = msg.role.upper()
|
|
content = msg.content.strip()
|
|
formatted_parts.append(f"### {role} ###\n\n{content}\n\n---{'--------' if role == 'ASSISTANT' else ''}")
|
|
|
|
text = "\n\n".join(formatted_parts)
|
|
if not text:
|
|
text = "(Empty Chat)"
|
|
|
|
return {
|
|
"ui": {"text": (text,)},
|
|
"result": (text,),
|
|
}
|