from abc import ABC, abstractmethodfrom dataclasses import dataclass
@dataclassclass BaseChatbot(ABC):
mode: str = "API"
@abstractmethoddef ask(self, msg: str) -> str:"""Ask LLM a question and get an answer"""
@abstractmethoddef ask_batch(self,msgs:list) -> str:"""ask LLM a series of questions and get a series of answers"""
@abstractmethoddef ask_code(self, msgs: list) -> str:"""Ask GPT multiple questions and get a piece of code"""

def ask_batch(self, msgs: list) -> str:context = []for msg in msgs:umsg = self._user_msg(msg)context.append(umsg)rsp = self.completion(context)rsp_text = self.get_choice_text(rsp)context.append(self._assistant_msg(rsp_text))return self._extract_assistant_rsp(context)

def _user_msg(self, msg: str) -> dict[str, str]:return {"role": "user", "content": msg}
def _assistant_msg(self, msg: str) -> dict[str, str]:return {"role": "assistant", "content": msg}
def _system_msg(self, msg: str) -> dict[str, str]:return {"role": "system", "content": msg}
def _system_msgs(self, msgs: list[str]) -> list[dict[str, str]]:return [self._system_msg(msg) for msg in msgs]
def _default_system_msg(self):return self._system_msg(self.system_prompt)

import tiktoken
def count_string_tokens(string: str, model_name: str) -> int:encoding = tiktoken.encoding_for_model(model_name)return len(encoding.encode(string))


if __name__ == "__main__":test_sentence = "hello world azent is coming soon" # 7 tokentokens_len = count_string_tokens(test_sentence,"gpt-3.5-turbo")print(tokens_len)

if __name__ == "__main__":test_sentence = "您好, azent 框架是为了中小企业而生，是 LLM 走进您的企业和公司正确方式，感觉您的支持，虽然是一杯咖啡，却盛满对开发者支持!"tokens_len = count_string_tokens(test_sentence,"gpt-3.5-turbo")print(tokens_len)

TOKEN_COSTS = {"gpt-3.5-turbo": {"prompt": 0.002, "completion": 0.002},"gpt-3.5-turbo-0301": {"prompt": 0.002, "completion": 0.002},"gpt-4-0314": {"prompt": 0.03, "completion": 0.06},"gpt-4": {"prompt": 0.03, "completion": 0.06},"gpt-4-32k": {"prompt": 0.06, "completion": 0.12},"gpt-4-32k-0314": {"prompt": 0.06, "completion": 0.12},"text-embedding-ada-002": {"prompt": 0.0004, "completion": 0.0},

class Costs(NamedTuple):total_prompt_tokens: inttotal_completion_tokens: inttotal_cost: floattotal_budget: float

class CostManager(metaclass=Singleton):def __init__(self):self.total_prompt_tokens = 0self.total_completion_tokens = 0self.total_cost = 0self.total_budget = 0
self.config = Config()
def update_cost(self, prompt_tokens, completion_tokens, model): self.total_prompt_tokens += prompt_tokensself.total_completion_tokens += completion_tokens
cost = (prompt_tokens * TOKEN_COSTS[model]["prompt"]+ completion_tokens * TOKEN_COSTS[model]["completion"]) / 1000

self.total_cost += cost
logger.info(f"Total running cost: ${self.total_cost:.3f} | Max budget: ${self.config.max_budget:.3f} | "f"Current cost: ${cost:.3f}, {prompt_tokens=}, {completion_tokens=}")self.config.total_cost = self.total_cost
def get_total_prompt_tokens(self):return self.total_prompt_tokensdef get_total_completion_tokens(self):return self.total_completion_tokensdef get_total_cost(self):return self.total_costdef get_costs(self) -> Costs:return Costs(self.total_prompt_tokens, self.total_completion_tokens, self.total_cost, self.total_budget)