XThomasBU
commited on
Commit
·
8ee0d3b
1
Parent(s):
2c49234
fixes
Browse files- code/app.py +15 -6
- code/main.py +17 -6
- code/modules/chat_processor/helpers.py +21 -7
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(
|
|
|
|
|
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"] =
|
|
|
|
|
244 |
if "all_time_tokens_allocated" not in user_details.metadata:
|
245 |
-
user_details.metadata["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[
|
|
|
|
|
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(
|
|
|
|
|
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"] =
|
487 |
-
|
|
|
|
|
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 +=
|
|
|
|
|
|
|
|
|
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 |
-
|
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 * (
|
|
|
|
|
218 |
else:
|
219 |
-
time_to_full_regen = REGEN_TIME * (
|
|
|
|
|
220 |
|
221 |
# Calculate the proportion of this time that has elapsed
|
222 |
-
proportion_of_time_elapsed =
|
|
|
|
|
|
|
223 |
|
224 |
# Calculate the tokens to regenerate based on the elapsed proportion
|
225 |
-
tokens_to_regenerate = int(
|
|
|
|
|
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(
|
|
|
|
|
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
|