asaddi-ComfyUI-YALLM-node/nodes.py

import os
from typing import Optional

import openai
from pydantic import BaseModel
import yaml


BASE_PATH = os.path.dirname(os.path.realpath(__file__))


class ModelDefinition(BaseModel):
    name: str
    base_url: str
    api_key: Optional[str]
    model: str


class Models:
    LIST: list[ModelDefinition] = []
    BY_NAME: dict[str,ModelDefinition] = {}
    CHOICES: list[str] = []

    def load(self) -> None:
        with open(os.path.join(BASE_PATH, 'models.yaml')) as inp:
            d = yaml.load(inp, yaml.Loader)

        for value in d['models']:
            self.LIST.append(ModelDefinition.model_validate(value))
        if not self.LIST:
            raise RuntimeError('Need at least one model defined')
        self.BY_NAME = { d.name: d for d in self.LIST }
        self.CHOICES = [ d.name for d in self.LIST ]

    def get_base_url(self, name: str) -> str:
        base_url = self.BY_NAME[name].base_url
        return Models._get_env_or_value(base_url)

    def get_api_key(self, name: str) -> str | None:
        api_key = self.BY_NAME[name].api_key
        if api_key:
            api_key = Models._get_env_or_value(api_key)
        if not api_key or api_key == 'none':
            return None
        return api_key

    def get_model(self, name: str) -> str:
        return self.BY_NAME[name].model

    @staticmethod
    def _get_env_or_value(value: str) -> str:
        """
        Retrieves the value from either an environment variable or a provided string.

        Args:
            value (str): The value to retrieve. It can be a string literal or a string
                starting with 'os.environ/' followed by an environment variable name.

        Returns:
            str: The retrieved value.
        """
        PREFIX = 'os.environ/'
        if value.startswith(PREFIX):
            env_var = value[len(PREFIX):]
            value = os.environ[env_var]
        return value


MODELS = Models()
MODELS.load()


class SamplerDefinition(BaseModel):
    name: str
    is_int: bool
    min: Optional[float|int] = None
    default: float|int
    max: Optional[float|int] = None

    def to_input_type(self):
        constraints = {}
        if self.min is not None:
            constraints['min'] = self.min
        if self.max is not None:
            constraints['max'] = self.max
        constraints['default'] = self.default
        return {
            self.name: ('INT' if self.is_int else 'FLOAT', constraints)
        }


class LLMTemperature:
    _DEF = SamplerDefinition(
        name='temperature',
        is_int=False,
        min=0.0,
        default=0.8,
    )

    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': cls._DEF.to_input_type(),
            'optional': {
                'previous': ('LLMSAMPLER',),
            },
        }

    TITLE = 'LLM Temperature'

    RETURN_TYPES = ('LLMSAMPLER',)
    RETURN_NAMES = ('llmsampler',)

    FUNCTION = 'execute'

    CATEGORY = 'YALLM/samplers'

    def execute(self, previous=None, **args):
        value = args[self._DEF.name]

        if previous is None:
            previous = []

        previous.append((self._DEF.name, value))

        return (previous,)


class LLMTopP(LLMTemperature):
    _DEF = SamplerDefinition(
        name='top_p',
        is_int=False,
        min=0.0,
        default=0.95,
        max=1.0,
    )

    TITLE = 'LLM Top-P'

class LLMTopK(LLMTemperature):
    _DEF = SamplerDefinition(
        name='top_k',
        is_int=True,
        min=1,
        default=40,
    )

    TITLE = 'LLM Top-K'


class LLMMinP(LLMTemperature):
    _DEF = SamplerDefinition(
        name='min_p',
        is_int=False,
        min=0.0,
        default=0.05,
        max=1.0,
    )

    TITLE = 'LLM Min-P'


# TODO typical_p? Any others? mirostat?


class LLMChat:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': {
                'model': (MODELS.CHOICES,),
                'seed': ('INT', { # Thankfully, frontend has special handling for a widget of this name...
                    'max': 0xffffffff_ffffffff_ffffffff_ffffffff, # TODO what's the actual acceptable range??
                }),
                'user_prompt': ('STRING', {
                    'multiline': True,
                }),
            },
            'optional': {
                'llm_sampler': ('LLMSAMPLER', ),
                'system_prompt': ('STRING', {
                    'multiline': True,
                }),
            },
        }

    TITLE = 'LLM Chat'

    RETURN_TYPES = ('STRING',)
    RETURN_NAMES = ('completion',)

    FUNCTION = 'execute'

    CATEGORY = 'YALLM'

    def execute(self, model, seed, user_prompt, llm_sampler=None, system_prompt=None):
        if llm_sampler is None:
            llm_sampler = []

        llm = openai.OpenAI(
            base_url=MODELS.get_base_url(model),
            api_key=(MODELS.get_api_key(model) or 'none'),
        )

        messages = []
        if system_prompt:
            messages.append({'role': 'system', 'content': system_prompt})
        messages.append({'role': 'user', 'content': user_prompt})

        extra_args = {}
        extra_body = {}
        for k,v in llm_sampler:
            if k in ('temperature', 'top_p'): # The only ones supported directly by openai package
                extra_args[k] = v
            else:
                # The rest have to go into "extra_body"
                extra_body[k] = v

        # print(f'extra_args = {repr(extra_args)}')
        # print(f'extra_body = {repr(extra_body)}')

        output = llm.chat.completions.create(
            messages=messages,
            model=MODELS.get_model(model),
            seed=seed,
            extra_body=extra_body,
            **extra_args
        )

        # TODO Since we're not streaming, we'll need some sort of timeout. How to specify that?

        return (output.choices[0].message.content,)


NODE_CLASS_MAPPINGS = {
    'LLMChat': LLMChat,
    'LLMTemperature': LLMTemperature,
    'LLMTopP': LLMTopP,
    'LLMTopK': LLMTopK,
    'LLMMinP': LLMMinP,
}