diff --git a/README.md b/README.md index 71bb218..717388c 100644 --- a/README.md +++ b/README.md @@ -105,9 +105,9 @@ This bot supports per-user custom indexes. This means that users can upload file `/index add file: or 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:` - Query your current index for a given prompt. GPT will answer based on your current document/indedx +`/index query query: nodes: response_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 here. They do not work for deep composed indexes -`/index load index:` - Load a previously created index to query +`/index load user_index: or server_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:` - 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 diff --git a/cogs/commands.py b/cogs/commands.py index 63acc24..d4f8f67 100644 --- a/cogs/commands.py +++ b/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 diff --git a/cogs/index_service_cog.py b/cogs/index_service_cog.py index ef0b4a2..8de367b 100644 --- a/cogs/index_service_cog.py +++ b/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""" diff --git a/models/index_model.py b/models/index_model.py index 21737d3..0eb0904 100644 --- a/models/index_model.py +++ b/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) diff --git a/services/usage_service.py b/services/usage_service.py index 53fec35..64bfb76 100644 --- a/services/usage_service.py +++ b/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)))