diff --git a/cogs/commands.py b/cogs/commands.py index ca8e3b4..1183591 100644 --- a/cogs/commands.py +++ b/cogs/commands.py @@ -810,7 +810,14 @@ class Commands(discord.Cog, name="Commands"): guild_ids=ALLOWED_GUILDS, ) @discord.option(name="query", description="The query to search", required=True) - @discord.option(name="scope", description="How many top links to use for context", required=False, input_type=discord.SlashCommandOptionType.integer, max_value=8, min_value=1) + @discord.option( + name="scope", + description="How many top links to use for context", + required=False, + input_type=discord.SlashCommandOptionType.integer, + max_value=8, + min_value=1, + ) @discord.guild_only() async def search(self, ctx: discord.ApplicationContext, query: str, scope: int): await self.search_cog.search_command(ctx, query, scope) diff --git a/cogs/search_service_cog.py b/cogs/search_service_cog.py index 48fc950..e17282a 100644 --- a/cogs/search_service_cog.py +++ b/cogs/search_service_cog.py @@ -12,6 +12,7 @@ ALLOWED_GUILDS = EnvService.get_allowed_guilds() USER_INPUT_API_KEYS = EnvService.get_user_input_api_keys() USER_KEY_DB = EnvService.get_api_db() + class SearchService(discord.Cog, name="SearchService"): """Cog containing translation commands and retrieval of translation services""" @@ -36,7 +37,10 @@ class SearchService(discord.Cog, name="SearchService"): if not user_api_key: return - if not EnvService.get_google_search_api_key() or not EnvService.get_google_search_engine_id(): + if ( + not EnvService.get_google_search_api_key() + or not EnvService.get_google_search_engine_id() + ): await ctx.send("The search service is not enabled.") return diff --git a/models/search_model.py b/models/search_model.py index 03ea3f7..ba067a7 100644 --- a/models/search_model.py +++ b/models/search_model.py @@ -6,7 +6,12 @@ from functools import partial from bs4 import BeautifulSoup import aiohttp -from gpt_index import QuestionAnswerPrompt, GPTSimpleVectorIndex, BeautifulSoupWebReader, Document +from gpt_index import ( + QuestionAnswerPrompt, + GPTSimpleVectorIndex, + BeautifulSoupWebReader, + Document, +) from gpt_index.readers.web import DEFAULT_WEBSITE_EXTRACTOR from services.environment_service import EnvService @@ -30,6 +35,7 @@ class Search: "answer the question, say that you were unable to answer the question if there is not sufficient context to formulate a decisive answer. The search query was: {query_str}\n" ) self.openai_key = os.getenv("OPENAI_TOKEN") + def index_webpage(self, url) -> list[Document]: documents = BeautifulSoupWebReader( website_extractor=DEFAULT_WEBSITE_EXTRACTOR @@ -62,7 +68,9 @@ class Search: # Concatenate all the text for a given website into one string and save it into an array: documents = [] for link in links: - document = await self.loop.run_in_executor(None, partial(self.index_webpage, link)) + document = await self.loop.run_in_executor( + None, partial(self.index_webpage, link) + ) [documents.append(doc) for doc in document] index = GPTSimpleVectorIndex(documents) @@ -71,4 +79,3 @@ class Search: response = index.query(query, text_qa_template=self.qaprompt) return response -