Merge pull request #80 from Kav-K/multi-tenancy

Multi tenancy api key support

Dockerfile

@@ -26,4 +26,5 @@ COPY --from=builder /install /usr/local/lib/python${PY_VERSION}/site-packages
 RUN mkdir -p /opt/gpt3discord/etc
 COPY gpt3discord.py /opt/gpt3discord/bin/
 COPY image_optimizer_pretext.txt conversation_starter_pretext.txt conversation_starter_pretext_minimal.txt /opt/gpt3discord/share/
+COPY openers /opt/gpt3discord/share/openers
 CMD ["python3", "/opt/gpt3discord/bin/gpt3discord.py"]

README.md

@@ -4,7 +4,7 @@
 
 [![PyPi version](https://badgen.net/pypi/v/gpt3discord/)](https://pypi.com/project/gpt3discord)
 [![Latest release](https://badgen.net/github/release/Kav-K/GPT3Discord)](https://github.com/Kav-K/GPT3Discord/releases)
-[![Maintenance](https://img.shields.io/badge/Maintained%3F-yes-green.svg)](https://GitHub.com/Naereen/StrapDown.js/graphs/commit-activity)
+[![Maintenance](https://img.shields.io/badge/Maintained%3F-yes-green.svg)](https://GitHub.com/Kav-K/GPT3Discord/graphs/commit-activity)
 [![GitHub license](https://img.shields.io/github/license/Kav-K/GPT3Discord)](https://github.com/Kav-K/GPT3Discord/blob/master/LICENSE)
 [![PRs Welcome](https://img.shields.io/badge/PRs-welcome-brightgreen.svg?style=flat-square)](http://makeapullrequest.com)
 
@@ -23,6 +23,9 @@
 A big shoutout to `CrypticHeaven-Lab` for hitting our first sponsorship goal!
 
 # Recent Notable Updates
+- **Allow each individual user to enter their own API Key!** - Each request that a user makes will be made using their own API key! Check out the User-Input API Key section in this README for more details.
+
+
 - **Permanent memory with embeddings and PineconeDB finished!** - An initial alpha version of permanent memory is now done! This allows you to chat with GPT3 infinitely and accurately, and save tokens, by using embeddings. *Please read the Permanent Memory section for more information!*
 
 
@@ -35,8 +38,6 @@ A big shoutout to `CrypticHeaven-Lab` for hitting our first sponsorship goal!
 - **AUTOMATIC CHAT SUMMARIZATION!** - When the context limit of a conversation is reached, the bot will use GPT3 itself to summarize the conversation to reduce the tokens, and continue conversing with you, this allows you to chat for a long time!
 
 
-- Custom conversation openers from https://github.com/f/awesome-chatgpt-prompts were integrated into the bot, check out `/gpt converse opener_file`! The bot now has built in support to make GPT3 behave like various personalities, such as a life coach, python interpreter, interviewer, text based adventure game, and much more!
-
 # Features
 - **Directly prompt GPT3 with `/gpt ask <prompt>`**
 
@@ -147,6 +148,26 @@ Moreover, an important thing to keep in mind is: pinecone indexes are currently
 
 Permanent memory using pinecone is still in alpha, I will be working on cleaning up this work, adding auto-clearing, and optimizing for stability and reliability, any help and feedback is appreciated (**add me on Discord Kaveen#0001 for pinecone help**)! If at any time you're having too many issues with pinecone, simply remove the `PINECONE_TOKEN` line in your `.env` file and the bot will revert to using conversation summarizations.
 
+# User-Input API Keys (Multi-key tenancy)
+This bot supports multi-user tenancy in regards to API keys. This means that, if you wanted, you could make it such that each user needs to enter their own API key in order to use commands that use GPT3 and DALLE.
+
+To enable this, add the following line to the end of your `.env` file:
+```env
+USER_INPUT_API_KEYS="True"
+```
+
+Then, restart the bot, and it will set up the system for everyone to input their own API keys. 
+
+The bot will use SQLite to store API keys for the users, each user's key will be saved with a USER_ID <> API_KEY mapping in SQLite, and will be persistent across restarts. All the data will be saved in a file called `user_key_db.sqlite` in the current working directory of the bot.
+
+With this feature enabled, any attempt to use a GPT3 or DALL-E command without a valid API key set for the user will pop up the following modal for them to enter their API key:
+<img src="https://i.imgur.com/ZDScoWk.png"/>
+
+Once the user enters their key, the bot will send a small test request to OpenAI to validate that the key indeed works, if not, it will tell the user to try again and tell them why it did not work.
+
+After the user's key is validated, they will be able to use GPT3 and DALLE commands.
+
+The Moderations service still uses the main API key defined in the `.env` file. Pinecone and discord-tokens are also per-host tokens, not per-user.
 
 # Step-by-Step Guides for GPT3Discord
 

cogs/draw_image_generation.py

@@ -10,13 +10,21 @@
 
 
 # We don't use the converser cog here because we want to be able to redo for the last images and text prompts at the same time
+from sqlitedict import SqliteDict
+
+from cogs.gpt_3_commands_and_converser import GPT3ComCon
 from models.env_service_model import EnvService
 from models.user_model import RedoUser
 
 redo_users = {}
 users_to_interactions = {}
 ALLOWED_GUILDS = EnvService.get_allowed_guilds()
 
+USER_INPUT_API_KEYS = EnvService.get_user_input_api_keys()
+USER_KEY_DB = None
+if USER_INPUT_API_KEYS:
+    USER_KEY_DB = SqliteDict("user_key_db.sqlite")
+
 
 class DrawDallEService(discord.Cog, name="DrawDallEService"):
     def __init__(
@@ -40,14 +48,15 @@ async def encapsulated_send(
         response_message=None,
         vary=None,
         draw_from_optimizer=None,
+        custom_api_key=None,
     ):
         await asyncio.sleep(0)
         # send the prompt to the model
         from_context = isinstance(ctx, discord.ApplicationContext)
 
         try:
             file, image_urls = await self.model.send_image_request(
-                ctx, prompt, vary=vary if not draw_from_optimizer else None
+                ctx, prompt, vary=vary if not draw_from_optimizer else None, custom_api_key=custom_api_key
             )
         except ValueError as e:
             (
@@ -87,7 +96,7 @@ async def encapsulated_send(
             )
 
             await result_message.edit(
-                view=SaveView(ctx, image_urls, self, self.converser_cog, result_message)
+                view=SaveView(ctx, image_urls, self, self.converser_cog, result_message, custom_api_key=custom_api_key)
             )
 
             self.converser_cog.users_to_interactions[user_id] = []
@@ -106,7 +115,7 @@ async def encapsulated_send(
                     file=file,
                 )
                 await message.edit(
-                    view=SaveView(ctx, image_urls, self, self.converser_cog, message)
+                    view=SaveView(ctx, image_urls, self, self.converser_cog, message, custom_api_key=custom_api_key)
                 )
             else:  # Varying case
                 if not draw_from_optimizer:
@@ -123,6 +132,7 @@ async def encapsulated_send(
                             self.converser_cog,
                             result_message,
                             True,
+                            custom_api_key=custom_api_key,
                         )
                     )
 
@@ -134,7 +144,7 @@ async def encapsulated_send(
                     )
                     await result_message.edit(
                         view=SaveView(
-                            ctx, image_urls, self, self.converser_cog, result_message
+                            ctx, image_urls, self, self.converser_cog, result_message, custom_api_key=custom_api_key
                         )
                     )
 
@@ -155,6 +165,12 @@ async def encapsulated_send(
     )
     @discord.option(name="prompt", description="The prompt to draw from", required=True)
     async def draw(self, ctx: discord.ApplicationContext, prompt: str):
+        user_api_key = None
+        if USER_INPUT_API_KEYS:
+            user_api_key = await GPT3ComCon.get_user_api_key(ctx.user.id, ctx)
+            if not user_api_key:
+                return
+
         await ctx.defer()
 
         user = ctx.user
@@ -163,7 +179,7 @@ async def draw(self, ctx: discord.ApplicationContext, prompt: str):
             return
 
         try:
-            asyncio.ensure_future(self.encapsulated_send(user.id, prompt, ctx))
+            asyncio.ensure_future(self.encapsulated_send(user.id, prompt, ctx, custom_api_key=user_api_key))
 
         except Exception as e:
             print(e)
@@ -226,6 +242,7 @@ def __init__(
         message,
         no_retry=False,
         only_save=None,
+        custom_api_key=None,
     ):
         super().__init__(
             timeout=3600 if not only_save else None
@@ -236,15 +253,16 @@ def __init__(
         self.no_retry = no_retry
         self.converser_cog = converser_cog
         self.message = message
+        self.custom_api_key = custom_api_key
         for x in range(1, len(image_urls) + 1):
             self.add_item(SaveButton(x, image_urls[x - 1]))
         if not only_save:
             if not no_retry:
-                self.add_item(RedoButton(self.cog, converser_cog=self.converser_cog))
+                self.add_item(RedoButton(self.cog, converser_cog=self.converser_cog, custom_api_key=self.custom_api_key))
             for x in range(1, len(image_urls) + 1):
                 self.add_item(
                     VaryButton(
-                        x, image_urls[x - 1], self.cog, converser_cog=self.converser_cog
+                        x, image_urls[x - 1], self.cog, converser_cog=self.converser_cog, custom_api_key=self.custom_api_key
                     )
                 )
 
@@ -270,12 +288,13 @@ async def on_timeout(self):
 
 
 class VaryButton(discord.ui.Button):
-    def __init__(self, number, image_url, cog, converser_cog):
+    def __init__(self, number, image_url, cog, converser_cog, custom_api_key):
         super().__init__(style=discord.ButtonStyle.blurple, label="Vary " + str(number))
         self.number = number
         self.image_url = image_url
         self.cog = cog
         self.converser_cog = converser_cog
+        self.custom_api_key = custom_api_key
 
     async def callback(self, interaction: discord.Interaction):
         user_id = interaction.user.id
@@ -318,6 +337,7 @@ async def callback(self, interaction: discord.Interaction):
                     interaction.message,
                     response_message=response_message,
                     vary=self.image_url,
+                    custom_api_key=self.custom_api_key,
                 )
             )
 
@@ -354,10 +374,11 @@ async def callback(self, interaction: discord.Interaction):
 
 
 class RedoButton(discord.ui.Button["SaveView"]):
-    def __init__(self, cog, converser_cog):
+    def __init__(self, cog, converser_cog, custom_api_key):
         super().__init__(style=discord.ButtonStyle.danger, label="Retry")
         self.cog = cog
         self.converser_cog = converser_cog
+        self.custom_api_key = custom_api_key
 
     async def callback(self, interaction: discord.Interaction):
         user_id = interaction.user.id
@@ -383,5 +404,5 @@ async def callback(self, interaction: discord.Interaction):
             self.converser_cog.users_to_interactions[user_id].append(message.id)
 
             asyncio.ensure_future(
-                self.cog.encapsulated_send(user_id, prompt, ctx, response_message)
+                self.cog.encapsulated_send(user_id, prompt, ctx, response_message, custom_api_key=self.custom_api_key)
             )