update bedrock claude usage

This commit is contained in:
better629 2025-03-17 11:22:32 +08:00
parent 58608a9422
commit 1b8df6ff1d
4 changed files with 15 additions and 3 deletions

View file

@ -1,11 +1,12 @@
import json
from abc import ABC, abstractmethod
from typing import Union
from typing import Optional, Union
class BaseBedrockProvider(ABC):
# to handle different generation kwargs
max_tokens_field_name = "max_tokens"
usage: Optional[dict] = None
def __init__(self, reasoning: bool = False, reasoning_max_token: int = 4000):
self.reasoning = reasoning

View file

@ -68,8 +68,12 @@ class AnthropicProvider(BaseBedrockProvider):
elif delta_type == "signature_delta":
completions = ""
return reasoning, completions
else:
return False, ""
elif rsp_dict["type"] == "message_stop":
self.usage = {
"prompt_tokens": rsp_dict.get("amazon-bedrock-invocationMetrics", {}).get("inputTokenCount", 0),
"completion_tokens": rsp_dict.get("amazon-bedrock-invocationMetrics", {}).get("outputTokenCount", 0),
}
return False, ""
class CohereProvider(BaseBedrockProvider):

View file

@ -130,6 +130,9 @@ class BedrockLLM(BaseLLM):
collected_content = await self._get_stream_response_body(stream_response)
log_llm_stream("\n")
full_text = ("".join(collected_content)).lstrip()
if self.__provider.usage:
# if provider provide usage, update it
self._update_costs(self.__provider.usage, self.config.model)
return full_text
def _get_response_body(self, response) -> dict:

View file

@ -373,6 +373,10 @@ BEDROCK_TOKEN_COSTS = {
"anthropic.claude-3-haiku-20240307-v1:0:200k": {"prompt": 0.00025, "completion": 0.00125},
# currently (2024-4-29) only available at US West (Oregon) AWS Region.
"anthropic.claude-3-opus-20240229-v1:0": {"prompt": 0.015, "completion": 0.075},
"anthropic.claude-3-5-sonnet-20241022-v2:0": {"prompt": 0.003, "completion": 0.015},
"us.anthropic.claude-3-5-sonnet-20241022-v2:0": {"prompt": 0.003, "completion": 0.015},
"anthropic.claude-3-7-sonnet-20250219-v1:0": {"prompt": 0.003, "completion": 0.015},
"us.anthropic.claude-3-7-sonnet-20250219-v1:0": {"prompt": 0.003, "completion": 0.015},
"cohere.command-text-v14": {"prompt": 0.0015, "completion": 0.0015},
"cohere.command-text-v14:7:4k": {"prompt": 0.0015, "completion": 0.0015},
"cohere.command-light-text-v14": {"prompt": 0.0003, "completion": 0.0003},