Merge pull request #1010 from better629/feat_yi

Feat add offcial yi-34b

metagpt/configs/llm_config.py

@@ -29,6 +29,7 @@ class LLMType(Enum):
     DASHSCOPE = "dashscope"  # Aliyun LingJi DashScope
     MOONSHOT = "moonshot"
     MISTRAL = "mistral"
+    YI = "yi"  # lingyiwanwu
 
     def __missing__(self, key):
         return self.OPENAI

metagpt/provider/base_llm.py

@@ -10,10 +10,9 @@ from __future__ import annotations
 
 import json
 from abc import ABC, abstractmethod
-from typing import Dict, Optional, Union
+from typing import Optional, Union
 
 from openai import AsyncOpenAI
-from openai.types import CompletionUsage
 from pydantic import BaseModel
 from tenacity import (
     after_log,
@@ -28,7 +27,6 @@ from metagpt.logs import logger
 from metagpt.schema import Message
 from metagpt.utils.common import log_and_reraise
 from metagpt.utils.cost_manager import CostManager, Costs
-from metagpt.utils.exceptions import handle_exception
 
 
 class BaseLLM(ABC):
@@ -88,6 +86,7 @@ class BaseLLM(ABC):
             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 or self.pricing_plan
         model = model or self.model
         usage = usage.model_dump() if isinstance(usage, BaseModel) else usage
         if calc_usage and self.cost_manager:
@@ -225,20 +224,6 @@ class BaseLLM(ABC):
         """
         return json.loads(self.get_choice_function(rsp)["arguments"], strict=False)
 
-    @handle_exception
-    def _update_costs(self, usage: CompletionUsage | Dict):
-        """
-        Updates the costs based on the provided usage information.
-        """
-        if self.config.calc_usage and usage and self.cost_manager:
-            if isinstance(usage, Dict):
-                prompt_tokens = int(usage.get("prompt_tokens", 0))
-                completion_tokens = int(usage.get("completion_tokens", 0))
-            else:
-                prompt_tokens = usage.prompt_tokens
-                completion_tokens = usage.completion_tokens
-            self.cost_manager.update_cost(prompt_tokens, completion_tokens, self.pricing_plan)
-
     def messages_to_prompt(self, messages: list[dict]):
         """[{"role": "user", "content": msg}] to user: <msg> etc."""
         return "\n".join([f"{i['role']}: {i['content']}" for i in messages])

metagpt/provider/openai_api.py

@@ -44,7 +44,7 @@ from metagpt.utils.token_counter import (
 )
 
 
-@register_provider([LLMType.OPENAI, LLMType.FIREWORKS, LLMType.OPEN_LLM, LLMType.MOONSHOT, LLMType.MISTRAL])
+@register_provider([LLMType.OPENAI, LLMType.FIREWORKS, LLMType.OPEN_LLM, LLMType.MOONSHOT, LLMType.MISTRAL, LLMType.YI])
 class OpenAILLM(BaseLLM):
     """Check https://platform.openai.com/examples for examples"""
 

metagpt/utils/token_counter.py

@@ -48,6 +48,8 @@ TOKEN_COSTS = {
     "claude-2.1": {"prompt": 0.008, "completion": 0.024},
     "claude-3-sonnet-20240229": {"prompt": 0.003, "completion": 0.015},
     "claude-3-opus-20240229": {"prompt": 0.015, "completion": 0.075},
+    "yi-34b-chat-0205": {"prompt": 0.0003, "completion": 0.0003},
+    "yi-34b-chat-200k": {"prompt": 0.0017, "completion": 0.0017},
 }
 
 
@@ -176,6 +178,8 @@ TOKEN_MAX = {
     "claude-2.1": 200000,
     "claude-3-sonnet-20240229": 200000,
     "claude-3-opus-20240229": 200000,
+    "yi-34b-chat-0205": 4000,
+    "yi-34b-chat-200k": 200000,
 }