294 lines
7.4 KiB
Python
294 lines
7.4 KiB
Python
# Copyright (c) 2024 Allan Saddi <allan@saddi.com>
|
|
import os
|
|
from typing import Any, 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 = ('llm_sampler',)
|
|
|
|
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? tfs_z? Any others? mirostat?
|
|
|
|
|
|
# TODO Will we need a more complex definition? (e.g. multi-modal)
|
|
ChatMessage = dict[str,str]
|
|
ChatHistory = list[ChatMessage]
|
|
|
|
SamplerSetting = tuple[str,Any]
|
|
|
|
|
|
class LLMModel:
|
|
def __init__(self, model: str):
|
|
self._llm = openai.OpenAI(
|
|
base_url=MODELS.get_base_url(model),
|
|
api_key=(MODELS.get_api_key(model) or 'none'), # openai package requires something for API key
|
|
)
|
|
self._model = MODELS.get_model(model)
|
|
|
|
def chat_completion(self, messages: ChatHistory, samplers: list[SamplerSetting]|None=None, seed: int|None=None) -> str:
|
|
if samplers is None:
|
|
samplers = []
|
|
|
|
sampler_order = []
|
|
extra_args = {}
|
|
extra_body = {}
|
|
for k,v in samplers:
|
|
sampler_order.append(k)
|
|
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
|
|
extra_body['samplers'] = sampler_order # Hopefully won't cause an issue for non-llama.cpp endpoints
|
|
|
|
if seed is not None:
|
|
extra_args['seed'] = seed
|
|
|
|
# print(f'extra_args = {repr(extra_args)}')
|
|
# print(f'extra_body = {repr(extra_body)}')
|
|
|
|
output = self._llm.chat.completions.create(
|
|
messages=messages,
|
|
model=self._model,
|
|
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
|
|
|
|
|
|
class LLMModelNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
'required': {
|
|
'model': (MODELS.CHOICES,),
|
|
},
|
|
}
|
|
|
|
TITLE = 'LLM Model (API)'
|
|
|
|
RETURN_TYPES = ('LLMMODEL',)
|
|
RETURN_NAMES = ('llm_model',)
|
|
|
|
FUNCTION = 'execute'
|
|
|
|
CATEGORY = 'YALLM'
|
|
|
|
def execute(self, model):
|
|
llm = LLMModel(model)
|
|
|
|
return (llm,)
|
|
|
|
|
|
class LLMChat:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
'required': {
|
|
'llm_model': ('LLMMODEL',),
|
|
'seed': ('INT', { # Thankfully, frontend has special handling for a widget of this name...
|
|
# min/max/default from the KSampler nodes
|
|
# llama.cpp is only 32-bit (with truncation, no doubt), OpenAI API spec doesn't specify a size
|
|
'min': 0,
|
|
'default': 0,
|
|
'max': 0xffffffff_ffffffff,
|
|
}),
|
|
'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, llm_model: LLMModel, seed, user_prompt, llm_sampler=None, system_prompt=None):
|
|
messages = []
|
|
if system_prompt:
|
|
messages.append({'role': 'system', 'content': system_prompt})
|
|
messages.append({'role': 'user', 'content': user_prompt})
|
|
|
|
result = llm_model.chat_completion(messages, samplers=llm_sampler, seed=seed)
|
|
|
|
return (result,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
'LLMModel': LLMModelNode,
|
|
'LLMChat': LLMChat,
|
|
'LLMTemperature': LLMTemperature,
|
|
'LLMTopP': LLMTopP,
|
|
'LLMTopK': LLMTopK,
|
|
'LLMMinP': LLMMinP,
|
|
}
|