LLM/DocsGPT
1
0
Fork 0
mirror of https://github.com/arc53/DocsGPT.git synced 2026-02-19 10:51:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

update: apply decorator to abstract classes

Browse Source
This commit is contained in:
Siddhant Rai
2024-04-15 18:57:28 +05:30
parent 262d160314
commit 590aa8b43f
3 changed files with 28 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
application/usage.py
View File

@@ -19,12 +19,10 @@ def update_token_usage(api_key, token_usage):
def gen_token_usage(func):
def wrapper(self, model, messages, *args, **kwargs):
context = messages[0]["content"]
user_question = messages[-1]["content"]
prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"
self.token_usage["prompt_tokens"] += count_tokens(prompt)
result = func(self, model, messages, *args, **kwargs)
def wrapper(self, model, messages, stream, **kwargs):
for message in messages:
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)
return result
@@ -33,13 +31,11 @@ def gen_token_usage(func):
def stream_token_usage(func):
def wrapper(self, model, messages, *args, **kwargs):
context = messages[0]["content"]
user_question = messages[-1]["content"]
prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"
self.token_usage["prompt_tokens"] += count_tokens(prompt)
def wrapper(self, model, messages, stream, **kwargs):
for message in messages:
self.token_usage["prompt_tokens"] += count_tokens(message["content"])
batch = []
result = func(self, model, messages, *args, **kwargs)
result = func(self, model, messages, stream, **kwargs)
for r in result:
batch.append(r)
yield r