diff --git a/src/ell/providers/bedrock.py b/src/ell/providers/bedrock.py index 77bc57c9..e99ff75a 100644 --- a/src/ell/providers/bedrock.py +++ b/src/ell/providers/bedrock.py @@ -147,15 +147,9 @@ def translate_from_provider( logger(tracked_results[0].text) - # usage = call_result.response.usage.dict() if call_result.response.get('usage') else {} - # metadata = call_result.response.model_dump() - # del metadata["content"] - - # process metadata for ell - # XXX: Unify an ell metadata format for ell studio. - usage["prompt_tokens"] = usage.get("inputTokens", 0) - usage["completion_tokens"] = usage.get("outputTokens", 0) - usage["total_tokens"] = usage['prompt_tokens'] + usage['completion_tokens'] + usage["prompt_tokens"] = provider_response.get('usage').get("inputTokens", 0) + usage["completion_tokens"] = provider_response.get('usage').get("outputTokens", 0) + usage["total_tokens"] = usage['prompt_tokens'] + usage['completion_tokens'] metadata["usage"] = usage return tracked_results, metadata @@ -228,4 +222,4 @@ def message_to_bedrock_message_format(message: Message) -> Dict[str, Any]: content_block_to_bedrock_format(c) for c in message.content ])) } - return converse_message \ No newline at end of file + return converse_message