vnyash
40 lines
1.4 KiB
Python
40 lines
1.4 KiB
Python
from openai import OpenAI
|
|
|
|
class LLM_node:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {
|
|
"required":{
|
|
"system_prompt" : ("STRING",{"default":"","multiline": True,},),
|
|
"user_query": ("STRING", {"default":"","multiline": True,},),
|
|
"model_name": (["gpt-3.5-turbo","gpt-4o","gpt-4-turbo","gpt-4",""],),
|
|
"api_key": ("STRING",{"default":""},)
|
|
},
|
|
"optional":{
|
|
"max_tokens" : ("INT", {"default":250,"min":10,"max":2000,"step":10},),
|
|
"temperature" : ("FLOAT", {"default":0,"min":0,"max":1,"step":0.1}),
|
|
"timeout": ("INT", {"default":10,"min":1,"max":200,"step":1},),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "LLM"
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
FUNCTION = "run_llm"
|
|
|
|
def run_llm(self, system_prompt, user_query, model_name,temperature,api_key,max_tokens,timeout):
|
|
client = OpenAI(api_key=api_key)
|
|
response = client.chat.completions.create(
|
|
model=model_name,
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_query}
|
|
],
|
|
# stop = stop.split(","),
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
timeout=timeout
|
|
)
|
|
return (text,)
|