Make search indexes saved and usable later with /index

Kaveen Kumarasinghe 2 years ago
parent 3998577341
commit d923e02700

@ -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)

@ -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)
#

@ -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):

@ -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(

@ -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"]

@ -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):

@ -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(

Loading…
Cancel
Save