Skip to content

Commit

Permalink
Bedrock usage metadata
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
cvanes authored Oct 17, 2024
1 parent 0161dee commit c3080b4
Showing 1 changed file with 4 additions and 10 deletions.
14 changes: 4 additions & 10 deletions src/ell/providers/bedrock.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
return converse_message

0 comments on commit c3080b4

Please sign in to comment.