This commit is contained in:
Wei-Jianan 2024-05-17 18:44:05 +08:00
parent 3338c304b5
commit e7c7b44a0b
2 changed files with 11 additions and 4 deletions

View file

@ -33,7 +33,14 @@ class AnthropicProvider(BaseBedrockProvider):
def get_request_body(self, messages: list[dict], generate_kwargs, *args, **kwargs) -> str:
system_message, user_messages = self._split_system_user_messages(messages)
body = json.dumps({"messages": user_messages, "anthropic_version": "bedrock-2023-05-31", "system": system_message, **generate_kwargs})
body = json.dumps(
{
"messages": user_messages,
"anthropic_version": "bedrock-2023-05-31",
"system": system_message,
**generate_kwargs,
}
)
return body
def _get_completion_from_dict(self, rsp_dict: dict) -> str:

View file

@ -132,7 +132,7 @@ class BedrockLLM(BaseLLM):
prompt_tokens = int(headers.get("x-amzn-bedrock-input-token-count", 0))
completion_tokens = int(headers.get("x-amzn-bedrock-output-token-count", 0))
usage = {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
}
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
}
return usage