Add setting top_k amount to the query command

Count LLM tokens for deep compose
Update presentation of cost in the console
Update readme
Set a prompthelper at 500 tokens output up from the default 256 I believe. ~4 characters per token so the discord limit is around 500 tokens.

README.md

@@ -105,9 +105,9 @@ This bot supports per-user custom indexes. This means that users can upload file
 
 `/index add file:<file> or link:<link>` - Use a document or use a link to create/add to your indexes. If you provide a youtube link, the transcript of the video will be used. If you provide a web url, the contents of the webpage will be used, if you provide an image, the image text will be extracted and used!
 
-`/index query query:<prompt>` - Query your current index for a given prompt. GPT will answer based on your current document/indedx
+`/index query query:<prompt> nodes:<number> response_mode:<mode>` - Query your current index for a given prompt. GPT will answer based on your current document/index. You can also set it to query over more nodes, further refining the output over each one. A description of the modes can be found <a href="https://gpt-index.readthedocs.io/en/latest/guides/usage_pattern.html#setting-response-mode">here</a>. They do not work for deep composed indexes
 
-`/index load index:<index>` - Load a previously created index to query
+`/index load user_index:<index> or server_index:<index>` - Load a previously created index you own yourself, or an index for the whole server.
 
 `/index compose` - Combine multiple saved indexes into one, or upgrade existing indexes into Deep Compositions.
 
@@ -115,7 +115,7 @@ This bot supports per-user custom indexes. This means that users can upload file
 
 `/index add_discord channel:<discord channel>` - Create an add an index based on a discord channel
 
-`/index discord_backup` - Use the last 3000 messages of every channel on your discord server as an index
+`/index discord_backup` - Use the last 3000 messages of every channel on your discord server as an index. Needs both an admin and a index role
 
 
 ### System and Settings

cogs/commands.py

@@ -539,7 +539,7 @@ class Commands(discord.Cog, name="Commands"):
     )
     @discord.option(
         name="server_index",
-        description="Which serever file to load the index from",
+        description="Which server file to load the index from",
         required=False,
         autocomplete=File_autocompleter.get_server_indexes,
     )
@@ -617,6 +617,7 @@ class Commands(discord.Cog, name="Commands"):
         name="discord_backup",
         description="Save an index made from the whole server",
         guild_ids=ALLOWED_GUILDS,
+        checks=[Check.check_admin_roles(), Check.check_index_roles()]
     )
     @discord.guild_only()
     async def discord_backup(self, ctx: discord.ApplicationContext):
@@ -628,6 +629,15 @@ class Commands(discord.Cog, name="Commands"):
     )
     @discord.guild_only()
     @discord.option(name="query", description="What to query the index", required=True)
+    @discord.option(
+        name="nodes",
+        description="How many nodes should the response be queried from, only non-deep indexes",
+        required=False,
+        default=1,
+        min_value=1,
+        max_value=3,
+        input_type=discord.SlashCommandOptionType.integer,
+    )
     @discord.option(
         name="response_mode",
         description="Response mode, doesn't work on deep composed indexes",
@@ -637,9 +647,9 @@ class Commands(discord.Cog, name="Commands"):
         choices=["default", "compact", "tree_summarize"],
     )
     async def query(
-        self, ctx: discord.ApplicationContext, query: str, response_mode: str
+        self, ctx: discord.ApplicationContext, query: str, nodes:int, response_mode: str
     ):
-        await self.index_cog.query_command(ctx, query, response_mode)
+        await self.index_cog.query_command(ctx, query, nodes, response_mode)
 
     #
     # DALLE commands

cogs/index_service_cog.py

@@ -126,7 +126,7 @@ class IndexService(discord.Cog, name="IndexService"):
         await ctx.defer(ephemeral=True)
         await self.index_handler.load_index(ctx, index, server, user_api_key)
 
-    async def query_command(self, ctx, query, response_mode):
+    async def query_command(self, ctx, query, nodes, response_mode):
         """Command handler to query your index"""
         user_api_key = None
         if USER_INPUT_API_KEYS:
@@ -137,7 +137,7 @@ class IndexService(discord.Cog, name="IndexService"):
                 return
 
         await ctx.defer()
-        await self.index_handler.query(ctx, query, response_mode, user_api_key)
+        await self.index_handler.query(ctx, query, response_mode, nodes, user_api_key)
 
     async def compose_command(self, ctx, name):
         """Command handler to compose from your index"""

models/index_model.py

@@ -13,6 +13,7 @@ from langchain import OpenAI
 
 from gpt_index.readers import YoutubeTranscriptReader
 from gpt_index.readers.schema.base import Document
+
 from gpt_index import (
     GPTSimpleVectorIndex,
     SimpleDirectoryReader,
@@ -25,6 +26,7 @@ from gpt_index import (
     MockLLMPredictor,
     LLMPredictor,
     QueryConfig,
+    PromptHelper,
     IndexStructType,
 )
 from gpt_index.readers.web import DEFAULT_WEBSITE_EXTRACTOR
@@ -36,14 +38,15 @@ from services.environment_service import EnvService, app_root_path
 SHORT_TO_LONG_CACHE = {}
 
 
-def get_and_query(user_id, index_storage, query, response_mode, llm_predictor):
+def get_and_query(user_id, index_storage, query, response_mode, nodes, llm_predictor):
     index: [GPTSimpleVectorIndex, ComposableGraph] = index_storage[
         user_id
     ].get_index_or_throw()
+    prompthelper = PromptHelper(4096, 500, 20)
     if isinstance(index, GPTTreeIndex):
-        response = index.query(query, verbose=True, child_branch_factor=2, llm_predictor=llm_predictor)
+        response = index.query(query, verbose=True, child_branch_factor=2, llm_predictor=llm_predictor, prompt_helper=prompthelper)
     else:
-        response = index.query(query, response_mode=response_mode, verbose=True, llm_predictor=llm_predictor)
+        response = index.query(query, response_mode=response_mode, verbose=True, llm_predictor=llm_predictor, similarity_top_k=nodes, prompt_helper=prompthelper)
     return response
 
 
@@ -308,7 +311,10 @@ class Index_handler:
                     for doc_id in [docmeta for docmeta in _index.docstore.docs.keys()]
                     if isinstance(_index.docstore.get_document(doc_id), Document)
                 ]
-            tree_index = GPTTreeIndex(documents=documents)
+            llm_predictor = LLMPredictor(llm=OpenAI(model_name="text-davinci-003"))
+            tree_index = GPTTreeIndex(documents=documents, llm_predictor=llm_predictor)
+            print("The last token usage was ", llm_predictor.last_token_usage)
+            await self.usage_service.update_usage(llm_predictor.last_token_usage)
 
             # Now we have a list of tree indexes, we can compose them
             if not name:
@@ -369,7 +375,7 @@ class Index_handler:
             traceback.print_exc()
 
     async def query(
-        self, ctx: discord.ApplicationContext, query: str, response_mode, user_api_key
+        self, ctx: discord.ApplicationContext, query: str, response_mode, nodes, user_api_key
     ):
         if not user_api_key:
             os.environ["OPENAI_API_KEY"] = self.openai_key
@@ -381,7 +387,7 @@ class Index_handler:
             response = await self.loop.run_in_executor(
                 None,
                 partial(
-                    get_and_query, ctx.user.id, self.index_storage, query, response_mode, llm_predictor
+                    get_and_query, ctx.user.id, self.index_storage, query, response_mode, nodes, llm_predictor
                 ),
             )
             print("The last token usage was ", llm_predictor.last_token_usage)

services/usage_service.py

@@ -18,7 +18,7 @@ class UsageService:
         tokens_used = int(tokens_used)
         price = (tokens_used / 1000) * 0.02
         usage = await self.get_usage()
-        print("The current usage is " + str(usage) + " credits")
+        print(f"Cost -> Old: {str(usage)} | New: {str(usage + float(price))}, used {str(float(price))} credits")
         # Do the same as above but with aiofiles
         async with aiofiles.open(self.usage_file_path, "w") as f:
             await f.write(str(usage + float(price)))