Feature add openai new models 20240125

council/llm/openai_token_counter.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import logging
 import tiktoken
 
@@ -13,6 +14,7 @@ class OpenAITokenCounter(LLMessageTokenCounterBase):
     """
     See https://github.com/openai/openai-python/blob/main/chatml.md for information on
         how messages are converted to tokens.
+        https://platform.openai.com/docs/models/overview for tokens
     """
 
     def __init__(
@@ -105,18 +107,20 @@ def token_limit(self) -> int:
         return self._limit
 
     @staticmethod
-    def from_model(model: str) -> Optional["OpenAITokenCounter"]:
+    def from_model(model: str) -> Optional[OpenAITokenCounter]:
         try:
             encoding = tiktoken.encoding_for_model(model)
         except KeyError:
             logger.warning(f"model {model} not found. Using cl100k_base encoding.")
             encoding = tiktoken.get_encoding("cl100k_base")
 
         if model in {
+            "gpt-3.5-turbo-0301",
             "gpt-3.5-turbo-0613",
+            "gpt-3.5-turbo-1106",
             "gpt-3.5-turbo-16k-0613",
         }:
-            tokens_limit = 16384 if "-16k-" in model else 4096
+            tokens_limit = 16384 if ("-16k-" in model) or ("-1106" in model) else 4096
             tokens_per_message = 3
             tokens_per_name = 1
         elif model in {
@@ -128,16 +132,27 @@ def from_model(model: str) -> Optional["OpenAITokenCounter"]:
             tokens_limit = 32768 if "-32k-" in model else 8192
             tokens_per_message = 3
             tokens_per_name = 1
+        elif model in {
+            "gpt-4-1106-preview",
+            "gpt-4-0125-preview",
+        }:
+            tokens_limit = 128000
+            tokens_per_message = 3
+            tokens_per_name = 1
         elif model == "gpt-3.5-turbo-0301":
             tokens_limit = 4096
             tokens_per_message = 4  # every message follows <|start|>{role/name}\n{content}<|end|>\n
             tokens_per_name = -1  # if there's a name, the role is omitted
-        elif "gpt-3.5-turbo" in model:
-            logger.warning("gpt-3.5-turbo may change over time. Returning num tokens assuming gpt-3.5-turbo-0613.")
-            return OpenAITokenCounter.from_model(model="gpt-3.5-turbo-0613")
-        elif "gpt-4" in model:
-            logger.warning("gpt-4 may change over time. Returning num tokens assuming gpt-4-0613.")
-            return OpenAITokenCounter.from_model(model="gpt-4-0613")
+        elif model == "gpt-3.5-turbo":
+            return OpenAITokenCounter._return_alias(model, "gpt-3.5-turbo-0613")
+        elif model == "gpt-3.5-turbo-16k":
+            return OpenAITokenCounter._return_alias(model, "gpt-3.5-turbo-16k-0613")
+        elif model == "gpt-4":
+            return OpenAITokenCounter._return_alias(model, "gpt-4-0613")
+        elif model == "gpt-4-32k":
+            return OpenAITokenCounter._return_alias(model, "gpt-4-32k-0613")
+        elif model == "gpt-4-turbo-preview":
+            return OpenAITokenCounter._return_alias(model, "gpt-4-0125-preview")
         else:
             return None
 
@@ -148,3 +163,8 @@ def from_model(model: str) -> Optional["OpenAITokenCounter"]:
             tokens_per_message=tokens_per_message,
             tokens_per_name=tokens_per_name,
         )
+
+    @staticmethod
+    def _return_alias(alias: str, last: str) -> OpenAITokenCounter:
+        logger.warning(f"{alias} may change over time. Returning num tokens assuming {last}.")
+        return OpenAITokenCounter.from_model(model=last)

tests/unit/llm/test_openai_token_counter.py

@@ -4,13 +4,18 @@
 
 
 class TestLlmOpenAI(unittest.TestCase):
-    def test_token_counter(self):
+    def test_token_counter_gpt_35(self):
         model = "gpt-3.5-turbo"
         counter = OpenAITokenCounter.from_model(model)
         messages = self._get_messages()
 
         self.assertEqual(129, counter.count_messages_token(messages))
 
+    def test_token_counter_gpt_4_turbo(self):
+        model = "gpt-4-turbo-preview"
+        counter = OpenAITokenCounter.from_model(model)
+        self.assertEqual(counter.token_limit, 128000)
+
     def test_token_counter_exception(self):
         model = "gpt-4"
         counter = OpenAITokenCounter.from_model(model)