jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

change return to dict for token cases

Browse source
This commit is contained in:
italojohnny 2024-05-29 16:01:54 -03:00
commit c1cf8fe564
1 changed files with 18 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

27
src/backend/base/langflow/base/models/model.py
View file

@ -53,19 +53,28 @@ class LCModelComponent(CustomComponent):
key in response_metadata["token_usage"] for key in inner_openai_keys
):
token_usage = response_metadata["token_usage"]
completion_tokens = token_usage["completion_tokens"]
prompt_tokens = token_usage["prompt_tokens"]
total_tokens = token_usage["total_tokens"]
finish_reason = response_metadata["finish_reason"]
status_message = f"Tokens:\nInput: {prompt_tokens}\nOutput: {completion_tokens}\nTotal Tokens: {total_tokens}\nStop Reason: {finish_reason}\nResponse: {content}"
status_message = {
"tokens": {
"input": token_usage["prompt_tokens"],
"output": token_usage["completion_tokens"],
"total": token_usage["total_tokens"],
"stop_reason": response_metadata["finish_reason"],
"response": content,
}
}
elif all(key in response_metadata for key in anthropic_keys) and all(
key in response_metadata["usage"] for key in inner_anthropic_keys
):
usage = response_metadata["usage"]
input_tokens = usage["input_tokens"]
output_tokens = usage["output_tokens"]
stop_reason = response_metadata["stop_reason"]
status_message = f"Tokens:\nInput: {input_tokens}\nOutput: {output_tokens}\nStop Reason: {stop_reason}\nResponse: {content}"
status_message = {
"tokens": {
"input": usage["input_tokens"],
"output": usage["output_tokens"],
"stop_reason": response_metadata["stop_reason"],
"response": content,
}
}
else:
status_message = f"Response: {content}"
else: