Merge pull request #942 from better629/feat_new_qianfan

Feat simplify provider ut code

metagpt/provider/fireworks_api.py

@@ -19,7 +19,7 @@ from metagpt.configs.llm_config import LLMConfig, LLMType
 from metagpt.logs import log_llm_stream, logger
 from metagpt.provider.llm_provider_registry import register_provider
 from metagpt.provider.openai_api import OpenAILLM, log_and_reraise
-from metagpt.utils.cost_manager import CostManager, Costs
+from metagpt.utils.cost_manager import CostManager
 
 MODEL_GRADE_TOKEN_COSTS = {
     "-1": {"prompt": 0.0, "completion": 0.0},  # abnormal condition
@@ -81,17 +81,6 @@ class FireworksLLM(OpenAILLM):
         kwargs = dict(api_key=self.config.api_key, base_url=self.config.base_url)
         return kwargs
 
-    def _update_costs(self, usage: CompletionUsage):
-        if self.config.calc_usage and usage:
-            try:
-                # use FireworksCostManager not context.cost_manager
-                self.cost_manager.update_cost(usage.prompt_tokens, usage.completion_tokens, self.model)
-            except Exception as e:
-                logger.error(f"updating costs failed!, exp: {e}")
-
-    def get_costs(self) -> Costs:
-        return self.cost_manager.get_costs()
-
     async def _achat_completion_stream(self, messages: list[dict], timeout=3) -> str:
         response: AsyncStream[ChatCompletionChunk] = await self.aclient.chat.completions.create(
             **self._cons_kwargs(messages), stream=True

metagpt/provider/google_gemini_api.py

@@ -72,16 +72,6 @@ class GeminiLLM(BaseLLM):
         kwargs = {"contents": messages, "generation_config": GenerationConfig(temperature=0.3), "stream": stream}
         return kwargs
 
-    def _update_costs(self, usage: dict):
-        """update each request's token cost"""
-        if self.config.calc_usage:
-            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, self.model)
-            except Exception as e:
-                logger.error(f"google gemini updats costs failed! exp: {e}")
-
     def get_choice_text(self, resp: GenerateContentResponse) -> str:
         return resp.text
 

metagpt/provider/ollama_api.py

@@ -46,16 +46,6 @@ class OllamaLLM(BaseLLM):
         kwargs = {"model": self.model, "messages": messages, "options": {"temperature": 0.3}, "stream": stream}
         return kwargs
 
-    def _update_costs(self, usage: dict):
-        """update each request's token cost"""
-        if self.config.calc_usage:
-            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, self.model)
-            except Exception as e:
-                logger.error(f"ollama updats costs failed! exp: {e}")
-
     def get_choice_text(self, resp: dict) -> str:
         """get the resp content from llm response"""
         assist_msg = resp.get("message", {})

metagpt/provider/open_llm_api.py

@@ -8,7 +8,7 @@ from metagpt.configs.llm_config import LLMConfig, 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.cost_manager import Costs, TokenCostManager
+from metagpt.utils.cost_manager import TokenCostManager
 from metagpt.utils.token_counter import count_message_tokens, count_string_tokens
 
 
@@ -34,14 +34,3 @@ class OpenLLM(OpenAILLM):
             logger.error(f"usage calculation failed!: {e}")
 
         return usage
-
-    def _update_costs(self, usage: CompletionUsage):
-        if self.config.calc_usage and usage:
-            try:
-                # use OpenLLMCostManager not CONFIG.cost_manager
-                self._cost_manager.update_cost(usage.prompt_tokens, usage.completion_tokens, self.model)
-            except Exception as e:
-                logger.error(f"updating costs failed!, exp: {e}")
-
-    def get_costs(self) -> Costs:
-        return self._cost_manager.get_costs()

metagpt/provider/openai_api.py

@@ -30,7 +30,7 @@ from metagpt.provider.constant import GENERAL_FUNCTION_SCHEMA
 from metagpt.provider.llm_provider_registry import register_provider
 from metagpt.schema import Message
 from metagpt.utils.common import CodeParser, decode_image
-from metagpt.utils.cost_manager import CostManager, Costs
+from metagpt.utils.cost_manager import CostManager
 from metagpt.utils.exceptions import handle_exception
 from metagpt.utils.token_counter import (
     count_message_tokens,
@@ -56,16 +56,13 @@ class OpenAILLM(BaseLLM):
 
     def __init__(self, config: LLMConfig):
         self.config = config
-        self._init_model()
         self._init_client()
         self.auto_max_tokens = False
         self.cost_manager: Optional[CostManager] = None
 
-    def _init_model(self):
-        self.model = self.config.model  # Used in _calc_usage & _cons_kwargs
-
     def _init_client(self):
         """https://github.com/openai/openai-python#async-usage"""
+        self.model = self.config.model  # Used in _calc_usage & _cons_kwargs
         kwargs = self._make_client_kwargs()
         self.aclient = AsyncOpenAI(**kwargs)
 
@@ -272,16 +269,6 @@ class OpenAILLM(BaseLLM):
 
         return usage
 
-    @handle_exception
-    def _update_costs(self, usage: CompletionUsage):
-        if self.config.calc_usage and usage and self.cost_manager:
-            self.cost_manager.update_cost(usage.prompt_tokens, usage.completion_tokens, self.model)
-
-    def get_costs(self) -> Costs:
-        if not self.cost_manager:
-            return Costs(0, 0, 0, 0)
-        return self.cost_manager.get_costs()
-
     def _get_max_tokens(self, messages: list[dict]):
         if not self.auto_max_tokens:
             return self.config.max_token

metagpt/provider/zhipuai_api.py

@@ -53,16 +53,6 @@ class ZhiPuAILLM(BaseLLM):
         kwargs = {"model": self.model, "messages": messages, "stream": stream, "temperature": 0.3}
         return kwargs
 
-    def _update_costs(self, usage: dict):
-        """update each request's token cost"""
-        if self.config.calc_usage:
-            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, self.model)
-            except Exception as e:
-                logger.error(f"zhipuai updats costs failed! exp: {e}")
-
     def completion(self, messages: list[dict], timeout=3) -> dict:
         resp: Completion = self.llm.chat.completions.create(**self._const_kwargs(messages))
         usage = resp.usage.model_dump()