File size: 1,432 Bytes
8f969b7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | 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
}
|