82 lines
2.9 KiB
Python
82 lines
2.9 KiB
Python
import uuid
|
|
|
|
from custom_nodes.comfyui_prompt_bench.api import _fetch_ollama_models, post_ollama_chat, OllamaChatMessage
|
|
from custom_nodes.comfyui_prompt_bench.prompt import build_prompt, refine_prompt, build_feedback
|
|
|
|
_SESSION_CHATS: dict[str, list[OllamaChatMessage]] = {}
|
|
|
|
|
|
class PBTextModelChat:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"session_seed": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xffffffffffffffff,
|
|
"step": 1,
|
|
"tooltip": "Session key. Reuse to continue chat; change to start a new session.",
|
|
}),
|
|
"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", "CHAT",)
|
|
RETURN_NAMES = ("enhanced_prompt", "model", "text", "chat",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Prompt Bench"
|
|
|
|
def process(self, session_seed: int, system_prompt: str, rating: dict, model: str, text: str = ""):
|
|
global _SESSION_CHATS
|
|
|
|
session_id = str(session_seed)
|
|
current_session = _SESSION_CHATS.get(session_id, [])
|
|
has_session_messages = len(current_session) > 0
|
|
previous_prompt = current_session[-1].content if has_session_messages else ""
|
|
|
|
prompt = build_prompt(
|
|
previous_prompt if has_session_messages else "",
|
|
"" if has_session_messages else system_prompt,
|
|
build_feedback(rating, previous_prompt=previous_prompt),
|
|
text
|
|
)
|
|
|
|
current_session.append(OllamaChatMessage(role="user", content=prompt))
|
|
|
|
# Override here since we get feedback loop text
|
|
assistant_message = post_ollama_chat(model, current_session)
|
|
if assistant_message:
|
|
current_session.append(assistant_message)
|
|
|
|
_SESSION_CHATS[session_id] = current_session
|
|
|
|
return assistant_message.content, model, text, _SESSION_CHATS[session_id],
|