Skip to content

Commit

Permalink
Merge pull request #1 from cvanes/fix-bedrock-usage-metadata
Browse files Browse the repository at this point in the history
Bedrock usage metadata
  • Loading branch information
cvanes authored Oct 17, 2024
2 parents 0161dee + c3080b4 commit 237c8d4
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 237c8d4

Please sign in to comment.