support o1-preview and o1-mini

binary-husky · Sep 13, 2024 · 0d0575a · 0d0575a
1 parent 4e041e1
commit 0d0575a
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/request_llms/bridge_all.py b/request_llms/bridge_all.py
@@ -248,6 +248,23 @@ def decode(self, *args, **kwargs):
         "token_cnt": get_token_num_gpt4,
     },
 
+    "o1-preview": {
+        "fn_with_ui": chatgpt_ui,
+        "fn_without_ui": chatgpt_noui,
+        "endpoint": openai_endpoint,
+        "max_token": 128000,
+        "tokenizer": tokenizer_gpt4,
+        "token_cnt": get_token_num_gpt4,
+    },
+    "o1-mini": {
+        "fn_with_ui": chatgpt_ui,
+        "fn_without_ui": chatgpt_noui,
+        "endpoint": openai_endpoint,
+        "max_token": 128000,
+        "tokenizer": tokenizer_gpt4,
+        "token_cnt": get_token_num_gpt4,
+    },
+
     "gpt-4-turbo": {
         "fn_with_ui": chatgpt_ui,
         "fn_without_ui": chatgpt_noui,

diff --git a/shared_utils/key_pattern_manager.py b/shared_utils/key_pattern_manager.py
@@ -74,7 +74,7 @@ def select_api_key(keys, llm_model):
     avail_key_list = []
     key_list = keys.split(',')
 
-    if llm_model.startswith('gpt-') or llm_model.startswith('one-api-'):
+    if llm_model.startswith('gpt-') or llm_model.startswith('one-api-') or llm_model.startswith('o1-'):
         for k in key_list:
             if is_openai_api_key(k): avail_key_list.append(k)