text_gen_node.py

from __future__ import annotations

from typing import Any

from comfy_api.latest import io, ui

from .client import (
    build_chat_messages,
    build_responses_input_text,
    create_openai_client,
    extract_chat_completion_text,
    extract_responses_text,
    resolve_request_seed,
    strip_think_content,
)
from .iotypes import ParamConnection
from .models import LMStudioConnectionPayload


class LMStudioTextGen(io.ComfyNode):
    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="LMStudio_TextGen",
            display_name="LMStudio - Text Gen",
            category="LMStudio",
            description="Generate text from an LMStudio connection with system and user prompts.",
            inputs=[
                ParamConnection.Input(
                    id="connection",
                    display_name="Connection",
                    tooltip="Connection output from LMStudio - Connect node.",
                ),
                io.String.Input(
                    id="system_prompt",
                    display_name="System Prompt (LM Souls)",
                    multiline=True,
                    default="You are LM Souls, an expert prompt engineer and creative assistant.",
                    tooltip="System instructions that define assistant behavior.",
                ),
                io.String.Input(
                    id="user_prompt",
                    display_name="User Prompt",
                    multiline=True,
                    default="Write a concise prompt for a cinematic portrait scene.",
                    tooltip="User input to transform into a final response.",
                ),
                io.Int.Input(
                    id="seed",
                    display_name="Seed",
                    default=-1,
                    min=-1,
                    max=0x7FFFFFFFFFFFFFFF,
                    control_after_generate=io.ControlAfterGenerate.randomize,
                    tooltip="-1 uses a secure random seed.",
                ),
            ],
            outputs=[
                io.String.Output(
                    id="response_text",
                    display_name="Response",
                    tooltip="Generated text output.",
                ),
            ],
        )

    @classmethod
    def validate_inputs(cls, user_prompt: str | None = None) -> bool | str:
        # user_prompt is often linked from another node; defer strict checks to execute().
        return True

    @classmethod
    def _responses_kwargs(
        cls,
        *,
        connection: LMStudioConnectionPayload,
        system_prompt: str,
        user_prompt: str,
        seed: int,
    ) -> dict[str, Any]:
        kwargs: dict[str, Any] = {
            "model": connection.model,
            "input": build_responses_input_text(user_prompt),
            "temperature": connection.temperature,
            "max_output_tokens": connection.max_tokens,
            "seed": seed,
        }
        if system_prompt.strip():
            kwargs["instructions"] = system_prompt
        if connection.reasoning_enabled:
            kwargs["reasoning"] = {"effort": "medium"}

        # MCP/tooling requires additional tool definitions and server details.
        # Keep an explicit metadata signal for future extension without sending invalid tool payloads.
        if connection.use_tooling_mcp:
            kwargs["metadata"] = {"lmstudio_tooling_mcp_requested": "true"}

        return kwargs

    @classmethod
    def execute(
        cls,
        connection: LMStudioConnectionPayload,
        system_prompt: str,
        user_prompt: str,
        seed: int,
    ) -> io.NodeOutput:
        system_prompt = system_prompt or ""
        user_prompt = user_prompt or ""
        if not user_prompt.strip():
            raise ValueError("user_prompt must not be empty")

        resolved_seed = resolve_request_seed(seed)
        client = create_openai_client(
            server_url=connection.server_url,
            api_key=connection.api_key,
            timeout_seconds=connection.timeout_seconds,
        )

        via_endpoint = "responses"
        fallback_reason: str | None = None

        try:
            response = client.responses.create(
                **cls._responses_kwargs(
                    connection=connection,
                    system_prompt=system_prompt,
                    user_prompt=user_prompt,
                    seed=resolved_seed,
                )
            )
            text = extract_responses_text(response).strip()
            if not text:
                raise ValueError("responses endpoint returned no text output")
        except Exception as responses_error:
            via_endpoint = "chat.completions"
            fallback_reason = str(responses_error)

            completion = client.chat.completions.create(
                model=connection.model,
                messages=build_chat_messages(system_prompt, user_prompt),
                seed=resolved_seed,
                temperature=connection.temperature,
                max_tokens=connection.max_tokens,
                n=1,
            )
            text = extract_chat_completion_text(completion).strip()
            if not text:
                raise ValueError("chat.completions fallback returned no text output")

        text = strip_think_content(text)

        status = f"Model '{connection.model}' via {via_endpoint} (seed={resolved_seed})."
        if fallback_reason:
            status += f" Fallback reason: {fallback_reason}"

        return io.NodeOutput(
            text,
            ui=ui.PreviewText(status),
        )