Merge pull request #1765 from better629/main

update bedrock claude usage

Author: better629

metagpt/provider/bedrock/base_provider.py

@@ -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

metagpt/provider/bedrock/bedrock_provider.py

@@ -68,8 +68,12 @@ def get_choice_text_from_stream(self, event) -> Union[bool, str]:
             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):

metagpt/provider/bedrock/utils.py

@@ -48,6 +48,8 @@
     "anthropic.claude-3-opus-20240229-v1:0": 4096,
     # Claude 3.5 Sonnet
     "anthropic.claude-3-5-sonnet-20240620-v1:0": 8192,
+    "anthropic.claude-3-5-sonnet-20241022-v2:0": 8192,
+    "us.anthropic.claude-3-5-sonnet-20241022-v2:0": 8192,
     # Claude 3.7 Sonnet
     "us.anthropic.claude-3-7-sonnet-20250219-v1:0": 131072,
     "anthropic.claude-3-7-sonnet-20250219-v1:0": 131072,

metagpt/provider/bedrock_api.py

@@ -130,6 +130,9 @@ async def _achat_completion_stream(self, messages: list[dict], timeout=USE_CONFI
         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:

metagpt/utils/token_counter.py

@@ -93,6 +93,8 @@
     "openai/o1-preview": {"prompt": 0.015, "completion": 0.06},
     "openai/o1-mini": {"prompt": 0.003, "completion": 0.012},
     "anthropic/claude-3-opus": {"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": {"prompt": 0.003, "completion": 0.015},
     "anthropic/claude-3.7-sonnet:beta": {"prompt": 0.003, "completion": 0.015},
     "anthropic/claude-3.7-sonnet:thinking": {"prompt": 0.003, "completion": 0.015},
@@ -371,6 +373,10 @@
     "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},
@@ -403,15 +409,24 @@
 }
 
 
+def count_claude_message_tokens(messages: list[dict], model: str) -> int:
+    # rough estimation for models newer than claude-2.1, needs api_key or auth_token
+    ac = anthropic.Client()
+    system_prompt = ""
+    new_messages = []
+    for msg in messages:
+        if msg.get("role") == "system":
+            system_prompt = msg.get("content")
+        else:
+            new_messages.append(msg)
+    num_tokens = ac.beta.messages.count_tokens(messages=new_messages, model=model, system=system_prompt)
+    return num_tokens.input_tokens
+
+
 def count_message_tokens(messages, model="gpt-3.5-turbo-0125"):
     """Return the number of tokens used by a list of messages."""
     if "claude" in model:
-        # rough estimation for models newer than claude-2.1
-        vo = anthropic.Client()
-        num_tokens = 0
-        for message in messages:
-            for key, value in message.items():
-                num_tokens += vo.count_tokens(str(value))
+        num_tokens = count_claude_message_tokens(messages, model)
         return num_tokens
     try:
         encoding = tiktoken.encoding_for_model(model)
@@ -500,8 +515,8 @@ def count_output_tokens(string: str, model: str) -> int:
         int: The number of tokens in the text string.
     """
     if "claude" in model:
-        vo = anthropic.Client()
-        num_tokens = vo.count_tokens(string)
+        messages = [{"role": "assistant", "content": string}]
+        num_tokens = count_claude_message_tokens(messages, model)
         return num_tokens
     try:
         encoding = tiktoken.encoding_for_model(model)

setup.py

@@ -98,7 +98,7 @@ def run(self):
 
 setup(
     name="metagpt",
-    version="0.8.1",
+    version="1.0.0",
     description="The Multi-Agent Framework",
     long_description=long_description,
     long_description_content_type="text/markdown",