open-source-mirrors
/
GPT3Discord
mirror of https://github.com/Kav-K/GPT3Discord
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Make search indexes saved and usable later with /index

Browse Source
Kaveen Kumarasinghe 2 years ago
parent 3998577341
commit d923e02700
7 changed files with 72 additions and 25 deletions
Show Stats Download Patch File Download Diff File

12
cogs/commands.py
Unescape View File

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

34
cogs/index_service_cog.py
Unescape View File

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

5
cogs/search_service_cog.py
Unescape View File

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

17
models/autocomplete_model.py
Unescape View File

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

11
models/index_model.py
Unescape View File

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

16
models/search_model.py
Unescape View File

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

Write Preview
Loading…
Cancel
Save