glonlas-comfyUI-LMStudio-nodes/client.py

from __future__ import annotations

import base64
import math
import re
import secrets
from io import BytesIO
from typing import Any, TYPE_CHECKING
from urllib.parse import urlparse

import numpy as np
from PIL import Image

if TYPE_CHECKING:
    from openai import OpenAI


DEFAULT_TIMEOUT_SECONDS = 600
DEFAULT_MAX_TOKENS = 1024
DEFAULT_TEMPERATURE = 0.7
DEFAULT_API_KEY_PLACEHOLDER = "-"
DEFAULT_API_KEY_FALLBACK = "lm-studio"
MODEL_PLACEHOLDER = "<refresh models>"

# User-facing thinking/reasoning choices exposed on the Connect node.
THINKING_AUTO = "auto"
THINKING_ON = "on"
THINKING_OFF = "off"
THINKING_OPTIONS = [THINKING_AUTO, THINKING_ON, THINKING_OFF]
DEFAULT_THINKING = THINKING_AUTO


def normalize_server_url(server_url: str) -> str:
    value = (server_url or "").strip()
    if not value:
        raise ValueError("server_url must not be empty")

    parsed = urlparse(value)
    if parsed.scheme not in {"http", "https"}:
        raise ValueError("server_url must start with http:// or https://")
    if not parsed.netloc:
        raise ValueError("server_url must include host and optional port")

    # Ensure no trailing slash and no duplicate /v1 (case-insensitive).
    value = value.rstrip("/")
    if value.lower().endswith("/v1"):
        value = value[:-3]
    return value


def to_openai_base_url(server_url: str) -> str:
    normalized = normalize_server_url(server_url)
    return f"{normalized}/v1"


def normalize_api_key(api_key: str | None) -> str:
    value = (api_key or "").strip()
    if not value or value == DEFAULT_API_KEY_PLACEHOLDER:
        return DEFAULT_API_KEY_FALLBACK
    return value


def create_openai_client(
    *,
    server_url: str,
    api_key: str | None,
    timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
    max_retries: int = 1,
):
    # Local import to keep module importable in environments without openai installed.
    from openai import OpenAI

    base_url = to_openai_base_url(server_url)
    token = normalize_api_key(api_key)
    return OpenAI(
        api_key=token,
        base_url=base_url,
        timeout=timeout_seconds,
        max_retries=max_retries,
    )


def list_models(client: OpenAI) -> list[str]:
    response = client.models.list()
    ids: list[str] = []
    for model in getattr(response, "data", []):
        model_id = getattr(model, "id", None)
        if isinstance(model_id, str) and model_id:
            ids.append(model_id)
    # Stable ordering for deterministic dropdown behavior.
    return sorted(set(ids))


def get_server_models(
    *,
    server_url: str,
    api_key: str | None,
    timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
) -> list[str]:
    client = create_openai_client(
        server_url=server_url,
        api_key=api_key,
        timeout_seconds=timeout_seconds,
    )
    return list_models(client)


def coerce_seed(value: Any) -> int | None:
    if value is None:
        return None

    if isinstance(value, (list, tuple)):
        if len(value) != 1:
            raise ValueError(
                "seed provided as list/tuple must contain exactly one element"
            )
        return coerce_seed(value[0])

    # Support numpy scalar-like values.
    if hasattr(value, "item") and not isinstance(value, (str, bytes, bytearray)):
        try:
            value = value.item()
        except Exception:
            pass

    if isinstance(value, bool):
        return int(value)
    if isinstance(value, int):
        return value
    if isinstance(value, float):
        if not math.isfinite(value):
            raise ValueError("seed float must be finite")
        return int(value)
    if isinstance(value, (bytes, bytearray)):
        try:
            value = value.decode("utf-8")
        except Exception as exc:
            raise ValueError("seed bytes must be utf-8 decodable") from exc
    if isinstance(value, str):
        stripped = value.strip()
        if not stripped:
            return None
        try:
            return int(stripped, 10)
        except ValueError as exc:
            raise ValueError(f"seed string must be an integer, got: {value!r}") from exc

    raise ValueError(
        "Unsupported seed type. Allowed: int, float, str, bytes, bytearray, "
        "None, or single-item list/tuple."
    )


def resolve_request_seed(value: Any) -> int:
    seed = coerce_seed(value)
    if seed is None or seed == -1:
        return secrets.randbelow(2**63 - 1) + 1
    if seed < -1:
        raise ValueError("seed must be -1 or >= 0")
    return seed


def comfy_image_to_base64_png_url(image_tensor) -> str:
    image_np = image_tensor.cpu().numpy()
    if image_np.ndim == 4:
        # Comfy image batches are [B, H, W, C]; for single-image prompts use first frame.
        image_np = image_np[0]
    i = np.multiply(255.0, image_np)
    img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
    buffer = BytesIO()
    img.save(buffer, format="PNG")
    b64_png = base64.b64encode(buffer.getvalue()).decode("utf-8")
    return f"data:image/png;base64,{b64_png}"


def build_responses_input_text(user_prompt: str | None) -> list[dict[str, Any]]:
    prompt_text = user_prompt or ""
    return [
        {
            "role": "user",
            "content": [{"type": "input_text", "text": prompt_text}],
        }
    ]


def extract_chat_completion_text(completion: Any) -> str:
    choices = getattr(completion, "choices", None) or []
    if not choices:
        return ""

    message = getattr(choices[0], "message", None)
    if message is None:
        return ""

    content = getattr(message, "content", None)
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: list[str] = []
        for part in content:
            if isinstance(part, str):
                parts.append(part)
                continue
            text = getattr(part, "text", None)
            if isinstance(text, str):
                parts.append(text)
                continue
            if isinstance(part, dict):
                dict_text = part.get("text")
                if isinstance(dict_text, str):
                    parts.append(dict_text)
        return "\n".join(p for p in parts if p)
    return ""


def extract_responses_text(response: Any) -> str:
    output_text = getattr(response, "output_text", None)
    if isinstance(output_text, str) and output_text:
        return output_text

    output = getattr(response, "output", None)
    if isinstance(output, list):
        parts: list[str] = []
        for item in output:
            if isinstance(item, dict):
                # Common response shape: {"content": [{"type": "output_text", "text": "..."}]}
                content = item.get("content", [])
                if isinstance(content, list):
                    for entry in content:
                        if isinstance(entry, dict):
                            text = entry.get("text")
                            if isinstance(text, str):
                                parts.append(text)
                continue

            content = getattr(item, "content", None)
            if isinstance(content, list):
                for entry in content:
                    text = getattr(entry, "text", None)
                    if isinstance(text, str):
                        parts.append(text)
            else:
                text = getattr(item, "text", None)
                if isinstance(text, str):
                    parts.append(text)

        if parts:
            return "\n".join(parts)

    return ""


_THINK_BLOCK_RE = re.compile(r"<think>.*?</think>", flags=re.IGNORECASE | re.DOTALL)
_THINK_START_RE = re.compile(r"<think>.*$", flags=re.IGNORECASE | re.DOTALL)
_THINK_CLOSE_RE = re.compile(r"</think>", flags=re.IGNORECASE)

# Wire-level thinking modes resolved from the user's THINKING_* choice.
THINK_MODE_SUPPRESS = "suppress"  # actively turn reasoning OFF
THINK_MODE_FORCE = "force"  # actively turn reasoning ON
THINK_MODE_DEFAULT = "default"  # send nothing; defer to the model's own template

# Trailing assistant turn that opens AND closes the reasoning block so a
# Qwen-style chat template starts generation *past* the <think> phase and the
# model writes prose immediately. Paired with THINK_SUPPRESS_KWARGS below,
# which newer templates (Qwen3+) actually honor. Harmless on non-thinking models.
THINK_SUPPRESS_PREFILL: dict[str, Any] = {"role": "assistant", "content": "<think></think>"}

# Template-level thinking switch — the same mechanism LM Studio's own "thinking"
# toggle uses. Forwarded to the upstream Jinja chat template via
# `chat_template_kwargs`. A *top-level* enable_thinking field is dropped before it
# reaches the template; the nesting under chat_template_kwargs is what makes it
# take effect. Ignored by templates that don't define the kwarg.
THINK_SUPPRESS_KWARGS: dict[str, Any] = {"enable_thinking": False}
THINK_ENABLE_KWARGS: dict[str, Any] = {"enable_thinking": True}

# Families whose LM Studio chat template does NOT think by default, so thinking
# level "auto" must actively pass enable_thinking:true to get any reasoning.
# Gemma (2/3/4) is the known case; Qwen 3.x-style templates think by default.
_THINK_OFF_BY_DEFAULT_RE = re.compile(r"\bgemma\b|gemma[-_]?\d", flags=re.IGNORECASE)


def model_thinks_off_by_default(model: str | None) -> bool:
    """True when the model's family needs thinking forced on in "auto" mode."""
    return bool(_THINK_OFF_BY_DEFAULT_RE.search(model or ""))


def resolve_thinking_mode(thinking: str | None, model: str | None) -> str:
    """Resolve a user THINKING_* choice + model id into a wire-level mode.

    - "off"  -> suppress (enable_thinking:false + closed <think> prefill)
    - "on"   -> force    (enable_thinking:true)
    - "auto" -> force for families that don't think by default (Gemma),
                otherwise default (send nothing — the template already thinks).
    """
    choice = (thinking or DEFAULT_THINKING).strip().lower()
    if choice == THINKING_OFF:
        return THINK_MODE_SUPPRESS
    if choice == THINKING_ON:
        return THINK_MODE_FORCE
    return THINK_MODE_FORCE if model_thinks_off_by_default(model) else THINK_MODE_DEFAULT


def thinking_chat_template_kwargs(mode: str) -> dict[str, Any] | None:
    """chat_template_kwargs payload for a resolved mode (None = send nothing)."""
    if mode == THINK_MODE_SUPPRESS:
        return dict(THINK_SUPPRESS_KWARGS)
    if mode == THINK_MODE_FORCE:
        return dict(THINK_ENABLE_KWARGS)
    return None


def apply_thinking_prefill(
    messages: list[dict[str, Any]], mode: str
) -> list[dict[str, Any]]:
    """Append the closed-<think> suppression prefill when actively suppressing."""
    if mode == THINK_MODE_SUPPRESS:
        return [*messages, dict(THINK_SUPPRESS_PREFILL)]
    return messages


def strip_think_content(text: str | None) -> str:
    """
    Remove reasoning blocks enclosed in <think>...</think>.
    Also removes dangling <think> blocks without a closing tag and any
    orphan </think> closing tags left by nested or malformed think blocks.
    """
    if not text:
        return ""

    cleaned = _THINK_BLOCK_RE.sub("", text)
    cleaned = _THINK_START_RE.sub("", cleaned)
    # After the above two passes, any remaining </think> tags are orphaned
    # closing tags produced by nested blocks (e.g. <think><think>…</think></think>).
    # The non-greedy block regex only removes up to the *first* </think>, leaving
    # the outer one behind; strip it here so it never appears in visible output.
    cleaned = _THINK_CLOSE_RE.sub("", cleaned)
    return cleaned.strip()


def build_chat_messages(system_prompt: str | None, user_prompt: str | None) -> list[dict[str, Any]]:
    system_text = system_prompt or ""
    user_text = user_prompt or ""
    messages: list[dict[str, Any]] = []
    if system_text.strip():
        messages.append({"role": "system", "content": system_text})
    messages.append({"role": "user", "content": user_text})
    return messages


def build_chat_messages_with_image(
    system_prompt: str | None,
    user_prompt: str | None,
    image_data_url: str,
) -> list[dict[str, Any]]:
    system_text = system_prompt or ""
    user_text = user_prompt or ""
    messages: list[dict[str, Any]] = []
    if system_text.strip():
        messages.append({"role": "system", "content": system_text})
    messages.append(
        {
            "role": "user",
            "content": [
                {"type": "text", "text": user_text},
                {"type": "image_url", "image_url": {"url": image_data_url}},
            ],
        }
    )
    return messages