update openai api key pattern

binary-husky · Aug 19, 2024 · e07caf7 · e07caf7
1 parent 4873e9d
commit e07caf7
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 4 deletions.
diff --git a/shared_utils/key_pattern_manager.py b/shared_utils/key_pattern_manager.py
@@ -8,13 +8,20 @@
 pj = os.path.join
 default_user_name = 'default_user'
 
-
+# match openai keys
+openai_regex = re.compile(
+    r"sk-[a-zA-Z0-9_-]{48}$|" +
+    r"sk-[a-zA-Z0-9_-]{92}$|" +
+    r"sk-proj-[a-zA-Z0-9_-]{48}$|"+
+    r"sk-proj-[a-zA-Z0-9_-]{124}$|"+
+    r"sess-[a-zA-Z0-9]{40}$"
+)
 def is_openai_api_key(key):
     CUSTOM_API_KEY_PATTERN = get_conf('CUSTOM_API_KEY_PATTERN')
     if len(CUSTOM_API_KEY_PATTERN) != 0:
         API_MATCH_ORIGINAL = re.match(CUSTOM_API_KEY_PATTERN, key)
     else:
-        API_MATCH_ORIGINAL = re.match(r"sk-[a-zA-Z0-9]{48}$|sk-proj-[a-zA-Z0-9]{48}$|sess-[a-zA-Z0-9]{40}$", key)
+        API_MATCH_ORIGINAL = openai_regex.match(key)
     return bool(API_MATCH_ORIGINAL)
 
 

diff --git a/tests/test_key_pattern_manager.py b/tests/test_key_pattern_manager.py
@@ -21,10 +21,32 @@ def test_is_openai_api_key_with_valid_key(self):
         key = "sx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
         self.assertFalse(is_openai_api_key(key))
 
-        key = "sess-wg61ZafYHpNz7FFwIH7HGZlbVqUVaeV5tatHCWpl"
+        key = "sess-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
         self.assertTrue(is_openai_api_key(key))
 
-        key = "sess-wg61ZafYHpNz7FFwIH7HGZlbVqUVa5tatHCWpl"
+        key = "sess-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
+        self.assertFalse(is_openai_api_key(key))
+
+        key = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx_xxxxxxxxxxxxxxxxxx-xxxxxxxxxxxxxxxxxxxx"
+        self.assertTrue(is_openai_api_key(key))
+        key = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-xxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxx-xxxxxxxxxxxxxxxxxxxx"
+        self.assertTrue(is_openai_api_key(key))
+        key = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-xxx-xxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxx-xxxxxx-xxxxxxxxxxxxx"
+        self.assertTrue(is_openai_api_key(key))
+        key = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-xxx-xxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxx-xxxxxxxxxxxxxxxxxx"
+        self.assertFalse(is_openai_api_key(key))
+        key = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-xxx-xxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxx-xxxxxxxxxxxxxxxxxxxxx"
+        self.assertFalse(is_openai_api_key(key))
+
+        key = "sk-proj-xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx_xxxxxxxxxxxxxxxxxxxxx-xxxxxxxx"
+        self.assertTrue(is_openai_api_key(key))
+        key = "sk-proj-xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxxxxxxxxxxxxxxxxxx-xxxxxxxx"
+        self.assertTrue(is_openai_api_key(key))
+        key = "sk-proj-xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxxxxxxxxxxxxxxx-xxxxxxxx"
+        self.assertFalse(is_openai_api_key(key))
+        key = "sk-proj-xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxxxxxxxxxxxxxxx-xxxxxxxxxxxxx"
+        self.assertFalse(is_openai_api_key(key))
+        key = "sk-proj-xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxxxxxxxxxxxxxxx-xxx啊xxxxxxx"
         self.assertFalse(is_openai_api_key(key))