Sam Khoze
45 lines
1.6 KiB
Python
45 lines
1.6 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 = "DeepFuze"
|
|
RETURN_TYPES = ("NEW_STRING",)
|
|
RETURN_NAMES = ("LLM_RESPONSE",)
|
|
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)
|
|
try:
|
|
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
|
|
).to_dict()
|
|
print(response)
|
|
result = response["choices"][0]["message"]["content"]
|
|
except Exception as e:
|
|
raise ValueError(e)
|
|
return (result,)
|