From 8b8ebc35aa3770ed1c4acde9043bfc71246ad667 Mon Sep 17 00:00:00 2001
From: Marco Loewe <marco@h2867845.stratoserver.net>
Date: Wed, 15 Nov 2023 12:05:36 +0100
Subject: [PATCH] Added GPT-3.5: Commands /ask for questions (without context
 atm) and /image for image generation

---
 .gitignore  |  7 ++++++-
 bot.py      | 21 ++++++++++++++++++++-
 chatgpt.py  | 33 +++++++++++++++++++++++++++++++++
 config.json |  3 ++-
 4 files changed, 61 insertions(+), 3 deletions(-)
 create mode 100644 chatgpt.py

diff --git a/.gitignore b/.gitignore
index 51b9aa7..f0fbd89 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,6 @@
-config.local.json
\ No newline at end of file
+config.local.json
+
+__pycache__/
+.vscode/
+
+test.py
\ No newline at end of file
diff --git a/bot.py b/bot.py
index 2ce6e20..f208184 100644
--- a/bot.py
+++ b/bot.py
@@ -5,6 +5,8 @@ import json
 import discord
 from discord import app_commands
 
+from chatgpt import ChatGPT
+
 with open('config.json', 'r') as f:
     data = json.load(f)
 
@@ -30,7 +32,6 @@ async def on_ready():
         await client.get_channel(1173561093579362305).send('Update done')
         write_json('update', 'none')
 
-
 @tree.command(name = 'cool', description = 'Check how cool you are!', guild=discord.Object(id='745287426402156555'))
 async def cool(ctx: discord.Interaction):
     cool_factor = round(random.uniform(0, 100), 2)
@@ -43,6 +44,24 @@ async def update(ctx: discord.Interaction):
     write_json('update', 'pending')
     subprocess.run('./update.sh')
 
+@tree.command(name = 'ask', description = 'Query anything, powered by GPT-3.5', guild=discord.Object(id='745287426402156555'))
+async def ask(ctx: discord.Interaction, query: str):
+    await ctx.response.defer(thinking=True)
+
+    gpt = ChatGPT(data['openAI'])
+    response = gpt.completion(query)
+    
+    await ctx.followup.send(response)
+
+@tree.command(name = 'image', description = 'Generate an image from a text prompt, powered by DALL-E-3', guild=discord.Object(id='745287426402156555'))
+async def image(ctx: discord.Interaction, prompt: str):
+    await ctx.response.defer(thinking=True)
+
+    gpt = ChatGPT(data['openAI'])
+    response = gpt.image(prompt)
+
+    await ctx.followup.send(response)
+
 @update.error
 async def update_error(ctx: discord.Interaction, error):
     if isinstance(error, discord.app_commands.MissingRole):
diff --git a/chatgpt.py b/chatgpt.py
new file mode 100644
index 0000000..52df90b
--- /dev/null
+++ b/chatgpt.py
@@ -0,0 +1,33 @@
+from openai import OpenAI
+
+class ChatGPT:
+
+    def __init__(self, api_key):
+        self.api_key = api_key
+        self.client = OpenAI(
+            api_key=api_key
+        )
+        self.role = {'role': 'system', 'content': 'You are a helpful personal assistent called Steve.'}
+
+    def completion(self, message: str) -> str:
+        messages = [
+            self.role,
+            {'role': 'user', 'content': message}
+        ]
+        response = self.client.chat.completions.create(
+            model='gpt-3.5-turbo',
+            messages=messages
+        )
+
+        return response.choices[0].message.content
+    
+    def image(self, prompt: str) -> str:
+        response = self.client.images.generate(
+            model="dall-e-3",
+            prompt=prompt,
+            size="1024x1024",
+            quality="standard",
+            n=1
+        )
+
+        return response.data[0].url
diff --git a/config.json b/config.json
index e66e4ae..25a3c64 100644
--- a/config.json
+++ b/config.json
@@ -1,3 +1,4 @@
 {
-    "token": "Your token here"
+    "token": "Your token here",
+    "openAI": "OpenAI API Key"
 }
\ No newline at end of file