LLM/DocsGPT
1
0
Fork 0
mirror of https://github.com/arc53/DocsGPT.git synced 2025-11-29 16:43:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: user_api_key capturing

Browse Source
This commit is contained in:
Siddhant Rai
2024-04-16 15:31:11 +05:30
parent 333b6e60e1
commit af5e73c8cb
14 changed files with 69 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
application/usage.py
View File

@@ -9,11 +9,11 @@ db = mongo["docsgpt"]
usage_collection = db["token_usage"]
def update_token_usage(api_key, token_usage):
def update_token_usage(user_api_key, token_usage):
if "pytest" in sys.modules:
return
usage_data = {
"api_key": api_key,
"api_key": user_api_key,
"prompt_tokens": token_usage["prompt_tokens"],
"generated_tokens": token_usage["generated_tokens"],
"timestamp": datetime.now(),
@@ -27,7 +27,7 @@ def gen_token_usage(func):
self.token_usage["prompt_tokens"] += count_tokens(message["content"])
result = func(self, model, messages, stream, **kwargs)
self.token_usage["generated_tokens"] += count_tokens(result)
update_token_usage(self.api_key, self.token_usage)
update_token_usage(self.user_api_key, self.token_usage)
return result
return wrapper
@@ -44,6 +44,6 @@ def stream_token_usage(func):
yield r
for line in batch:
self.token_usage["generated_tokens"] += count_tokens(line)
update_token_usage(self.api_key, self.token_usage)
update_token_usage(self.user_api_key, self.token_usage)
return wrapper