From 28727a03d841c5c950c8b7db129c17d7dec3bd8a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=8E=98=E6=9D=83=20=E9=A9=AC?= Date: Fri, 2 Feb 2024 14:15:17 +0800 Subject: [PATCH] feat: +azure pricing --- metagpt/provider/azure_openai_api.py | 34 ++++++++++++++++++++++++++-- metagpt/utils/token_counter.py | 9 +++++++- 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/metagpt/provider/azure_openai_api.py b/metagpt/provider/azure_openai_api.py index 99063f7cf..c215acf45 100644 --- a/metagpt/provider/azure_openai_api.py +++ b/metagpt/provider/azure_openai_api.py @@ -11,8 +11,10 @@ from openai._base_client import AsyncHttpxClientWrapper from openai.types import CompletionUsage from metagpt.configs.llm_config import LLMType +from metagpt.logs import logger from metagpt.provider.llm_provider_registry import register_provider from metagpt.provider.openai_api import OpenAILLM +from metagpt.utils import TOKEN_COSTS, count_message_tokens, count_string_tokens from metagpt.utils.exceptions import handle_exception @@ -43,8 +45,36 @@ class AzureOpenAILLM(OpenAILLM): return kwargs def _calc_usage(self, messages: list[dict], rsp: str) -> CompletionUsage: - return CompletionUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + usage = CompletionUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + if not self.config.calc_usage: + return usage + + model_name = "gpt-35-turbo" if "gpt-3" in self.model.lower() else "gpt-4-turbo-preview" + try: + usage.prompt_tokens = count_message_tokens(messages, model_name) + usage.completion_tokens = count_string_tokens(rsp, model_name) + except Exception as e: + logger.error(f"usage calculation failed: {e}") + + return usage @handle_exception def _update_costs(self, usage: CompletionUsage): - pass + if self.config.calc_usage and usage and self.cost_manager: + model_name = self._get_azure_model() + # More about pricing: https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/ + self.cost_manager.update_cost(usage.prompt_tokens, usage.completion_tokens, model_name) + + def _get_azure_model(self) -> str: + models = [i.lower() for i in TOKEN_COSTS.keys() if "azure" in i] + mappings = {i: set(i.split("-")) for i in models} + words = self.model.lower().split("-") + weights = [] + for k, v in mappings.items(): + count = 0 + for i in words: + if i in v: + count += 1 + weights.append((k, count)) + sorted_list = sorted(weights, key=lambda x: x[1], reverse=True) + return sorted_list[0][0] diff --git a/metagpt/utils/token_counter.py b/metagpt/utils/token_counter.py index 94506e373..c6bf2f0d7 100644 --- a/metagpt/utils/token_counter.py +++ b/metagpt/utils/token_counter.py @@ -34,9 +34,16 @@ TOKEN_COSTS = { "glm-3-turbo": {"prompt": 0.0, "completion": 0.0007}, # 128k version, prompt + completion tokens=0.005￥/k-tokens "glm-4": {"prompt": 0.0, "completion": 0.014}, # 128k version, prompt + completion tokens=0.1￥/k-tokens "gemini-pro": {"prompt": 0.00025, "completion": 0.0005}, + # Azure + "azure-gpt-3.5-turbo-4k": {"prompt": 0.0015, "completion": 0.002}, + "azure-gpt-3.5-turbo-16k": {"prompt": 0.003, "completion": 0.004}, + "azure-gpt-3.5-turbo-1106": {"prompt": 0.001, "completion": 0.002}, + "azure-gpt-4-turbo": {"prompt": 0.01, "completion": 0.03}, + "azure-gpt-4-turbo-vision": {"prompt": 0.01, "completion": 0.03}, + "azure-gpt-4-8k": {"prompt": 0.03, "completion": 0.06}, + "azure-gpt-4-32k": {"prompt": 0.06, "completion": 0.12}, } - TOKEN_MAX = { "gpt-3.5-turbo": 4096, "gpt-3.5-turbo-0301": 4096,