comfyui-chatgpt-node.py

import requests
from typing import Dict, Any

class OpenAIChatNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "system_prompt": ("STRING", {"multiline": True}),
                "user_input": ("STRING", {"multiline": True}),
                "api_key": ("STRING", {"multiline": False}),
            },
        }

    RETURN_TYPES = ("STRING",)
    FUNCTION = "chat_with_openai"
    CATEGORY = "OpenAI"

    def chat_with_openai(self, system_prompt: str, user_input: str, api_key: str) -> (str,):
        url = "https://api.openai.com/v1/chat/completions"

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}",
        }

        payload = {
            "model": "gpt-4",
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_input},
            ],
        }

        response = requests.post(url, headers=headers, json=payload)

        if response.status_code == 200:
            data = response.json()
            assistant_message = data['choices'][0]['message']['content']
            return (assistant_message,)
        else:
            error_message = f"Error {response.status_code}: {response.text}"
            return (error_message,)

# Registrar el nodo
NODE_CLASS_MAPPINGS = {
    "OpenAIChatNode": OpenAIChatNode
}