change return to dict for token cases
This commit is contained in:
parent
67442774cf
commit
c1cf8fe564
1 changed files with 18 additions and 9 deletions
|
|
@ -53,19 +53,28 @@ class LCModelComponent(CustomComponent):
|
||||||
key in response_metadata["token_usage"] for key in inner_openai_keys
|
key in response_metadata["token_usage"] for key in inner_openai_keys
|
||||||
):
|
):
|
||||||
token_usage = response_metadata["token_usage"]
|
token_usage = response_metadata["token_usage"]
|
||||||
completion_tokens = token_usage["completion_tokens"]
|
status_message = {
|
||||||
prompt_tokens = token_usage["prompt_tokens"]
|
"tokens": {
|
||||||
total_tokens = token_usage["total_tokens"]
|
"input": token_usage["prompt_tokens"],
|
||||||
finish_reason = response_metadata["finish_reason"]
|
"output": token_usage["completion_tokens"],
|
||||||
status_message = f"Tokens:\nInput: {prompt_tokens}\nOutput: {completion_tokens}\nTotal Tokens: {total_tokens}\nStop Reason: {finish_reason}\nResponse: {content}"
|
"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(
|
elif all(key in response_metadata for key in anthropic_keys) and all(
|
||||||
key in response_metadata["usage"] for key in inner_anthropic_keys
|
key in response_metadata["usage"] for key in inner_anthropic_keys
|
||||||
):
|
):
|
||||||
usage = response_metadata["usage"]
|
usage = response_metadata["usage"]
|
||||||
input_tokens = usage["input_tokens"]
|
status_message = {
|
||||||
output_tokens = usage["output_tokens"]
|
"tokens": {
|
||||||
stop_reason = response_metadata["stop_reason"]
|
"input": usage["input_tokens"],
|
||||||
status_message = f"Tokens:\nInput: {input_tokens}\nOutput: {output_tokens}\nStop Reason: {stop_reason}\nResponse: {content}"
|
"output": usage["output_tokens"],
|
||||||
|
"stop_reason": response_metadata["stop_reason"],
|
||||||
|
"response": content,
|
||||||
|
}
|
||||||
|
}
|
||||||
else:
|
else:
|
||||||
status_message = f"Response: {content}"
|
status_message = f"Response: {content}"
|
||||||
else:
|
else:
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue