Add gpt-4o model

Source: https://platform.openai.com/docs/models/gpt-4o
mkydr · May 13, 2024 · aa3dc32 · aa3dc32
1 parent dc51eb0
commit aa3dc32
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/bot/openai_helper.py b/bot/openai_helper.py
@@ -28,7 +28,8 @@
 GPT_4_32K_MODELS = ("gpt-4-32k", "gpt-4-32k-0314", "gpt-4-32k-0613")
 GPT_4_VISION_MODELS = ("gpt-4-vision-preview",)
 GPT_4_128K_MODELS = ("gpt-4-1106-preview","gpt-4-0125-preview","gpt-4-turbo-preview")
-GPT_ALL_MODELS = GPT_3_MODELS + GPT_3_16K_MODELS + GPT_4_MODELS + GPT_4_32K_MODELS + GPT_4_VISION_MODELS + GPT_4_128K_MODELS
+GPT_4O_MODELS = ("gpt-4o")
+GPT_ALL_MODELS = GPT_3_MODELS + GPT_3_16K_MODELS + GPT_4_MODELS + GPT_4_32K_MODELS + GPT_4_VISION_MODELS + GPT_4_128K_MODELS + GPT_4O_MODELS
 
 
 def default_max_tokens(model: str) -> int:
@@ -42,7 +43,7 @@ def default_max_tokens(model: str) -> int:
         return base
     elif model in GPT_4_MODELS:
         return base * 2
-    elif model in GPT_3_16K_MODELS:    
+    elif model in GPT_3_16K_MODELS:
         if model == "gpt-3.5-turbo-1106":
             return 4096
         return base * 4
@@ -52,6 +53,8 @@ def default_max_tokens(model: str) -> int:
         return 4096
     elif model in GPT_4_128K_MODELS:
         return 4096
+    elif model in GPT_4O_MODELS:
+        return 4096
 
 
 def are_functions_available(model: str) -> bool:
@@ -634,6 +637,8 @@ def __max_model_tokens(self):
             return base * 31
         if self.config['model'] in GPT_4_128K_MODELS:
             return base * 31
+        if self.config['model'] in GPT_4O_MODELS:
+            return base * 31
         raise NotImplementedError(
             f"Max tokens for model {self.config['model']} is not implemented yet."
         )
@@ -654,7 +659,7 @@ def __count_tokens(self, messages) -> int:
         if model in GPT_3_MODELS + GPT_3_16K_MODELS:
             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 model in GPT_4_MODELS + GPT_4_32K_MODELS + GPT_4_VISION_MODELS + GPT_4_128K_MODELS:
+        elif model in GPT_4_MODELS + GPT_4_32K_MODELS + GPT_4_VISION_MODELS + GPT_4_128K_MODELS + GPT_4O_MODELS:
             tokens_per_message = 3
             tokens_per_name = 1
         else: