60 lines
2.3 KiB
Python
60 lines
2.3 KiB
Python
from .categories import CATE_GEN
|
|
from .utils.ask_stream import openai_ask_stream
|
|
from .config.configs import config
|
|
from .utils.node_defs import FEAlwaysChangeNode
|
|
|
|
|
|
class FEGenStringGPT(FEAlwaysChangeNode):
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"system_prompt": (
|
|
"STRING",
|
|
{"default": "You are ChatGPT, a large language model trained by OpenAI.", "multiline": True, "dynamicPrompts": False}
|
|
),
|
|
|
|
"user_prompt": ("STRING", {"default": "Hello World!", "multiline": True}),
|
|
"model": ("STRING", {"default": "gpt-4", "multiline": False}),
|
|
"api": (config["openai_host"],),
|
|
"temperature": ("FLOAT", {"default": 1, "min": 0, "max": 2, "step": 0.01}),
|
|
"presence_penalty": ("FLOAT", {"default": 0, "min": 0, "max": 10, "step": 0.01}),
|
|
"max_tokens": ("INT", {"default": 2048, "min": 0, "max": 2048, "step": 1}),
|
|
"async_infer": (["YES", "NO"], {"default": "NO"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "GPT_GEN_CTX")
|
|
RETURN_NAMES = ('prompt', 'generate_ctx')
|
|
FUNCTION = "query"
|
|
CATEGORY = CATE_GEN
|
|
OUTPUT_NODE = True
|
|
|
|
def query(self, system_prompt, user_prompt, api, model, temperature, presence_penalty, max_tokens, async_infer):
|
|
if async_infer == "YES":
|
|
args = {
|
|
"system_prompt": system_prompt,
|
|
"user_prompt": user_prompt,
|
|
"model": model,
|
|
"temperature": temperature,
|
|
"presence_penalty": presence_penalty,
|
|
"max_tokens": max_tokens,
|
|
# "seed": seed,
|
|
}
|
|
return ("", args,)
|
|
result = ""
|
|
print("generating prompt...", end="")
|
|
msgs = []
|
|
if system_prompt:
|
|
msgs.append({"role": "system", "content": system_prompt})
|
|
if user_prompt:
|
|
msgs.append({"role": "user", "content": user_prompt})
|
|
for _ in openai_ask_stream(
|
|
msgs, temperature=temperature, model=model, presence_penalty=presence_penalty, max_tokens=max_tokens,
|
|
api=api):
|
|
result += _
|
|
print(_, end="")
|
|
print("done")
|
|
return (result, {},)
|