Added new models' pricing. (#1)

* - imported functions, making module a bit easier to use;

* - updated pricing for new models;

* - removed unnecessary code;

* - fixed missing return type specification;

src/llm_utils/__init__.py

@@ -0,0 +1 @@
+from .llm_utils import *

src/llm_utils/llm_utils.py

@@ -1,7 +1,7 @@
 import textwrap
-
 import tiktoken
 
+
 # OpenAI specific.
 def count_tokens(model: str, string: str) -> int:
     """Returns the number of tokens in a text string."""
@@ -11,7 +11,7 @@ def count_tokens(model: str, string: str) -> int:
 
 
 # OpenAI specific.
-def calculate_cost(num_input_tokens: int, num_output_tokens: int, model_type: str):
+def calculate_cost(num_input_tokens: int, num_output_tokens: int, model_type: str) -> float:
     """
     Calculate the cost of processing a request based on model type.
 
@@ -23,10 +23,12 @@ def calculate_cost(num_input_tokens: int, num_output_tokens: int, model_type: st
     Returns:
         The cost of processing the request, in USD.
     """
-    # Latest pricing info from OpenAI (https://openai.com/pricing), as of Oct 3 2023.
+    # Latest pricing info from OpenAI (https://openai.com/pricing), as of November 9, 2023.
     PRICING_PER_1000 = [
+        ("gpt-3.5-turbo-1106", {"input": 0.001, "output": 0.002}),
         ("gpt-3.5-turbo-16k", {"input": 0.003, "output": 0.004}),
         ("gpt-3.5-turbo", {"input": 0.0015, "output": 0.002}),
+        ("gpt-4-1106-preview", {"input": 0.01, "output": 0.03}),
         ("gpt-4-32k", {"input": 0.06, "output": 0.12}),
         ("gpt-4", {"input": 0.03, "output": 0.06}),
     ]
@@ -36,7 +38,6 @@ def calculate_cost(num_input_tokens: int, num_output_tokens: int, model_type: st
             # Calculate total cost per token and total tokens.
             input_cost_per_token = pricing["input"] / 1000
             output_cost_per_token = pricing["output"] / 1000
-            total_tokens = num_input_tokens + num_output_tokens
 
             # Calculate cost for input and output separately.
             input_cost = num_input_tokens * input_cost_per_token

tests/test_llm_utils.py

@@ -1,6 +1,6 @@
 import unittest
 
-from llm_utils import llm_utils
+import llm_utils
 
 class TestStringMethods(unittest.TestCase):
     def test_calculate_cost(self):
@@ -9,6 +9,8 @@ def test_calculate_cost(self):
         self.assertAlmostEqual(llm_utils.calculate_cost(1000, 2000, "gpt-4"), 0.15)
         self.assertAlmostEqual(llm_utils.calculate_cost(1000, 2000, "gpt-4-32k"), 0.3)
         self.assertAlmostEqual(llm_utils.calculate_cost(1000, 2000, "gpt-4-32k"), 0.3)
+        self.assertAlmostEqual(llm_utils.calculate_cost(10000, 20000, "gpt-3.5-turbo-1106"), 0.05)
+        self.assertAlmostEqual(llm_utils.calculate_cost(10000, 2000, "gpt-4-1106-preview"), 0.16)
         self.assertRaises(ValueError, llm_utils.calculate_cost, 0, 0, "not-an-llm")
 
 if __name__ == '__main__':