Update models.py

models.py

@@ -45,55 +45,79 @@ class BaseTCOModel(ABC):
     def get_latency(self):
         return self.latency
 
-class OpenAIModel(BaseTCOModel):
+class OpenAIModelGPT4(BaseTCOModel):
 
     def __init__(self):
-        self.set_name("(SaaS) OpenAI")
+        self.set_name("(SaaS) OpenAI GPT4")
         self.latency = "15s" #Default value for GPT4
         super().__init__()
 
     def render(self):
-        def on_model_change(model):
-            
-            if model == "GPT-4":
-                self.latency = "15s"
-                return gr.Dropdown.update(choices=["8K", "32K"])
-            else:
-                self.latency = "5s"
-                return gr.Dropdown.update(choices=["4K", "16K"], value="4K")
-
-        def define_cost_per_token(model, context_length):
-            if model == "GPT-4" and context_length == "8K":
+        def define_cost_per_token(context_length):
+            if context_length == "8K":
                 cost_per_1k_input_tokens = 0.03
                 cost_per_1k_output_tokens = 0.06
-            elif model == "GPT-4" and context_length == "32K":
+            else:
                 cost_per_1k_input_tokens = 0.06
                 cost_per_1k_output_tokens = 0.12
-            elif model == "GPT-3.5" and context_length == "4K":
+            return cost_per_1k_input_tokens, cost_per_1k_output_tokens
+        
+        self.context_length = gr.Dropdown(["8K", "32K"], value="8K", interactive=True,
+                                          label="Context size",
+                                          visible=False, info="Number of tokens the model considers when processing text")
+        self.input_tokens_cost_per_second = gr.Number(0.03, visible=False,
+                                           label="($) Price/1K input prompt tokens",
+                                           interactive=False
+                                           )
+        self.output_tokens_cost_per_second = gr.Number(0.06, visible=False,
+                                           label="($) Price/1K output prompt tokens",
+                                           interactive=False
+                                           )
+        self.info = gr.Markdown("The cost per input and output tokens values are from OpenAI's [pricing web page](https://openai.com/pricing)", interactive=False, visible=False)
+        self.context_length.change(define_cost_per_token, inputs=self.context_length, outputs=[self.input_tokens_cost_per_second, self.output_tokens_cost_per_second])
+        
+        self.labor = gr.Number(0, visible=False, 
+                                label="($) Labor cost per month", 
+                                info="This is an estimate of the labor cost of the AI engineer in charge of deploying the model",
+                                interactive=True
+                                )
+
+    def compute_cost_per_token(self, input_tokens_cost_per_second, output_tokens_cost_per_second, labor):
+        cost_per_input_token = (input_tokens_cost_per_second / 1000) 
+        cost_per_output_token = (output_tokens_cost_per_second / 1000)
+
+        return cost_per_input_token, cost_per_output_token, labor
+
+class OpenAIModelGPT3_5(BaseTCOModel):
+
+    def __init__(self):
+        self.set_name("(SaaS) OpenAI GPT3.5 Turbo")
+        self.latency = "5s" #Default value for GPT3.5 Turbo
+        super().__init__()
+
+    def render(self):
+        def define_cost_per_token(context_length):
+            if context_length == "4K":
                 cost_per_1k_input_tokens = 0.0015
                 cost_per_1k_output_tokens = 0.002
             else:
                 cost_per_1k_input_tokens = 0.003
                 cost_per_1k_output_tokens = 0.004
             return cost_per_1k_input_tokens, cost_per_1k_output_tokens
-        
-        self.model = gr.Dropdown(["GPT-4", "GPT-3.5 Turbo"], value="GPT-4",
-                                 label="OpenAI models",
-                                 interactive=True, visible=False)
-        self.context_length = gr.Dropdown(["8K", "32K"], value="8K", interactive=True,
+
+        self.context_length = gr.Dropdown(choices=["4K", "16K"], value="4K", interactive=True,
                                           label="Context size",
                                           visible=False, info="Number of tokens the model considers when processing text")
-        self.input_tokens_cost_per_second = gr.Number(0.03, visible=False,
+        self.input_tokens_cost_per_second = gr.Number(0.0015, visible=False,
                                            label="($) Price/1K input prompt tokens",
                                            interactive=False
                                            )
-        self.output_tokens_cost_per_second = gr.Number(0.06, visible=False,
+        self.output_tokens_cost_per_second = gr.Number(0.002, visible=False,
                                            label="($) Price/1K output prompt tokens",
                                            interactive=False
                                            )
         self.info = gr.Markdown("The cost per input and output tokens values are from OpenAI's [pricing web page](https://openai.com/pricing)", interactive=False, visible=False)
-        self.model.change(on_model_change, inputs=self.model, outputs=self.context_length).then(define_cost_per_token, inputs=[self.model, self.context_length], outputs=[self.input_tokens_cost_per_second, self.output_tokens_cost_per_second])
-        self.context_length.change(define_cost_per_token, inputs=[self.model, self.context_length], outputs=[self.input_tokens_cost_per_second, self.output_tokens_cost_per_second])
+        self.context_length.change(define_cost_per_token, inputs=self.context_length, outputs=[self.input_tokens_cost_per_second, self.output_tokens_cost_per_second])
         
         self.labor = gr.Number(0, visible=False, 
                                 label="($) Labor cost per month",