From c3080b45d3048ba6723e31a02541e8b1d341f2f5 Mon Sep 17 00:00:00 2001 From: Chris van Es Date: Fri, 18 Oct 2024 00:53:48 +0100 Subject: [PATCH] Bedrock usage metadata For non streaming responses, Bedrock usage metadata was not being captured. Updated the metadata `dict` in line with how it was being handled for streaming responses. --- src/ell/providers/bedrock.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) 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