diff --git a/metagpt/configs/llm_config.py b/metagpt/configs/llm_config.py index 4ed741af3..36f5d7ae7 100644 --- a/metagpt/configs/llm_config.py +++ b/metagpt/configs/llm_config.py @@ -38,13 +38,18 @@ class LLMConfig(YamlModel): Optional Fields in pydantic: https://docs.pydantic.dev/latest/migration/#required-optional-and-nullable-fields """ - api_key: str + api_key: str = "sk-" api_type: LLMType = LLMType.OPENAI base_url: str = "https://api.openai.com/v1" api_version: Optional[str] = None model: Optional[str] = None # also stands for DEPLOYMENT_NAME + # For Cloud Service Provider like Baidu/ Alibaba + access_key: Optional[str] = None + secret_key: Optional[str] = None + endpoint: Optional[str] = None # for self-deployed model on the cloud + # For Spark(Xunfei), maybe remove later app_id: Optional[str] = None api_secret: Optional[str] = None diff --git a/metagpt/provider/base_llm.py b/metagpt/provider/base_llm.py index b144471b5..2f57b15aa 100644 --- a/metagpt/provider/base_llm.py +++ b/metagpt/provider/base_llm.py @@ -11,11 +11,12 @@ from abc import ABC, abstractmethod from typing import Optional, Union from openai import AsyncOpenAI +from pydantic import BaseModel from metagpt.configs.llm_config import LLMConfig from metagpt.logs import logger from metagpt.schema import Message -from metagpt.utils.cost_manager import CostManager +from metagpt.utils.cost_manager import CostManager, Costs class BaseLLM(ABC): @@ -67,6 +68,28 @@ class BaseLLM(ABC): def _default_system_msg(self): return self._system_msg(self.system_prompt) + def _update_costs(self, usage: Union[dict, BaseModel], model: str = None, local_calc_usage: bool = True): + """update each request's token cost + Args: + model (str): model name or in some scenarios called endpoint + local_calc_usage (bool): some models don't calculate usage, it will overwrite LLMConfig.calc_usage + """ + calc_usage = self.config.calc_usage and local_calc_usage + model = model if model else self.model + usage = usage.model_dump() if isinstance(usage, BaseModel) else usage + if calc_usage and self.cost_manager: + try: + prompt_tokens = int(usage.get("prompt_tokens", 0)) + completion_tokens = int(usage.get("completion_tokens", 0)) + self.cost_manager.update_cost(prompt_tokens, completion_tokens, model) + except Exception as e: + logger.error(f"{self.__class__.__name__} updats costs failed! exp: {e}") + + def get_costs(self) -> Costs: + if not self.cost_manager: + return Costs(0, 0, 0, 0) + return self.cost_manager.get_costs() + async def aask( self, msg: str, diff --git a/metagpt/utils/cost_manager.py b/metagpt/utils/cost_manager.py index c3aa7323f..efff07ae1 100644 --- a/metagpt/utils/cost_manager.py +++ b/metagpt/utils/cost_manager.py @@ -29,6 +29,7 @@ class CostManager(BaseModel): total_budget: float = 0 max_budget: float = 10.0 total_cost: float = 0 + token_costs: dict[str, dict[str, float]] = TOKEN_COSTS # different model's token cost def update_cost(self, prompt_tokens, completion_tokens, model): """ @@ -46,7 +47,8 @@ class CostManager(BaseModel): return cost = ( - prompt_tokens * TOKEN_COSTS[model]["prompt"] + completion_tokens * TOKEN_COSTS[model]["completion"] + prompt_tokens * self.token_costs[model]["prompt"] + + completion_tokens * self.token_costs[model]["completion"] ) / 1000 self.total_cost += cost logger.info(