From d923e02700a289ba70039ced6610e36670065361 Mon Sep 17 00:00:00 2001 From: Kaveen Kumarasinghe Date: Sat, 18 Feb 2023 22:26:34 -0500 Subject: [PATCH] Make search indexes saved and usable later with /index --- README.md | 2 +- cogs/commands.py | 12 ++++++++++-- cogs/index_service_cog.py | 34 ++++++++++++++++++---------------- cogs/search_service_cog.py | 5 +++-- models/autocomplete_model.py | 17 +++++++++++++++++ models/index_model.py | 11 ++++++++--- models/search_model.py | 16 +++++++++++++++- 7 files changed, 72 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index 01f00c6..77a79cf 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ [![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) -# Overview +# Overview A robust, all-in-one GPT3 interface for Discord. Chat just like ChatGPT right inside Discord! Generate beautiful AI art using DALL-E 2! Automatically moderate your server using AI! Upload documents, videos, and files to get AI-assisted insights! A thorough integration with permanent conversation memory, automatic request retry, fault tolerance and reliability for servers of any scale, and much more. SUPPORT SERVER FOR BOT SETUP: https://discord.gg/WvAHXDMS7Q (You can try out the bot here also in a limited fashion) diff --git a/cogs/commands.py b/cogs/commands.py index 7f6eb6b..b24d3c5 100644 --- a/cogs/commands.py +++ b/cogs/commands.py @@ -543,10 +543,17 @@ class Commands(discord.Cog, name="Commands"): required=False, autocomplete=File_autocompleter.get_server_indexes, ) + @discord.option( + name="search_index", + description="Which search index file to load the index from", + required=False, + autocomplete=File_autocompleter.get_user_search_indexes, + ) async def load_index( - self, ctx: discord.ApplicationContext, user_index: str, server_index: str + self, ctx: discord.ApplicationContext, user_index: str, server_index: str, search_index: str ): - await self.index_cog.load_index_command(ctx, user_index, server_index) + await ctx.defer() + await self.index_cog.load_index_command(ctx, user_index, server_index, search_index) @add_to_group("index") @discord.slash_command( @@ -655,6 +662,7 @@ class Commands(discord.Cog, name="Commands"): nodes: int, response_mode: str, ): + await ctx.defer() await self.index_cog.query_command(ctx, query, nodes, response_mode) # diff --git a/cogs/index_service_cog.py b/cogs/index_service_cog.py index 8de367b..f025f62 100644 --- a/cogs/index_service_cog.py +++ b/cogs/index_service_cog.py @@ -25,6 +25,7 @@ class IndexService(discord.Cog, name="IndexService"): async def set_index_command( self, ctx, file: discord.Attachment = None, link: str = None ): + await ctx.defer() """Command handler to set a file as your personal index""" if not file and not link: await ctx.respond("Please provide a file or a link") @@ -44,7 +45,6 @@ class IndexService(discord.Cog, name="IndexService"): if not user_api_key: return - await ctx.defer(ephemeral=True) if file: await self.index_handler.set_file_index( ctx, file, user_api_key=user_api_key @@ -56,6 +56,7 @@ class IndexService(discord.Cog, name="IndexService"): async def set_discord_command(self, ctx, channel: discord.TextChannel = None): """Command handler to set a channel as your personal index""" + await ctx.defer() user_api_key = None if USER_INPUT_API_KEYS: @@ -65,13 +66,12 @@ class IndexService(discord.Cog, name="IndexService"): if not user_api_key: return - await ctx.defer(ephemeral=True) await self.index_handler.set_discord_index( ctx, channel, user_api_key=user_api_key ) async def reset_command(self, ctx): - await ctx.defer(ephemeral=True) + await ctx.defer() try: self.index_handler.reset_indexes(ctx.user.id) await ctx.respond("Your indexes have been reset") @@ -83,6 +83,7 @@ class IndexService(discord.Cog, name="IndexService"): async def discord_backup_command(self, ctx): """Command handler to backup the entire server""" + await ctx.defer() user_api_key = None if USER_INPUT_API_KEYS: @@ -91,29 +92,32 @@ class IndexService(discord.Cog, name="IndexService"): ) if not user_api_key: return - - await ctx.defer(ephemeral=True) await self.index_handler.backup_discord(ctx, user_api_key=user_api_key) - async def load_index_command(self, ctx, user_index, server_index): - """Command handler to backup the entire server""" + async def load_index_command(self, ctx, user_index, server_index, search_index): + """Command handler to load indexes""" - if not user_index and not server_index: - await ctx.respond("Please provide a user or server index") + if not user_index and not server_index and not search_index: + await ctx.respond("Please provide a user or server or search index") return - if user_index and server_index: + if user_index and server_index or user_index and search_index or server_index and search_index: await ctx.respond( - "Please provide only one user index or server index. Only one or the other." + "Please only try to load one type of index. Either a user index, a server index or a search index." ) return + search = False if server_index: index = server_index server = True - else: + elif user_index: index = user_index server = False + else: + index = search_index + server = False + search = True user_api_key = None if USER_INPUT_API_KEYS: @@ -122,12 +126,11 @@ class IndexService(discord.Cog, name="IndexService"): ) if not user_api_key: return - - await ctx.defer(ephemeral=True) - await self.index_handler.load_index(ctx, index, server, user_api_key) + await self.index_handler.load_index(ctx, index, server, search, user_api_key) 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: user_api_key = await TextService.get_user_api_key( @@ -136,7 +139,6 @@ class IndexService(discord.Cog, name="IndexService"): if not user_api_key: return - await ctx.defer() await self.index_handler.query(ctx, query, response_mode, nodes, user_api_key) async def compose_command(self, ctx, name): diff --git a/cogs/search_service_cog.py b/cogs/search_service_cog.py index 5fefd4b..320469e 100644 --- a/cogs/search_service_cog.py +++ b/cogs/search_service_cog.py @@ -90,7 +90,9 @@ class SearchService(discord.Cog, name="SearchService"): redo=None, from_followup=None, ): - """Command handler for the translation command""" + """Command handler for the search command""" + await ctx.defer() if not redo else None + user_api_key = None if USER_INPUT_API_KEYS: user_api_key = await TextService.get_user_api_key( @@ -106,7 +108,6 @@ class SearchService(discord.Cog, name="SearchService"): await ctx.respond("The search service is not enabled.") return - await ctx.defer() if not redo else None try: response, refined_text = await self.model.search( diff --git a/models/autocomplete_model.py b/models/autocomplete_model.py index 6e2cc2d..2a8271b 100644 --- a/models/autocomplete_model.py +++ b/models/autocomplete_model.py @@ -183,3 +183,20 @@ class File_autocompleter: ] # returns the 25 first files from your current input except Exception: return ["No server indexes found, add an index"] + + async def get_user_search_indexes(ctx: discord.AutocompleteContext): + """get all files in the indexes folder""" + try: + return [ + file + for file in os.listdir( + EnvService.find_shared_file( + f"indexes/{str(ctx.interaction.user.id)}_search/" + ) + ) + if file.startswith(ctx.value.lower()) + ][ + :25 + ] # returns the 25 first files from your current input + except Exception: + return ["No user indexes found, add an index"] diff --git a/models/index_model.py b/models/index_model.py index 1b57e8b..870c470 100644 --- a/models/index_model.py +++ b/models/index_model.py @@ -343,11 +343,11 @@ class Index_handler: if response.status == 200: content_type = response.headers.get("content-type") else: - await ctx.respond("Failed to get link", ephemeral=True) + await ctx.respond("Failed to get link") return except Exception: traceback.print_exc() - await ctx.respond("Failed to get link", ephemeral=True) + await ctx.respond("Failed to get link") return # Check if the link contains youtube in it @@ -416,7 +416,7 @@ class Index_handler: traceback.print_exc() async def load_index( - self, ctx: discord.ApplicationContext, index, server, user_api_key + self, ctx: discord.ApplicationContext, index, server, search, user_api_key ): if not user_api_key: os.environ["OPENAI_API_KEY"] = self.openai_key @@ -428,6 +428,10 @@ class Index_handler: index_file = EnvService.find_shared_file( f"indexes/{ctx.guild.id}/{index}" ) + elif search: + index_file = EnvService.find_shared_file( + f"indexes/{ctx.user.id}_search/{index}" + ) else: index_file = EnvService.find_shared_file( f"indexes/{ctx.user.id}/{index}" @@ -438,6 +442,7 @@ class Index_handler: self.index_storage[ctx.user.id].queryable_index = index await ctx.respond("Loaded index") except Exception as e: + traceback.print_exc() await ctx.respond(e) async def compose_indexes(self, user_id, indexes, name, deep_compose): diff --git a/models/search_model.py b/models/search_model.py index 75c329a..80df00b 100644 --- a/models/search_model.py +++ b/models/search_model.py @@ -4,8 +4,9 @@ import random import re import tempfile import traceback -from datetime import datetime +from datetime import datetime, date from functools import partial +from pathlib import Path import discord from bs4 import BeautifulSoup @@ -47,6 +48,16 @@ class Search: self.openai_key = os.getenv("OPENAI_TOKEN") self.EMBED_CUTOFF = 2000 + def add_search_index(self, index, user_id, query): + # Create a folder called "indexes/{USER_ID}" if it doesn't exist already + Path(f"{app_root_path()}/indexes/{user_id}_search").mkdir(parents=True, exist_ok=True) + # Save the index to file under the user id + file = f"{query[:20]}_{date.today().month}_{date.today().day}" + + index.save_to_disk( + app_root_path() / "indexes" / f"{str(user_id)}_search" / f"{file}.json" + ) + def build_search_started_embed(self): embed = discord.Embed( title="Searching the web...", @@ -295,6 +306,9 @@ class Search: None, partial(GPTSimpleVectorIndex, documents, embed_model=embedding_model), ) + # save the index to disk if not a redo + if not redo: + self.add_search_index(index, ctx.user.id if isinstance(ctx, discord.ApplicationContext) else ctx.author.id, query) else: print("Doing a deep search") llm_predictor_deep = LLMPredictor(