JasonHoku-comfyui-llm-tools/llm_chat_node.py

import json
import time

import folder_paths
from llama_cpp import Llama

from .llm_utils import resolve_gpu_layers

# Reuse the folder registration from llm_gguf_node (already registered at import time)


class LLMChatStart:
    """Initialize a chat conversation with a system prompt and first user message."""

    _model_cache = {}

    @classmethod
    def INPUT_TYPES(cls):
        model_list = folder_paths.get_filename_list("LLMs")
        return {
            "required": {
                "model": (model_list, {"tooltip": "Select a .gguf model from ComfyUI/models/LLM"}),
                "system_prompt": ("STRING", {
                    "multiline": True,
                    "default": "You are a helpful assistant.",
                    "tooltip": "System prompt to set the model's behavior",
                }),
                "user_message": ("STRING", {
                    "multiline": True,
                    "default": "",
                    "tooltip": "First user message to start the conversation",
                }),
                "max_tokens": ("INT", {"default": 512, "min": 1, "max": 4096, "step": 1}),
                "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.05}),
                "top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.05}),
                "context_length": ("INT", {"default": 2048, "min": 512, "max": 32768, "step": 256}),
                "gpu_layers": ("INT", {"default": -1, "min": -1, "max": 200, "step": 1}),
                "keep_model_loaded": ("BOOLEAN", {"default": True}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING",)
    RETURN_NAMES = ("response", "chat_history",)
    FUNCTION = "start_chat"
    CATEGORY = "LLM"

    def start_chat(
        self, model, system_prompt, user_message,
        max_tokens, temperature, top_p,
        context_length, gpu_layers, keep_model_loaded, seed,
    ):
        llm = _get_or_load_model(
            self._model_cache, model, context_length, gpu_layers,
        )

        messages = []
        if system_prompt.strip():
            messages.append({"role": "system", "content": system_prompt.strip()})
        messages.append({"role": "user", "content": user_message})

        response_text, elapsed = _run_inference(
            llm, messages, max_tokens, temperature, top_p, seed,
        )

        messages.append({"role": "assistant", "content": response_text})
        chat_history = json.dumps(messages, ensure_ascii=False)

        if not keep_model_loaded:
            _evict_cache(self._model_cache)

        return (response_text, chat_history,)


class LLMChatContinue:
    """Continue an existing chat conversation with a new user message."""

    _model_cache = {}

    @classmethod
    def INPUT_TYPES(cls):
        model_list = folder_paths.get_filename_list("LLMs")
        return {
            "required": {
                "model": (model_list, {"tooltip": "Select a .gguf model from ComfyUI/models/LLM"}),
                "chat_history": ("STRING", {
                    "multiline": True,
                    "default": "",
                    "tooltip": "Chat history JSON from a previous Chat Start or Chat Continue node",
                }),
                "user_message": ("STRING", {
                    "multiline": True,
                    "default": "",
                    "tooltip": "Next user message in the conversation",
                }),
                "max_tokens": ("INT", {"default": 512, "min": 1, "max": 4096, "step": 1}),
                "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.05}),
                "top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.05}),
                "context_length": ("INT", {"default": 2048, "min": 512, "max": 32768, "step": 256}),
                "gpu_layers": ("INT", {"default": -1, "min": -1, "max": 200, "step": 1}),
                "keep_model_loaded": ("BOOLEAN", {"default": True}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING",)
    RETURN_NAMES = ("response", "chat_history",)
    FUNCTION = "continue_chat"
    CATEGORY = "LLM"

    def continue_chat(
        self, model, chat_history, user_message,
        max_tokens, temperature, top_p,
        context_length, gpu_layers, keep_model_loaded, seed,
    ):
        try:
            messages = json.loads(chat_history)
            if not isinstance(messages, list):
                raise ValueError("chat_history must be a JSON array")
        except (json.JSONDecodeError, ValueError) as e:
            raise ValueError(
                f"Invalid chat_history: {e}. "
                "Connect this to the chat_history output of a Chat Start or Chat Continue node."
            )

        messages.append({"role": "user", "content": user_message})

        llm = _get_or_load_model(
            self._model_cache, model, context_length, gpu_layers,
        )

        response_text, elapsed = _run_inference(
            llm, messages, max_tokens, temperature, top_p, seed,
        )

        messages.append({"role": "assistant", "content": response_text})
        updated_history = json.dumps(messages, ensure_ascii=False)

        if not keep_model_loaded:
            _evict_cache(self._model_cache)

        return (response_text, updated_history,)


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

def _get_or_load_model(cache, model, context_length, gpu_layers):
    model_path = folder_paths.get_full_path_or_raise("LLMs", model)
    actual_gpu_layers = resolve_gpu_layers(model_path, context_length, gpu_layers)
    cache_key = (model_path, context_length, actual_gpu_layers)
    llm = cache.get(cache_key)

    if llm is None:
        _evict_cache(cache)
        print(f"[LLM-Tools] Loading model: {model}  gpu_layers={actual_gpu_layers}")
        llm = Llama(
            model_path=model_path,
            n_ctx=context_length,
            n_gpu_layers=actual_gpu_layers,
            verbose=False,
        )
        cache[cache_key] = llm

    return llm


def _run_inference(llm, messages, max_tokens, temperature, top_p, seed):
    start = time.perf_counter()
    result = llm.create_chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        seed=seed if seed != 0 else None,
    )
    elapsed = time.perf_counter() - start

    response_text = ""
    if result and "choices" in result and len(result["choices"]) > 0:
        response_text = result["choices"][0].get("message", {}).get("content", "")

    tokens_used = result.get("usage", {})
    print(
        f"[LLM-Tools] Generated in {elapsed:.1f}s | "
        f"prompt_tokens={tokens_used.get('prompt_tokens', '?')} "
        f"completion_tokens={tokens_used.get('completion_tokens', '?')}"
    )

    return response_text, elapsed


def _evict_cache(cache):
    for key, llm in list(cache.items()):
        try:
            llm.close()
        except Exception:
            pass
        del cache[key]
    cache.clear()

    import gc
    gc.collect()
    gc.collect()

    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
    except ImportError:
        pass


NODE_CLASS_MAPPINGS = {
    "LLMChatStart": LLMChatStart,
    "LLMChatContinue": LLMChatContinue,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LLMChatStart": "LLM Chat Start",
    "LLMChatContinue": "LLM Chat Continue",
}