parameterize discord indexing message limit and child branch factor for /query

cogs/commands.py

@@ -622,10 +622,16 @@ class Commands(discord.Cog, name="Commands"):
         required=False,
         input_type=discord.SlashCommandOptionType.channel,
     )
+    @discord.option(
+        name="message_limit",
+        description="The number of messages to index",
+        required=False,
+        input_type=discord.SlashCommandOptionType.integer,
+    )
     async def set_discord(
-        self, ctx: discord.ApplicationContext, channel: discord.TextChannel
+        self, ctx: discord.ApplicationContext, channel: discord.TextChannel, message_limit: int
     ):
-        await self.index_cog.set_discord_command(ctx, channel)
+        await self.index_cog.set_discord_command(ctx, channel, message_limit=message_limit)
 
     @add_to_group("index")
     @discord.slash_command(
@@ -634,8 +640,14 @@ class Commands(discord.Cog, name="Commands"):
         guild_ids=ALLOWED_GUILDS,
         checks=[Check.check_admin_roles(), Check.check_index_roles()],
     )
+    @discord.option(
+        name="message_limit",
+        description="The number of messages to index per channel",
+        required=False,
+        input_type=discord.SlashCommandOptionType.integer,
+    )
     @discord.guild_only()
-    async def discord_backup(self, ctx: discord.ApplicationContext):
+    async def discord_backup(self, ctx: discord.ApplicationContext, message_limit: int):
         await self.index_cog.discord_backup_command(ctx)
 
     @add_to_group("index")
@@ -661,15 +673,25 @@ class Commands(discord.Cog, name="Commands"):
         default="default",
         choices=["default", "compact", "tree_summarize"],
     )
+    @discord.option(
+        name="child_branch_factor",
+        description="Only for deep indexes, how deep to go, higher is expensive.",
+        required=False,
+        default=1,
+        min_value=1,
+        max_value=3,
+        input_type=discord.SlashCommandOptionType.integer,
+    )
     async def query(
         self,
         ctx: discord.ApplicationContext,
         query: str,
         nodes: int,
         response_mode: str,
+        child_branch_factor: int,
     ):
         await ctx.defer()
-        await self.index_cog.query_command(ctx, query, nodes, response_mode)
+        await self.index_cog.query_command(ctx, query, nodes, response_mode, child_branch_factor)
 
     #
     # DALLE commands

cogs/index_service_cog.py

@@ -56,7 +56,7 @@ class IndexService(discord.Cog, name="IndexService"):
                 ctx, link, user_api_key=user_api_key
             )
 
-    async def set_discord_command(self, ctx, channel: discord.TextChannel = None):
+    async def set_discord_command(self, ctx, channel: discord.TextChannel = None, message_limit: int = 2500):
         """Command handler to set a channel as your personal index"""
         await ctx.defer()
 
@@ -69,7 +69,7 @@ class IndexService(discord.Cog, name="IndexService"):
                 return
 
         await self.index_handler.set_discord_index(
-            ctx, channel, user_api_key=user_api_key
+            ctx, channel, user_api_key=user_api_key, message_limit=message_limit
         )
 
     async def reset_command(self, ctx):
@@ -83,7 +83,7 @@ class IndexService(discord.Cog, name="IndexService"):
                 "Something went wrong while resetting your indexes. Contact the server admin."
             )
 
-    async def discord_backup_command(self, ctx):
+    async def discord_backup_command(self, ctx, message_limit: int = 2500):
         """Command handler to backup the entire server"""
         await ctx.defer()
 
@@ -94,7 +94,7 @@ class IndexService(discord.Cog, name="IndexService"):
             )
             if not user_api_key:
                 return
-        await self.index_handler.backup_discord(ctx, user_api_key=user_api_key)
+        await self.index_handler.backup_discord(ctx, user_api_key=user_api_key, message_limit=message_limit)
 
     async def load_index_command(self, ctx, user_index, server_index, search_index):
         """Command handler to load indexes"""
@@ -137,7 +137,7 @@ class IndexService(discord.Cog, name="IndexService"):
                 return
         await self.index_handler.load_index(ctx, index, server, search, user_api_key)
 
-    async def query_command(self, ctx, query, nodes, response_mode):
+    async def query_command(self, ctx, query, nodes, response_mode, child_branch_factor):
         """Command handler to query your index"""
 
         user_api_key = None
@@ -153,7 +153,7 @@ class IndexService(discord.Cog, name="IndexService"):
             if await Moderation.simple_moderate_and_respond(query, ctx):
                 return
 
-        await self.index_handler.query(ctx, query, response_mode, nodes, user_api_key)
+        await self.index_handler.query(ctx, query, response_mode, nodes, user_api_key, child_branch_factor)
 
     async def compose_command(self, ctx, name):
         """Command handler to compose from your index"""

detailed_guides/CUSTOM-INDEXES.md

@@ -14,4 +14,6 @@ Indexes can be combined with other indexes through a composition. To combine ind
   
 You can also compose a singular index with itself with "Deep Compose", this will give you a more detailed version of the index, but will be costly and will sometimes take multiple minutes to compose. **Deep compositions are useless for very short documents!**  
 
+Doing a deep composition will also allow you to use the `child_branch_factor` parameter for `/index query`, increasing this past 1 will take a much longer time to query and will be much more expensive for large documents, so be wary.
+
 **When doing Deep Compositions, it's highly reccomended to keep the document size small, or only do deep compositions on single documents.** This is because a deep composition reorganizes the simple index into a tree structure and uses GPT3 to summarize different nodes of the tree, which will lead to high costs. For example, a deep composition of a 300 page lab manual and the contents of my personal website at https://kaveenk.com cost me $2 USD roughly.

models/index_model.py

@@ -48,7 +48,7 @@ SHORT_TO_LONG_CACHE = {}
 
 
 def get_and_query(
-    user_id, index_storage, query, response_mode, nodes, llm_predictor, embed_model
+    user_id, index_storage, query, response_mode, nodes, llm_predictor, embed_model, child_branch_factor
 ):
     index: [GPTSimpleVectorIndex, ComposableGraph] = index_storage[
         user_id
@@ -56,7 +56,7 @@ def get_and_query(
     if isinstance(index, GPTTreeIndex):
         response = index.query(
             query,
-            child_branch_factor=1,
+            child_branch_factor=child_branch_factor,
             llm_predictor=llm_predictor,
             embed_model=embed_model,
             use_async=True,
@@ -421,6 +421,7 @@ class Index_handler:
         ctx: discord.ApplicationContext,
         channel: discord.TextChannel,
         user_api_key,
+        message_limit: int = 2500,
     ):
         if not user_api_key:
             os.environ["OPENAI_API_KEY"] = self.openai_key
@@ -429,7 +430,7 @@ class Index_handler:
 
         try:
             document = await self.load_data(
-                channel_ids=[channel.id], limit=1000, oldest_first=False
+                channel_ids=[channel.id], limit=message_limit, oldest_first=False
             )
             embedding_model = OpenAIEmbedding()
             index = await self.loop.run_in_executor(
@@ -575,7 +576,7 @@ class Index_handler:
             simple_index.save_to_disk(f"indexes/{user_id}/{name}.json")
             self.index_storage[user_id].queryable_index = simple_index
 
-    async def backup_discord(self, ctx: discord.ApplicationContext, user_api_key):
+    async def backup_discord(self, ctx: discord.ApplicationContext, user_api_key, message_limit):
         if not user_api_key:
             os.environ["OPENAI_API_KEY"] = self.openai_key
         else:
@@ -586,7 +587,7 @@ class Index_handler:
             for c in ctx.guild.text_channels:
                 channel_ids.append(c.id)
             document = await self.load_data(
-                channel_ids=channel_ids, limit=3000, oldest_first=False
+                channel_ids=channel_ids, limit=message_limit, oldest_first=False
             )
             embedding_model = OpenAIEmbedding()
             index = await self.loop.run_in_executor(
@@ -617,6 +618,7 @@ class Index_handler:
         response_mode,
         nodes,
         user_api_key,
+        child_branch_factor,
     ):
         if not user_api_key:
             os.environ["OPENAI_API_KEY"] = self.openai_key
@@ -638,6 +640,7 @@ class Index_handler:
                     nodes,
                     llm_predictor,
                     embedding_model,
+                    child_branch_factor
                 ),
             )
             print("The last token usage was ", llm_predictor.last_token_usage)