From 35a4162185ff7178811c0a64c0b1348071c836fb Mon Sep 17 00:00:00 2001
From: Logan zoellner <nagolinc@gmail.com>
Date: Fri, 14 Jun 2024 09:06:59 -0400
Subject: [PATCH] initial commit

---
 README.md        |  8 ++++++
 __init__.py      | 63 ++++++++++++++++++++++++++++++++++++++++++++++++
 requirements.txt |  1 +
 3 files changed, 72 insertions(+)
 create mode 100644 README.md
 create mode 100644 __init__.py
 create mode 100644 requirements.txt

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..b9a1593
--- /dev/null
+++ b/README.md
@@ -0,0 +1,8 @@
+This provides a single node openai > Open AI query node
+
+that takes a system prompt and user message and sends them to chatGPT 3.5
+
+Note, you MUST have an OPEN AI API key stored in the environment variable OPENAI_API_KEY
+ in order for this to work.
+
+ 
\ No newline at end of file
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..21489c4
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,63 @@
+import openai
+import os
+
+class OpenAINode:
+    FUNCTION = "execute"
+    CATEGORY = "OpenAI"
+    RETURN_TYPES = ("STRING",)
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "system_prompt": ("STRING", {
+                    "multiline": True,
+                    "default": "Enter system prompt here..."
+                }),
+                "user_string": ("STRING", {
+                    "multiline": True,
+                    "default": "Enter user input here..."
+                }),
+                "seed": ("INT", {
+                    "default": -1
+                })
+            }
+        }
+
+    def __init__(self):
+        #load key from environment variable
+        self.api_key = os.getenv("OPENAI_API_KEY")
+        openai.api_key = self.api_key
+
+    def execute(self, system_prompt, user_string,seed=-1):
+        try:
+            
+            print("DOING CHAT WITH",user_string)
+            
+            chat_completion = openai.chat.completions.create(
+                messages=[
+                    {
+                        "role": "system",
+                        "content": system_prompt,
+                    },
+                    {
+                        "role": "user",
+                        "content": user_string,
+                    },
+                ],
+                model="gpt-3.5-turbo",
+            )
+            
+            output=chat_completion.choices[0].message.content
+            
+            return (output.strip(),)
+        except Exception as e:
+            return (f"Error: {str(e)}",)
+
+NODE_CLASS_MAPPINGS = {
+    "OpenAINode": OpenAINode
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "OpenAINode": "OpenAI Query Node"
+}
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..f0dd0ae
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1 @@
+openai
\ No newline at end of file