XThomasBU commited on
Commit
8ee0d3b
·
1 Parent(s): 2c49234
code/app.py CHANGED
@@ -15,7 +15,7 @@ from modules.config.constants import (
15
  GITHUB_REPO,
16
  DOCS_WEBSITE,
17
  ALL_TIME_TOKENS_ALLOCATED,
18
- TOKENS_LEFT
19
  )
20
  from fastapi.middleware.cors import CORSMiddleware
21
  from fastapi.staticfiles import StaticFiles
@@ -209,7 +209,9 @@ async def cooldown(request: Request):
209
  user_info = await get_user_info_from_cookie(request)
210
  user_details = await get_user_details(user_info["email"])
211
  current_datetime = get_time()
212
- cooldown, cooldown_end_time = await check_user_cooldown(user_details, current_datetime)
 
 
213
  print(f"User in cooldown: {cooldown}")
214
  print(f"Cooldown end time: {cooldown_end_time}")
215
  if cooldown and "admin" not in get_user_role(user_info["email"]):
@@ -240,9 +242,13 @@ async def post_signin(request: Request):
240
  user_details.metadata["last_login"] = current_datetime
241
  # if new user, set the number of tries
242
  if "tokens_left" not in user_details.metadata:
243
- user_details.metadata["tokens_left"] = TOKENS_LEFT # set the number of tokens left for the new user
 
 
244
  if "all_time_tokens_allocated" not in user_details.metadata:
245
- user_details.metadata["all_time_tokens_allocated"] = ALL_TIME_TOKENS_ALLOCATED
 
 
246
  if "in_cooldown" not in user_details.metadata:
247
  user_details.metadata["in_cooldown"] = False
248
  await update_user_info(user_details)
@@ -252,8 +258,10 @@ async def post_signin(request: Request):
252
  ):
253
  cooldown, _ = await check_user_cooldown(user_details, current_datetime)
254
  if cooldown:
 
255
  return RedirectResponse("/cooldown")
256
  else:
 
257
  await reset_tokens_for_user(user_details)
258
 
259
  if user_info:
@@ -268,7 +276,9 @@ async def post_signin(request: Request):
268
  "role": role,
269
  "jwt_token": jwt_token,
270
  "tokens_left": user_details.metadata["tokens_left"],
271
- "all_time_tokens_allocated": user_details.metadata["all_time_tokens_allocated"],
 
 
272
  "total_tokens_allocated": ALL_TIME_TOKENS_ALLOCATED,
273
  },
274
  )
@@ -331,7 +341,6 @@ async def get_tokens_left(request: Request):
331
  except Exception as e:
332
  print(f"Error getting tokens left: {e}")
333
  return {"tokens_left": 0}
334
-
335
 
336
 
337
  mount_chainlit(app=app, target="main.py", path=CHAINLIT_PATH)
 
15
  GITHUB_REPO,
16
  DOCS_WEBSITE,
17
  ALL_TIME_TOKENS_ALLOCATED,
18
+ TOKENS_LEFT,
19
  )
20
  from fastapi.middleware.cors import CORSMiddleware
21
  from fastapi.staticfiles import StaticFiles
 
209
  user_info = await get_user_info_from_cookie(request)
210
  user_details = await get_user_details(user_info["email"])
211
  current_datetime = get_time()
212
+ cooldown, cooldown_end_time = await check_user_cooldown(
213
+ user_details, current_datetime
214
+ )
215
  print(f"User in cooldown: {cooldown}")
216
  print(f"Cooldown end time: {cooldown_end_time}")
217
  if cooldown and "admin" not in get_user_role(user_info["email"]):
 
242
  user_details.metadata["last_login"] = current_datetime
243
  # if new user, set the number of tries
244
  if "tokens_left" not in user_details.metadata:
245
+ user_details.metadata["tokens_left"] = (
246
+ TOKENS_LEFT # set the number of tokens left for the new user
247
+ )
248
  if "all_time_tokens_allocated" not in user_details.metadata:
249
+ user_details.metadata["all_time_tokens_allocated"] = (
250
+ ALL_TIME_TOKENS_ALLOCATED
251
+ )
252
  if "in_cooldown" not in user_details.metadata:
253
  user_details.metadata["in_cooldown"] = False
254
  await update_user_info(user_details)
 
258
  ):
259
  cooldown, _ = await check_user_cooldown(user_details, current_datetime)
260
  if cooldown:
261
+ user_details.metadata["in_cooldown"] = True
262
  return RedirectResponse("/cooldown")
263
  else:
264
+ user_details.metadata["in_cooldown"] = False
265
  await reset_tokens_for_user(user_details)
266
 
267
  if user_info:
 
276
  "role": role,
277
  "jwt_token": jwt_token,
278
  "tokens_left": user_details.metadata["tokens_left"],
279
+ "all_time_tokens_allocated": user_details.metadata[
280
+ "all_time_tokens_allocated"
281
+ ],
282
  "total_tokens_allocated": ALL_TIME_TOKENS_ALLOCATED,
283
  },
284
  )
 
341
  except Exception as e:
342
  print(f"Error getting tokens left: {e}")
343
  return {"tokens_left": 0}
 
344
 
345
 
346
  mount_chainlit(app=app, target="main.py", path=CHAINLIT_PATH)
code/main.py CHANGED
@@ -222,6 +222,7 @@ class Chatbot:
222
  "view_sources": llm_settings.get("view_sources"),
223
  "follow_up_questions": llm_settings.get("follow_up_questions"),
224
  }
 
225
  await cl.Message(
226
  author=SYSTEM,
227
  content="LLM settings have been updated. You can continue with your Query!",
@@ -373,9 +374,10 @@ class Chatbot:
373
  # if not, return message saying they have run out of tokens
374
  if user.metadata["tokens_left"] <= 0 and "admin" not in user.metadata["role"]:
375
  current_datetime = get_time()
376
- cooldown, cooldown_end_time = await check_user_cooldown(user, current_datetime)
 
 
377
  if cooldown:
378
- user.metadata["in_cooldown"] = True
379
  # get time left in cooldown
380
  # convert both to datetime objects
381
  cooldown_end_time = datetime.fromisoformat(cooldown_end_time).replace(
@@ -392,7 +394,6 @@ class Chatbot:
392
  minutes, seconds = divmod(remainder, 60)
393
  # Format the time as 00 hrs 00 mins 00 secs
394
  formatted_time = f"{hours:02} hrs {minutes:02} mins {seconds:02} secs"
395
- await update_user_info(user)
396
  await cl.Message(
397
  content=(
398
  "Ah, seems like you have run out of tokens...Click "
@@ -402,8 +403,12 @@ class Chatbot:
402
  ),
403
  author=SYSTEM,
404
  ).send()
 
 
405
  return
406
 
 
 
407
  llm_settings = cl.user_session.get("llm_settings", {})
408
  view_sources = llm_settings.get("view_sources", False)
409
  stream = llm_settings.get("stream_response", False)
@@ -483,15 +488,21 @@ class Chatbot:
483
  # # update user info with token count
484
  if "admin" not in user.metadata["role"]:
485
  user.metadata["tokens_left"] = user.metadata["tokens_left"] - token_count
486
- user.metadata["all_time_tokens_allocated"] = user.metadata["all_time_tokens_allocated"] - token_count
487
- await reset_tokens_for_user(user) # regenerate tokens for the user
 
 
488
  user.metadata["last_message_time"] = get_time()
489
  await update_user_info(user)
490
 
491
  tokens_left = user.metadata["tokens_left"]
492
  if tokens_left < 0:
493
  tokens_left = 0
494
- answer_with_sources += '\n\n<footer><span style="font-size: 0.8em; text-align: right; display: block;">Tokens Left: ' + str(tokens_left) + '</span></footer>\n'
 
 
 
 
495
 
496
  await cl.Message(
497
  content=answer_with_sources,
 
222
  "view_sources": llm_settings.get("view_sources"),
223
  "follow_up_questions": llm_settings.get("follow_up_questions"),
224
  }
225
+ print("Settings Dict: ", settings_dict)
226
  await cl.Message(
227
  author=SYSTEM,
228
  content="LLM settings have been updated. You can continue with your Query!",
 
374
  # if not, return message saying they have run out of tokens
375
  if user.metadata["tokens_left"] <= 0 and "admin" not in user.metadata["role"]:
376
  current_datetime = get_time()
377
+ cooldown, cooldown_end_time = await check_user_cooldown(
378
+ user, current_datetime
379
+ )
380
  if cooldown:
 
381
  # get time left in cooldown
382
  # convert both to datetime objects
383
  cooldown_end_time = datetime.fromisoformat(cooldown_end_time).replace(
 
394
  minutes, seconds = divmod(remainder, 60)
395
  # Format the time as 00 hrs 00 mins 00 secs
396
  formatted_time = f"{hours:02} hrs {minutes:02} mins {seconds:02} secs"
 
397
  await cl.Message(
398
  content=(
399
  "Ah, seems like you have run out of tokens...Click "
 
403
  ),
404
  author=SYSTEM,
405
  ).send()
406
+ user.metadata["in_cooldown"] = True
407
+ await update_user_info(user)
408
  return
409
 
410
+ user.metadata["in_cooldown"] = False
411
+
412
  llm_settings = cl.user_session.get("llm_settings", {})
413
  view_sources = llm_settings.get("view_sources", False)
414
  stream = llm_settings.get("stream_response", False)
 
488
  # # update user info with token count
489
  if "admin" not in user.metadata["role"]:
490
  user.metadata["tokens_left"] = user.metadata["tokens_left"] - token_count
491
+ user.metadata["all_time_tokens_allocated"] = (
492
+ user.metadata["all_time_tokens_allocated"] - token_count
493
+ )
494
+ await reset_tokens_for_user(user) # regenerate tokens for the user
495
  user.metadata["last_message_time"] = get_time()
496
  await update_user_info(user)
497
 
498
  tokens_left = user.metadata["tokens_left"]
499
  if tokens_left < 0:
500
  tokens_left = 0
501
+ answer_with_sources += (
502
+ '\n\n<footer><span style="font-size: 0.8em; text-align: right; display: block;">Tokens Left: '
503
+ + str(tokens_left)
504
+ + "</span></footer>\n"
505
+ )
506
 
507
  await cl.Message(
508
  content=answer_with_sources,
code/modules/chat_processor/helpers.py CHANGED
@@ -186,12 +186,13 @@ async def check_user_cooldown(user_info, current_time):
186
  if elapsed_time_in_seconds < COOLDOWN_TIME:
187
  return True, cooldown_end_time_iso # Return in ISO 8601 format
188
 
189
- user.metadata["in_cooldown"] = False
190
  # If not in cooldown, regenerate tokens
191
  await reset_tokens_for_user(user_info)
192
 
193
  return False, None
194
 
 
195
  async def reset_tokens_for_user(user_info):
196
  user_info = convert_to_dict(user_info)
197
  last_message_time_str = user_info["metadata"].get("last_message_time")
@@ -212,29 +213,42 @@ async def reset_tokens_for_user(user_info):
212
 
213
  # Calculate how many tokens should have been regenerated proportionally
214
  if current_tokens < max_tokens:
215
- # Determine the time required to fully regenerate tokens from current state
216
  if current_tokens < 0:
217
- time_to_full_regen = REGEN_TIME * (1 + abs(current_tokens) / max_tokens) # more time to regenerate if tokens left is negative
 
 
218
  else:
219
- time_to_full_regen = REGEN_TIME * (1 - current_tokens / max_tokens) # less time to regenerate if tokens left is positive
 
 
220
 
221
  # Calculate the proportion of this time that has elapsed
222
- proportion_of_time_elapsed = min(elapsed_time_in_seconds / time_to_full_regen, 1.0)
 
 
 
223
 
224
  # Calculate the tokens to regenerate based on the elapsed proportion
225
- tokens_to_regenerate = int(proportion_of_time_elapsed * (max_tokens - current_tokens))
 
 
226
 
227
  # Ensure tokens_to_regenerate is positive and doesn't exceed the maximum
228
  if tokens_to_regenerate > 0:
229
  new_token_count = min(current_tokens + tokens_to_regenerate, max_tokens)
230
 
231
- print(f"\n\n Adding {tokens_to_regenerate} tokens to the user, Time left for full credits: {time_to_full_regen - elapsed_time_in_seconds} \n\n")
 
 
232
 
233
  # Update the user's token count
234
  user_info["metadata"]["tokens_left"] = new_token_count
235
 
236
  await update_user_info(user_info)
237
 
 
 
238
  async def get_thread_step_info(thread_id):
239
  step = await literal_client.api.get_step(thread_id)
240
  return step
 
186
  if elapsed_time_in_seconds < COOLDOWN_TIME:
187
  return True, cooldown_end_time_iso # Return in ISO 8601 format
188
 
189
+ user_info["metadata"]["in_cooldown"] = False
190
  # If not in cooldown, regenerate tokens
191
  await reset_tokens_for_user(user_info)
192
 
193
  return False, None
194
 
195
+
196
  async def reset_tokens_for_user(user_info):
197
  user_info = convert_to_dict(user_info)
198
  last_message_time_str = user_info["metadata"].get("last_message_time")
 
213
 
214
  # Calculate how many tokens should have been regenerated proportionally
215
  if current_tokens < max_tokens:
216
+ # Determine the time required to fully regenerate tokens from the current state
217
  if current_tokens < 0:
218
+ time_to_full_regen = REGEN_TIME * (
219
+ 1 + abs(current_tokens) / max_tokens
220
+ ) # more time to regenerate if tokens left is negative
221
  else:
222
+ time_to_full_regen = REGEN_TIME * (
223
+ 1 - current_tokens / max_tokens
224
+ ) # less time to regenerate if tokens left is positive
225
 
226
  # Calculate the proportion of this time that has elapsed
227
+ proportion_of_time_elapsed = elapsed_time_in_seconds / time_to_full_regen
228
+
229
+ # Ensure the proportion doesn't exceed 1.0
230
+ proportion_of_time_elapsed = min(proportion_of_time_elapsed, 1.0)
231
 
232
  # Calculate the tokens to regenerate based on the elapsed proportion
233
+ tokens_to_regenerate = int(
234
+ proportion_of_time_elapsed * (max_tokens - current_tokens)
235
+ )
236
 
237
  # Ensure tokens_to_regenerate is positive and doesn't exceed the maximum
238
  if tokens_to_regenerate > 0:
239
  new_token_count = min(current_tokens + tokens_to_regenerate, max_tokens)
240
 
241
+ print(
242
+ f"\n\n Adding {tokens_to_regenerate} tokens to the user, Time left for full credits: {max(0, time_to_full_regen - elapsed_time_in_seconds)} \n\n"
243
+ )
244
 
245
  # Update the user's token count
246
  user_info["metadata"]["tokens_left"] = new_token_count
247
 
248
  await update_user_info(user_info)
249
 
250
+
251
+
252
  async def get_thread_step_info(thread_id):
253
  step = await literal_client.api.get_step(thread_id)
254
  return step