minor fixes

Kav-K · Feb 6, 2023 · 1f40020 · 1f40020
1 parent 01fb514
commit 1f40020
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 20 deletions.
diff --git a/models/index_model.py b/models/index_model.py
@@ -417,11 +417,11 @@ async def compose_indexes(self, user_id, indexes, name, deep_compose):
                 ]
             llm_predictor = LLMPredictor(llm=OpenAI(model_name="text-davinci-003", max_tokens=-1))
             embedding_model = OpenAIEmbedding()
-            tree_index = GPTTreeIndex(
-                documents=documents,
-                llm_predictor=llm_predictor,
-                embed_model=embedding_model,
+
+            tree_index = await self.loop.run_in_executor(
+                None, partial(GPTTreeIndex, documents=documents, llm_predictor=llm_predictor, embed_model=embedding_model)
             )
+
             await self.usage_service.update_usage(llm_predictor.last_token_usage)
             await self.usage_service.update_usage(
                 embedding_model.last_token_usage, embeddings=True
@@ -447,10 +447,11 @@ async def compose_indexes(self, user_id, indexes, name, deep_compose):
                 ]
 
             embedding_model = OpenAIEmbedding()
-            # Add everything into a simple vector index
-            simple_index = GPTSimpleVectorIndex(
-                documents=documents, embed_model=embedding_model
+
+            simple_index = await self.loop.run_in_executor(
+                None, partial(GPTSimpleVectorIndex, documents=documents, embed_model=embedding_model)
             )
+
             await self.usage_service.update_usage(
                 embedding_model.last_token_usage, embeddings=True
             )
@@ -737,7 +738,7 @@ async def interaction_check(self, interaction: discord.Interaction) -> bool:
                 )
             else:
                 composing_message = await interaction.response.send_message(
-                    "Composing indexes, this may take a long time...",
+                    "Composing indexes, this may take a long time, you will be DMed when it's ready!",
                     ephemeral=True,
                     delete_after=120,
                 )
@@ -751,9 +752,18 @@ async def interaction_check(self, interaction: discord.Interaction) -> bool:
                     else True,
                 )
                 await interaction.followup.send(
-                    "Composed indexes", ephemeral=True, delete_after=10
+                    "Composed indexes", ephemeral=True, delete_after=180
                 )
 
+                # Try to direct message the user that their composed index is ready
+                try:
+                    await self.index_cog.bot.get_user(self.user_id).send(
+                        f"Your composed index is ready! You can load it with /index load now in the server."
+                    )
+                except discord.Forbidden:
+                    pass
+
+
                 try:
                     await composing_message.delete()
                 except:

diff --git a/models/search_model.py b/models/search_model.py
@@ -70,15 +70,15 @@ async def index_pdf(self, url) -> list[Document]:
         # Delete the temporary file
         return documents
 
-    async def get_links(self, query, search_scope=3):
+    async def get_links(self, query, search_scope=2):
         """Search the web for a query"""
         async with aiohttp.ClientSession() as session:
             async with session.get(
                 f"https://www.googleapis.com/customsearch/v1?key={self.google_search_api_key}&cx={self.google_search_engine_id}&q={query}"
             ) as response:
                 if response.status == 200:
                     data = await response.json()
-                    # Return a list of the top 5 links
+                    # Return a list of the top 2 links
                     return ([item["link"] for item in data["items"][:search_scope]], [
                         item["link"] for item in data["items"]
                     ])
@@ -156,22 +156,19 @@ async def search(self, query, user_api_key, search_scope, nodes):
                 traceback.print_exc()
 
         embedding_model = OpenAIEmbedding()
-        index = GPTSimpleVectorIndex(documents, embed_model=embedding_model)
+
+        index = await self.loop.run_in_executor(None, partial(GPTSimpleVectorIndex, documents, embed_model=embedding_model))
+
         await self.usage_service.update_usage(
             embedding_model.last_token_usage, embeddings=True
         )
 
         llm_predictor = LLMPredictor(llm=OpenAI(model_name="text-davinci-003", max_tokens=-1))
         # Now we can search the index for a query:
         embedding_model.last_token_usage = 0
-        response = index.query(
-            query,
-            verbose=True,
-            embed_model=embedding_model,
-            llm_predictor=llm_predictor,
-            similarity_top_k=nodes or DEFAULT_SEARCH_NODES,
-            text_qa_template=self.qaprompt,
-        )
+
+        response = await self.loop.run_in_executor(None, partial(index.query, query, verbose=True, embed_model=embedding_model, llm_predictor=llm_predictor, similarity_top_k=nodes or DEFAULT_SEARCH_NODES, text_qa_template=self.qaprompt))
+
         await self.usage_service.update_usage(llm_predictor.last_token_usage)
         await self.usage_service.update_usage(
             embedding_model.last_token_usage, embeddings=True