Added gpt-index commands

cogs/commands.py

@@ -26,6 +26,7 @@ class Commands(discord.Cog, name="Commands"):
         image_draw_cog,
         image_service_cog,
         moderations_cog,
+        index_cog,
         translations_cog=None,
         search_cog=None,
     ):
@@ -39,6 +40,7 @@ class Commands(discord.Cog, name="Commands"):
         self.image_draw_cog = image_draw_cog
         self.image_service_cog = image_service_cog
         self.moderations_cog = moderations_cog
+        self.index_cog = index_cog
         self.translations_cog = translations_cog
         self.search_cog = search_cog
 
@@ -67,6 +69,12 @@ class Commands(discord.Cog, name="Commands"):
         guild_ids=ALLOWED_GUILDS,
         checks=[Check.check_admin_roles()],
     )
+    index = discord.SlashCommandGroup(
+        name="index",
+        description="gpt-index commands",
+        guild_ids=ALLOWED_GUILDS,
+        checks=[Check.check_gpt_roles()],
+    )
 
     #
     # System commands
@@ -489,6 +497,35 @@ class Commands(discord.Cog, name="Commands"):
     async def end(self, ctx: discord.ApplicationContext):
         await self.converser_cog.end_command(ctx)
 
+    #
+    # Index commands
+    #
+
+    @add_to_group("index")
+    @discord.slash_command(
+        name="set",
+        description="Set an index to query from",
+        guild_ids=ALLOWED_GUILDS
+    )
+    @discord.guild_only()
+    @discord.option(name="file", description="A file to create the index from", required=True, input_type=discord.Attachment)
+    async def set(self, ctx:discord.ApplicationContext, file: discord.Attachment):
+        await self.index_cog.set_index_command(ctx, file)
+
+
+    @add_to_group("index")
+    @discord.slash_command(
+        name="query",
+        description="Query from your index",
+        guild_ids=ALLOWED_GUILDS
+    )
+    @discord.guild_only()
+    @discord.option(name="query", description="What to query the index", required=True)
+    async def query(self, ctx:discord.ApplicationContext, query: str):
+        await self.index_cog.query_command(ctx, query)
+
+
+
     #
     # DALLE commands
     #

cogs/index_service_cog.py

@@ -0,0 +1,42 @@
+import discord
+
+from services.environment_service import EnvService
+from services.text_service import TextService
+from models.index_model import Index_handler
+
+USER_INPUT_API_KEYS = EnvService.get_user_input_api_keys()
+USER_KEY_DB = EnvService.get_api_db()
+
+class IndexService(discord.Cog, name="IndexService"):
+    """Cog containing gpt-index commands"""
+    def __init__(
+        self,
+        bot,
+    ):
+        super().__init__()
+        self.bot = bot
+        self.index_handler = Index_handler()
+    
+    async def set_index_command(self, ctx, file: discord.Attachment):
+        """Command handler to set a file as your personal index"""
+
+        user_api_key = None
+        if USER_INPUT_API_KEYS:
+            user_api_key = await TextService.get_user_api_key(ctx.user.id, ctx, USER_KEY_DB)
+            if not user_api_key:
+                return
+
+        await ctx.defer(ephemeral=True)
+        await self.index_handler.set_index(ctx, file, user_api_key=user_api_key)
+
+
+    async def query_command(self, ctx, query):
+        """Command handler to query your index"""
+        user_api_key = None
+        if USER_INPUT_API_KEYS:
+            user_api_key = await TextService.get_user_api_key(ctx.user.id, ctx, USER_KEY_DB)
+            if not user_api_key:
+                return
+
+        await ctx.defer()
+        await self.index_handler.query(ctx, query, user_api_key=user_api_key)

gpt3discord.py

@@ -18,6 +18,7 @@ from cogs.prompt_optimizer_cog import ImgPromptOptimizer
 from cogs.moderations_service_cog import ModerationsService
 from cogs.commands import Commands
 from cogs.translation_service_cog import TranslationService
+from cogs.index_service_cog import IndexService
 from models.deepl_model import TranslationModel
 from services.health_service import HealthService
 
@@ -169,6 +170,12 @@ async def main():
         )
     )
 
+    bot.add_cog(
+        IndexService(
+            bot
+        )
+    )
+
     if EnvService.get_deepl_token():
         bot.add_cog(TranslationService(bot, TranslationModel()))
         print("The translation service is enabled.")
@@ -191,6 +198,7 @@ async def main():
             bot.get_cog("DrawDallEService"),
             bot.get_cog("ImgPromptOptimizer"),
             bot.get_cog("ModerationsService"),
+            bot.get_cog("IndexService"),
             bot.get_cog("TranslationService"),
             bot.get_cog("SearchService"),
         )

models/index_model.py

@@ -0,0 +1,55 @@
+import os
+import traceback
+import asyncio
+import tempfile
+from functools import partial
+import discord
+
+from gpt_index import GPTSimpleVectorIndex, SimpleDirectoryReader
+
+
+
+class Index_handler:
+    def __init__(self):
+        self.openai_key = os.getenv("OPENAI_TOKEN")
+        self.index_storage = {}
+    
+    def index_file(self, file):
+        document = SimpleDirectoryReader(file).load_data()
+        index = GPTSimpleVectorIndex(document)
+        return index
+    
+    
+    async def set_index(self, ctx: discord.ApplicationContext, file: discord.Attachment, user_api_key):
+        loop = asyncio.get_running_loop()
+
+        if not user_api_key:
+            os.environ["OPENAI_API_KEY"] = self.openai_key
+        else:
+            os.environ["OPENAI_API_KEY"] = user_api_key
+    
+        try:
+            temp_path = tempfile.TemporaryDirectory()
+            temp_file = tempfile.NamedTemporaryFile(suffix=".txt", dir=temp_path.name, delete=False)
+            await file.save(temp_file.name)
+            index = await loop.run_in_executor(None, partial(self.index_file, temp_path.name))
+            self.index_storage[ctx.user.id] = index
+            temp_path.cleanup()
+            await ctx.respond("Index set")
+        except Exception:
+            await ctx.respond("Failed to set index")
+            traceback.print_exc()
+
+    async def query(self, ctx: discord.ApplicationContext, query, user_api_key):
+        if not user_api_key:
+            os.environ["OPENAI_API_KEY"] = self.openai_key
+        else:
+            os.environ["OPENAI_API_KEY"] = user_api_key
+
+        if not self.index_storage[ctx.user.id]:
+            await ctx.respond("You need to set an index", ephemeral=True, delete_after=5)
+            return
+        
+        index: GPTSimpleVectorIndex = self.index_storage[ctx.user.id]
+        response = index.query(query, verbose=True)
+        await ctx.respond(f"Query response: {response}")

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
   "flask",
   "flask",
   "beautifulsoup4",
+  "gpt-index",
 ]
 dynamic = ["version"]
 [project.scripts]

requirements.txt

@@ -9,4 +9,5 @@ pinecone-client==2.1.0
 sqlitedict==2.1.0
 backoff==2.2.1
 flask==2.2.2
-beautifulsoup4==4.11.1
+beautifulsoup4==4.11.1
+gpt-index==0.2.16