From 890bdfd44bd571e35eac8f5ade260aab111c0970 Mon Sep 17 00:00:00 2001 From: Hekmon <3418467+hekmon@users.noreply.github.com> Date: Wed, 2 Apr 2025 10:38:11 +0200 Subject: [PATCH] cleanup --- __init__.py | 30 +++++++++++++++--------------- completions.py | 15 +++++++++------ iotypes.py | 3 ++- options.py | 42 +++++++++++++++++++++--------------------- 4 files changed, 47 insertions(+), 43 deletions(-) diff --git a/__init__.py b/__init__.py index bf779a7..ca2e0e7 100644 --- a/__init__.py +++ b/__init__.py @@ -1,29 +1,29 @@ from .client import Client from .completions import ChatCompletion -from .options import OptionTemperature, OptionMaxTokens, OptionTopP, OptionFrequencyPenalty, OptionPresencePenalty, OptionExtraBody, OptionDebug +from .options import OptionTemperature, OptionMaxTokens, OptionTopP, OptionFrequencyPenalty, OptionPresencePenalty, OptionExtraBody, OptionsDebug NODE_CLASS_MAPPINGS = { "OAIAPIClient": Client, "OAIAPIChatCompletion": ChatCompletion, - "OAIAPIOptionTemperature": OptionTemperature, - "OAIAPIOptionMaxTokens": OptionMaxTokens, - "OAIAPIOptionTopP": OptionTopP, - "OAIAPIOptionFrequencyPenalty": OptionFrequencyPenalty, - "OAIAPIOptionPresencePenalty": OptionPresencePenalty, - "OAIAPIOptionExtraBody": OptionExtraBody, - "OAIAPIOptionDebug": OptionDebug, + "OAIAPITemperature": OptionTemperature, + "OAIAPIMaxTokens": OptionMaxTokens, + "OAIAPITopP": OptionTopP, + "OAIAPIFrequencyPenalty": OptionFrequencyPenalty, + "OAIAPIPresencePenalty": OptionPresencePenalty, + "OAIAPIExtraBody": OptionExtraBody, + "OAIAPIDebug": OptionsDebug, } NODE_DISPLAY_NAME_MAPPINGS = { "OAIAPIClient": "OpenAI API - Client", "OAIAPIChatCompletion": "OpenAI API - Chat Completion", - "OAIAPIOptionTemperature": "OpenAI API - Option - Temperature", - "OAIAPIOptionMaxTokens": "OpenAI API - Option - Max Tokens", - "OAIAPIOptionTopP": "OpenAI API - Option - Top P", - "OAIAPIOptionFrequencyPenalty": "OpenAI API - Option - Frequency Penalty", - "OAIAPIOptionPresencePenalty": "OpenAI API - Option - Presence Penalty", - "OAIAPIOptionExtraBody": "OpenAI API - Option - Extra Body", - "OAIAPIOptionDebug": "OpenAI API - Option - Debug", + "OAIAPITemperature": "OpenAI API - Temperature", + "OAIAPIMaxTokens": "OpenAI API - Max Tokens", + "OAIAPITopP": "OpenAI API - Top P", + "OAIAPIFrequencyPenalty": "OpenAI API - Frequency Penalty", + "OAIAPIPresencePenalty": "OpenAI API - Presence Penalty", + "OAIAPIExtraBody": "OpenAI API - Extra Body", + "OAIAPIDebug": "OpenAI API - Debug", } __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/completions.py b/completions.py index 7e8e74d..db55ab5 100644 --- a/completions.py +++ b/completions.py @@ -7,8 +7,8 @@ from .iotypes import OAIAPIIO class ChatCompletion: CATEGORY = "OpenAI API" - RETURN_TYPES = (IO.STRING,) - RETURN_NAMES = ("Response",) + RETURN_TYPES = (IO.STRING,OAIAPIIO.HISTORY) + RETURN_NAMES = ("RESPONSE","HISTORY") FUNCTION = "generate" @classmethod @@ -36,7 +36,7 @@ class ChatCompletion: "default": False, "tooltip": "With o1 models and newer, OpenAI has changed the 'system' prompt role to 'developper' prompt role. Set this switch to true to set the system prompt as 'developper'.", }), - "options": (OAIAPIIO.CHAT_COMPLETION_OPTIONS, { + "options": (OAIAPIIO.OPTIONS, { "default": None, "tooltip": "Chat completion options. This can be used to specify additional parameters for the chat completion request.", }) @@ -93,6 +93,7 @@ class ChatCompletion: delattr(options, "presence_penalty") else: presence_penalty = None + print(f"final options: {options}") # Create the completion completion = client.chat.completions.create( model=model, @@ -102,7 +103,9 @@ class ChatCompletion: top_p=top_p, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, - extra_body=options - + extra_body=options, + n=1 ) - return (completion.choices[0].message.content,) + print(f"Completion: {completion.choices[0]}") + # Return the response + return (completion.choices[0].message.content,"") diff --git a/iotypes.py b/iotypes.py index 7b32738..75c0288 100644 --- a/iotypes.py +++ b/iotypes.py @@ -5,5 +5,6 @@ class OAIAPIIO(StrEnum): Node input/output data types. """ CLIENT = "CLIENT" - CHAT_COMPLETION_OPTIONS = "CHAT_COMPLETION_OPTIONS" + OPTIONS = "OPTIONS" + HISTORY = "HISTORY" diff --git a/options.py b/options.py index a2fda6b..356c153 100644 --- a/options.py +++ b/options.py @@ -5,8 +5,8 @@ from comfy.comfy_types import IO from .iotypes import OAIAPIIO class OptionTemperature: - CATEGORY = "OpenAI API" - RETURN_TYPES = (OAIAPIIO.CHAT_COMPLETION_OPTIONS,) + CATEGORY = "OpenAI API/Options" + RETURN_TYPES = (OAIAPIIO.OPTIONS,) RETURN_NAMES = ("OPTIONS",) FUNCTION = "merge" @@ -23,7 +23,7 @@ class OptionTemperature: }), }, "optional": { - "options": (OAIAPIIO.CHAT_COMPLETION_OPTIONS, { + "options": (OAIAPIIO.OPTIONS, { "tooltip": "Others options to merge with", }), }, @@ -38,8 +38,8 @@ class OptionTemperature: class OptionMaxTokens: - CATEGORY = "OpenAI API" - RETURN_TYPES = (OAIAPIIO.CHAT_COMPLETION_OPTIONS,) + CATEGORY = "OpenAI API/Options" + RETURN_TYPES = (OAIAPIIO.OPTIONS,) RETURN_NAMES = ("OPTIONS",) FUNCTION = "merge" @@ -54,7 +54,7 @@ class OptionMaxTokens: }), }, "optional": { - "options": (OAIAPIIO.CHAT_COMPLETION_OPTIONS, { + "options": (OAIAPIIO.OPTIONS, { "tooltip": "Others options to merge with", }), }, @@ -69,8 +69,8 @@ class OptionMaxTokens: class OptionTopP: - CATEGORY = "OpenAI API" - RETURN_TYPES = (OAIAPIIO.CHAT_COMPLETION_OPTIONS,) + CATEGORY = "OpenAI API/Options" + RETURN_TYPES = (OAIAPIIO.OPTIONS,) RETURN_NAMES = ("OPTIONS",) FUNCTION = "merge" @@ -86,7 +86,7 @@ class OptionTopP: }) }, "optional": { - "options": (OAIAPIIO.CHAT_COMPLETION_OPTIONS, { + "options": (OAIAPIIO.OPTIONS, { "tooltip": "Others options to merge with", }), }, @@ -101,8 +101,8 @@ class OptionTopP: class OptionFrequencyPenalty: - CATEGORY = "OpenAI API" - RETURN_TYPES = (OAIAPIIO.CHAT_COMPLETION_OPTIONS,) + CATEGORY = "OpenAI API/Options" + RETURN_TYPES = (OAIAPIIO.OPTIONS,) RETURN_NAMES = ("OPTIONS",) FUNCTION = "merge" @@ -118,7 +118,7 @@ class OptionFrequencyPenalty: }), }, "optional": { - "options": (OAIAPIIO.CHAT_COMPLETION_OPTIONS, { + "options": (OAIAPIIO.OPTIONS, { "tooltip": "Others options to merge with", }), }, @@ -133,8 +133,8 @@ class OptionFrequencyPenalty: class OptionPresencePenalty: - CATEGORY = "OpenAI API" - RETURN_TYPES = (OAIAPIIO.CHAT_COMPLETION_OPTIONS,) + CATEGORY = "OpenAI API/Options" + RETURN_TYPES = (OAIAPIIO.OPTIONS,) RETURN_NAMES = ("OPTIONS",) FUNCTION = "merge" @@ -150,7 +150,7 @@ class OptionPresencePenalty: }), }, "optional": { - "options": (OAIAPIIO.CHAT_COMPLETION_OPTIONS, { + "options": (OAIAPIIO.OPTIONS, { "tooltip": "Others options to merge with", }), }, @@ -165,8 +165,8 @@ class OptionPresencePenalty: class OptionExtraBody: - CATEGORY = "OpenAI API" - RETURN_TYPES = (OAIAPIIO.CHAT_COMPLETION_OPTIONS,) + CATEGORY = "OpenAI API/Options" + RETURN_TYPES = (OAIAPIIO.OPTIONS,) RETURN_NAMES = ("OPTIONS",) FUNCTION = "merge" @@ -181,7 +181,7 @@ class OptionExtraBody: }), }, "optional": { - "options": (OAIAPIIO.CHAT_COMPLETION_OPTIONS, { + "options": (OAIAPIIO.OPTIONS, { "tooltip": "Others options to merge with", }), }, @@ -206,8 +206,8 @@ class OptionExtraBody: return (options,) -class OptionDebug: - CATEGORY = "OpenAI API" +class OptionsDebug: + CATEGORY = "OpenAI API/Options" RETURN_TYPES = (IO.STRING,) RETURN_NAMES = ("ENCODED_JSON",) FUNCTION = "marshall" @@ -217,7 +217,7 @@ class OptionDebug: def INPUT_TYPES(cls): return { "required": { - "options": (OAIAPIIO.CHAT_COMPLETION_OPTIONS, { + "options": (OAIAPIIO.OPTIONS, { "tooltip": "Options chain you want encode", }), },