import traceback
from pathlib import Path
import discord
from models . embed_statics_model import EmbedStatics
from services . environment_service import EnvService
from services . moderations_service import Moderation
from services . text_service import TextService
from models . index_model import Index_handler
USER_INPUT_API_KEYS = EnvService . get_user_input_api_keys ( )
USER_KEY_DB = EnvService . get_api_db ( )
PRE_MODERATE = EnvService . get_premoderate ( )
class IndexService ( discord . Cog , name = " IndexService " ) :
""" Cog containing gpt-index commands """
def __init__ (
self ,
bot ,
usage_service ,
) :
super ( ) . __init__ ( )
self . bot = bot
self . index_handler = Index_handler ( bot , usage_service )
async def rename_user_index_command ( self , ctx , user_index , new_name ) :
""" Command handler to rename a user index """
if not new_name :
await ctx . respond ( await EmbedStatics . get_index_rename_failure_embed ( user_index . split ( " / " ) [ - 1 ] , " None " , " Please provide a new name for this index " ) )
return
if await self . index_handler . rename_index (
ctx ,
f " indexes/ { ctx . user . id } / { user_index } " ,
f " indexes/ { ctx . user . id } / { new_name } " ,
) :
await ctx . respond ( embed = EmbedStatics . get_index_rename_success_embed ( user_index . split ( " / " ) [ - 1 ] , new_name ) )
else :
await ctx . respond ( embed = EmbedStatics . get_index_rename_failure_embed ( user_index . split ( " / " ) [ - 1 ] , new_name , " Please check the server console for more details. " ) )
async def rename_server_index_command ( self , ctx , server_index , new_name ) :
""" Command handler to rename a user index """
if not new_name :
await ctx . respond ( await EmbedStatics . get_index_rename_failure_embed ( server_index . split ( " / " ) [ - 1 ] , " None " , " Please provide a new name for this index " ) )
return
if await self . index_handler . rename_index (
ctx ,
f " indexes/ { ctx . guild . id } / { server_index } " ,
f " indexes/ { ctx . guild . id } / { new_name } " ,
) :
await ctx . respond ( embed = EmbedStatics . get_index_rename_success_embed ( server_index . split ( " / " ) [ - 1 ] , new_name ) )
else :
await ctx . respond ( embed = EmbedStatics . get_index_rename_failure_embed ( server_index . split ( " / " ) [ - 1 ] , new_name , " Please check the server console for more details. " ) )
async def rename_search_index_command ( self , ctx , search_index , new_name ) :
if not new_name :
await ctx . respond ( await EmbedStatics . get_index_rename_failure_embed ( search_index . split ( " / " ) [ - 1 ] , " None " , " Please provide a new name for this index " ) )
return
if await self . index_handler . rename_index (
ctx ,
f " indexes/ { ctx . user . id } _search/ { search_index } " ,
f " indexes/ { ctx . user . id } _search/ { new_name } " ,
) :
await ctx . respond ( embed = EmbedStatics . get_index_rename_success_embed ( search_index . split ( " / " ) [ - 1 ] , new_name ) )
else :
await ctx . respond ( embed = EmbedStatics . get_index_rename_failure_embed ( search_index . split ( " / " ) [ - 1 ] , new_name , " Please check the server console for more details. " ) )
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 " )
return
if file and link :
await ctx . respond (
" Please provide only one file or link. Only one or the other. "
)
return
user_api_key = None
if USER_INPUT_API_KEYS :
user_api_key = await TextService . get_user_api_key (
ctx . user . id , ctx , USER_KEY_DB
)
if not user_api_key :
return
if file :
await self . index_handler . set_file_index (
ctx , file , user_api_key = user_api_key
)
elif link :
await self . index_handler . set_link_index (
ctx , link , user_api_key = user_api_key
)
async def set_discord_command (
self , ctx , channel : discord . TextChannel = None , message_limit : int = 2500
) :
""" Command handler to set a channel as your personal index """
await ctx . defer ( )
user_api_key = None
if USER_INPUT_API_KEYS :
user_api_key = await TextService . get_user_api_key (
ctx . user . id , ctx , USER_KEY_DB
)
if not user_api_key :
return
await self . index_handler . set_discord_index (
ctx , channel , user_api_key = user_api_key , message_limit = message_limit
)
async def reset_command ( self , ctx ) :
await ctx . defer ( )
try :
self . index_handler . reset_indexes ( ctx . user . id )
await ctx . respond ( " Your indexes have been reset " )
except :
traceback . print_exc ( )
await ctx . respond (
" Something went wrong while resetting your indexes. Contact the server admin. "
)
async def discord_backup_command ( self , ctx , message_limit : int = 2500 ) :
""" Command handler to backup the entire server """
await ctx . defer ( )
user_api_key = None
if USER_INPUT_API_KEYS :
user_api_key = await TextService . get_user_api_key (
ctx . user . id , ctx , USER_KEY_DB
)
if not user_api_key :
return
await self . index_handler . backup_discord (
ctx , user_api_key = user_api_key , message_limit = message_limit
)
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 and not search_index :
await ctx . respond ( " Please provide a user or server or search index " )
return
if (
user_index
and server_index
or user_index
and search_index
or server_index
and search_index
) :
await ctx . respond (
" 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
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 :
user_api_key = await TextService . get_user_api_key (
ctx . user . id , ctx , USER_KEY_DB
)
if not user_api_key :
return
await self . index_handler . load_index ( ctx , index , server , search , user_api_key )
async def query_command (
self , ctx , query , nodes , response_mode , child_branch_factor
) :
""" Command handler to query your index """
user_api_key = None
if USER_INPUT_API_KEYS :
user_api_key = await TextService . get_user_api_key (
ctx . user . id , ctx , USER_KEY_DB
)
if not user_api_key :
return
# Check the opener for bad content.
if PRE_MODERATE :
if await Moderation . simple_moderate_and_respond ( query , ctx ) :
return
await self . index_handler . query (
ctx , query , response_mode , nodes , user_api_key , child_branch_factor
)
async def compose_command ( self , ctx , name ) :
""" Command handler to compose from your index """
user_api_key = None
if USER_INPUT_API_KEYS :
user_api_key = await TextService . get_user_api_key (
ctx . user . id , ctx , USER_KEY_DB
)
if not user_api_key :
return
await self . index_handler . compose ( ctx , name , user_api_key )