Skip to content

Commit

Permalink
match sess-* key
Browse files Browse the repository at this point in the history
  • Loading branch information
binary-husky committed Feb 6, 2024
1 parent cdb5288 commit c27e559
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 1 deletion.
2 changes: 1 addition & 1 deletion shared_utils/key_pattern_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ def is_openai_api_key(key):
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}$", key)
API_MATCH_ORIGINAL = re.match(r"sk-[a-zA-Z0-9]{48}$|sess-[a-zA-Z0-9]{40}$", key)
return bool(API_MATCH_ORIGINAL)


Expand Down
41 changes: 41 additions & 0 deletions tests/test_key_pattern_manager.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
import unittest

def validate_path():
import os, sys

os.path.dirname(__file__)
root_dir_assume = os.path.abspath(os.path.dirname(__file__) + "/..")
os.chdir(root_dir_assume)
sys.path.append(root_dir_assume)


validate_path() # validate path so you can run from base directory

from shared_utils.key_pattern_manager import is_openai_api_key

class TestKeyPatternManager(unittest.TestCase):
def test_is_openai_api_key_with_valid_key(self):
key = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
self.assertTrue(is_openai_api_key(key))

key = "sx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
self.assertFalse(is_openai_api_key(key))

key = "sess-wg61ZafYHpNz7FFwIH7HGZlbVqUVaeV5tatHCWpl"
self.assertTrue(is_openai_api_key(key))

key = "sess-wg61ZafYHpNz7FFwIH7HGZlbVqUVa5tatHCWpl"
self.assertFalse(is_openai_api_key(key))


def test_is_openai_api_key_with_invalid_key(self):
key = "invalid_key"
self.assertFalse(is_openai_api_key(key))

def test_is_openai_api_key_with_custom_pattern(self):
# Assuming you have set a custom pattern in your configuration
key = "custom-pattern-key"
self.assertFalse(is_openai_api_key(key))

if __name__ == '__main__':
unittest.main()

0 comments on commit c27e559

Please sign in to comment.