LLM/DocsGPT
1
0
Fork 0
mirror of https://github.com/arc53/DocsGPT.git synced 2026-02-11 08:41:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: pass decoded_token to llm and retrievers

Browse Source
This commit is contained in:
Siddhant Rai
2025-03-18 23:46:02 +05:30
parent f4ab85a2bb
commit ab95d90284
9 changed files with 75 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
application/usage.py
View File

@@ -9,10 +9,15 @@ db = mongo["docsgpt"]
usage_collection = db["token_usage"]
def update_token_usage(user_api_key, token_usage):
def update_token_usage(decoded_token, user_api_key, token_usage):
if "pytest" in sys.modules:
return
if decoded_token:
user_id = decoded_token["sub"]
else:
user_id = None
usage_data = {
"user_id": user_id,
"api_key": user_api_key,
"prompt_tokens": token_usage["prompt_tokens"],
"generated_tokens": token_usage["generated_tokens"],
@@ -35,7 +40,7 @@ def gen_token_usage(func):
self.token_usage["generated_tokens"] += num_tokens_from_object_or_list(
result
)
update_token_usage(self.user_api_key, self.token_usage)
update_token_usage(self.decoded_token, self.user_api_key, self.token_usage)
return result
return wrapper
@@ -54,6 +59,6 @@ def stream_token_usage(func):
yield r
for line in batch:
self.token_usage["generated_tokens"] += num_tokens_from_string(line)
update_token_usage(self.user_api_key, self.token_usage)
update_token_usage(self.decoded_token, self.user_api_key, self.token_usage)
return wrapper