From cd88cc163bb6ae0f2a185e26a7a1140bfebff2ac Mon Sep 17 00:00:00 2001
From: Rene Teigen <rene9917@gmail.com>
Date: Thu, 9 Mar 2023 07:23:12 +0000
Subject: [PATCH] Fix broken per user api key on query and compose

---
 models/index_model.py | 33 ++++++++++++++++++++-------------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/models/index_model.py b/models/index_model.py
index c1fcbe7..63b855d 100644
--- a/models/index_model.py
+++ b/models/index_model.py
@@ -166,13 +166,6 @@ class Index_handler:
     def __init__(self, bot, usage_service):
         self.bot = bot
         self.openai_key = os.getenv("OPENAI_TOKEN")
-        self.llm_predictor = LLMPredictor(
-            llm=OpenAIChat(
-                temperature=0,
-                model_name="gpt-3.5-turbo",
-                openai_api_key=self.openai_key,
-            )
-        )
         self.index_storage = defaultdict(IndexData)
         self.loop = asyncio.get_running_loop()
         self.usage_service = usage_service
@@ -752,6 +745,13 @@ class Index_handler:
             )
             index_objects.append(index)
 
+        llm_predictor = LLMPredictor(
+            llm=OpenAIChat(
+                temperature=0,
+                model_name="gpt-3.5-turbo"
+            )
+        )
+
         # For each index object, add its documents to a GPTTreeIndex
         if deep_compose:
             documents = []
@@ -790,14 +790,14 @@ class Index_handler:
                 partial(
                     GPTTreeIndex,
                     documents=documents,
-                    llm_predictor=self.llm_predictor,
+                    llm_predictor=llm_predictor,
                     embed_model=embedding_model,
                     use_async=True,
                 ),
             )
 
             await self.usage_service.update_usage(
-                self.llm_predictor.last_token_usage, chatgpt=True
+                llm_predictor.last_token_usage, chatgpt=True
             )
             await self.usage_service.update_usage(
                 embedding_model.last_token_usage, embeddings=True
@@ -914,6 +914,13 @@ class Index_handler:
         else:
             os.environ["OPENAI_API_KEY"] = user_api_key
 
+        llm_predictor = LLMPredictor(
+            llm=OpenAIChat(
+                temperature=0,
+                model_name="gpt-3.5-turbo"
+            )
+        )
+
         ctx_response = await ctx.respond(
             embed=EmbedStatics.build_index_query_progress_embed(query)
         )
@@ -930,14 +937,14 @@ class Index_handler:
                     query,
                     response_mode,
                     nodes,
-                    self.llm_predictor,
+                    llm_predictor,
                     embedding_model,
                     child_branch_factor,
                 ),
             )
-            print("The last token usage was ", self.llm_predictor.last_token_usage)
+            print("The last token usage was ", llm_predictor.last_token_usage)
             await self.usage_service.update_usage(
-                self.llm_predictor.last_token_usage, chatgpt=True
+                llm_predictor.last_token_usage, chatgpt=True
             )
             await self.usage_service.update_usage(
                 embedding_model.last_token_usage, embeddings=True
@@ -946,7 +953,7 @@ class Index_handler:
             try:
                 total_price = round(
                     await self.usage_service.get_price(
-                        self.llm_predictor.last_token_usage, chatgpt=True
+                        llm_predictor.last_token_usage, chatgpt=True
                     )
                     + await self.usage_service.get_price(
                         embedding_model.last_token_usage, embeddings=True