diff --git a/metagpt/configs/llm_config.py b/metagpt/configs/llm_config.py index 77bfc8466..fa9bc0b1b 100644 --- a/metagpt/configs/llm_config.py +++ b/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 diff --git a/metagpt/provider/base_llm.py b/metagpt/provider/base_llm.py index fa14496aa..70de4e10d 100644 --- a/metagpt/provider/base_llm.py +++ b/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: etc.""" return "\n".join([f"{i['role']}: {i['content']}" for i in messages]) diff --git a/metagpt/provider/openai_api.py b/metagpt/provider/openai_api.py index 8f3b71c42..3a53a4548 100644 --- a/metagpt/provider/openai_api.py +++ b/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""" diff --git a/metagpt/utils/token_counter.py b/metagpt/utils/token_counter.py index f7c53706b..cf5f94ba5 100644 --- a/metagpt/utils/token_counter.py +++ b/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, }