From f33939f4c0d4e87a88834eabe5e82a3bcd1e89cf Mon Sep 17 00:00:00 2001
From: Rene Teigen <rene9917@gmail.com>
Date: Thu, 2 Mar 2023 12:14:07 +0000
Subject: [PATCH] Fix adding discord index

---
 gpt3discord.py                |  2 +-
 models/embed_statics_model.py |  6 +++---
 models/index_model.py         | 18 ++++++++++++++++--
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/gpt3discord.py b/gpt3discord.py
index 8399028..3e133af 100644
--- a/gpt3discord.py
+++ b/gpt3discord.py
@@ -32,7 +32,7 @@ from services.environment_service import EnvService
 from models.openai_model import Model
 
 
-__version__ = "10.9.3"
+__version__ = "10.9.4"
 
 
 PID_FILE = Path("bot.pid")
diff --git a/models/embed_statics_model.py b/models/embed_statics_model.py
index 725bcaf..235c572 100644
--- a/models/embed_statics_model.py
+++ b/models/embed_statics_model.py
@@ -87,7 +87,7 @@ class EmbedStatics:
     def get_index_set_success_embed(price="Unknown"):
         embed = discord.Embed(
             title="Index Added",
-            description=f"This index can now be queried and loaded with `/index query` and `/index load`\n\n||Total cost: {round(float(price), 6)}||",
+            description=f"This index can now be queried and loaded with `/index query` and `/index load`\n\n||Total cost: {round(float(price), 6) if price != 'Unknown' else 'Unknown'}||",
             color=discord.Color.green(),
         )
         # thumbnail of https://i.imgur.com/7JF0oGD.png
@@ -141,7 +141,7 @@ class EmbedStatics:
     def get_index_compose_success_embed(price="Unknown"):
         embed = discord.Embed(
             title="Indexes Composed",
-            description=f"Indexes composed successfully, you can query and load this index with `/index query` and `/index load`\n\n||Total cost: {round(float(price), 6)}||",
+            description=f"Indexes composed successfully, you can query and load this index with `/index query` and `/index load`\n\n||Total cost: {round(float(price), 6) if price != 'Unknown' else 'Unknown'}||",
             color=discord.Color.green(),
         )
         # thumbnail of https://i.imgur.com/7JF0oGD.png
@@ -269,7 +269,7 @@ class EmbedStatics:
     def build_index_query_success_embed(query, price="Unknown"):
         embed = discord.Embed(
             title="Index Service",
-            description=f"Query:\n`{query}`\nThe index query was successful.\n\n||Total cost: {round(float(price), 6)}||",
+            description=f"Query:\n`{query}`\nThe index query was successful.\n\n||Total cost: {round(float(price), 6) if price != 'Unknown' else 'Unknown'}||",
             color=discord.Color.green(),
         )
         # thumbnail of https://i.imgur.com/7JF0oGD.png
diff --git a/models/index_model.py b/models/index_model.py
index 8ec4b31..17ac72a 100644
--- a/models/index_model.py
+++ b/models/index_model.py
@@ -517,11 +517,18 @@ class Index_handler:
             index = await self.loop.run_in_executor(
                 None, partial(self.index_discord, document, embedding_model)
             )
+            try:
+                price = await self.usage_service.get_price(
+                    embedding_model.last_token_usage, embeddings=True
+                )
+            except Exception:
+                traceback.print_exc()
+                price = "Unknown"
             await self.usage_service.update_usage(
                 embedding_model.last_token_usage, embeddings=True
             )
             self.index_storage[ctx.user.id].add_index(index, ctx.user.id, channel.name)
-            await ctx.respond(embed=EmbedStatics.get_index_set_success_embed())
+            await ctx.respond(embed=EmbedStatics.get_index_set_success_embed(price))
         except Exception as e:
             await ctx.respond(embed=EmbedStatics.get_index_set_failure_embed(str(e)))
             traceback.print_exc()
@@ -725,6 +732,13 @@ class Index_handler:
             await self.usage_service.update_usage(
                 embedding_model.last_token_usage, embeddings=True
             )
+            try:
+                price = await self.usage_service.get_price(
+                    embedding_model.last_token_usage, embeddings=True
+                )
+            except Exception:
+                traceback.print_exc()
+                price = "Unknown"
             Path(app_root_path() / "indexes" / str(ctx.guild.id)).mkdir(
                 parents=True, exist_ok=True
             )
@@ -735,7 +749,7 @@ class Index_handler:
                 / f"{ctx.guild.name.replace(' ', '-')}_{date.today().month}_{date.today().day}.json"
             )
 
-            await ctx.respond(embed=EmbedStatics.get_index_set_success_embed())
+            await ctx.respond(embed=EmbedStatics.get_index_set_success_embed(price))
         except Exception as e:
             await ctx.respond(embed=EmbedStatics.get_index_set_failure_embed((str(e))))
             traceback.print_exc()