barunsaha commited on
Commit
17a0c62
1 Parent(s): aa45cf2

Check API keys for proper format

Browse files
Files changed (2) hide show
  1. app.py +2 -1
  2. helpers/llm_helper.py +10 -3
app.py CHANGED
@@ -82,7 +82,8 @@ def are_all_inputs_valid(
82
  if not llm_helper.is_valid_llm_provider_model(selected_provider, selected_model, user_key):
83
  handle_error(
84
  'The LLM settings do not look correct. Make sure that an API key/access token'
85
- ' is provided if the selected LLM requires it.',
 
86
  False
87
  )
88
  return False
 
82
  if not llm_helper.is_valid_llm_provider_model(selected_provider, selected_model, user_key):
83
  handle_error(
84
  'The LLM settings do not look correct. Make sure that an API key/access token'
85
+ ' is provided if the selected LLM requires it. An API key should be 6-64 characters'
86
+ ' long, only containing alphanumeric characters, hyphens, and underscores.',
87
  False
88
  )
89
  return False
helpers/llm_helper.py CHANGED
@@ -17,6 +17,8 @@ from global_config import GlobalConfig
17
 
18
 
19
  LLM_PROVIDER_MODEL_REGEX = re.compile(r'\[(.*?)\](.*)')
 
 
20
  HF_API_HEADERS = {'Authorization': f'Bearer {GlobalConfig.HUGGINGFACEHUB_API_TOKEN}'}
21
  REQUEST_TIMEOUT = 35
22
 
@@ -70,9 +72,14 @@ def is_valid_llm_provider_model(provider: str, model: str, api_key: str) -> bool
70
  if not provider or not model or provider not in GlobalConfig.VALID_PROVIDERS:
71
  return False
72
 
73
- if provider in [GlobalConfig.PROVIDER_GOOGLE_GEMINI, GlobalConfig.PROVIDER_COHERE,]:
74
- if not api_key or len(api_key) < 5:
75
- return False
 
 
 
 
 
76
 
77
  return True
78
 
 
17
 
18
 
19
  LLM_PROVIDER_MODEL_REGEX = re.compile(r'\[(.*?)\](.*)')
20
+ # 6-64 characters long, only containing alphanumeric characters, hyphens, and underscores
21
+ API_KEY_REGEX = re.compile(r'^[a-zA-Z0-9\-_]{6,64}$')
22
  HF_API_HEADERS = {'Authorization': f'Bearer {GlobalConfig.HUGGINGFACEHUB_API_TOKEN}'}
23
  REQUEST_TIMEOUT = 35
24
 
 
72
  if not provider or not model or provider not in GlobalConfig.VALID_PROVIDERS:
73
  return False
74
 
75
+ if provider in [
76
+ GlobalConfig.PROVIDER_GOOGLE_GEMINI,
77
+ GlobalConfig.PROVIDER_COHERE,
78
+ ] and not api_key:
79
+ return False
80
+
81
+ if api_key:
82
+ return API_KEY_REGEX.match(api_key) is not None
83
 
84
  return True
85