change return to dict for token cases

This commit is contained in:
italojohnny 2024-05-29 16:01:54 -03:00
commit c1cf8fe564

View file

@ -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: