Update app.py

app.py

@@ -3,13 +3,55 @@ import gradio as gr
 import time
 import tiktoken
 
-def num_tokens_from_message(string: str, encoding_name: str = "cl100k_base") -> int:
-    """返回文本字符串中的Token数量"""
+def calculate_cost_and_tokens(string: str, model_name: str, encoding_name: str = "cl100k_base", models_info=None):
+    """
+    计算字符串的token数量和消耗费用。
+    :param string: 需要计算的字符串
+    :param model_name: 模型名称，用于获取token价格
+    :param encoding_name: 编码名称，默认为 "cl100k_base"
+    :param models_info: 包含模型价格和token数量的字典
+    :return: 格式化的字符串，包含token数量和消耗费用
+    """
+    if models_info is None:
+        models_info = {
+            'Model A': (15.000, 1000000),  # 价格和token数量
+            'Model B': (30.000, 2000000),
+            'Model C': (45.000, 3000000)
+        }
+    
+    # 获取token价格
+    price_per_token = models_info[model_name][0] / models_info[model_name][1]
+    
+    # 获取token数量
     encoding = tiktoken.get_encoding(encoding_name)
     num_tokens = len(encoding.encode(string))
-    return num_tokens
+    
+    # 计算总费用
+    total_cost = num_tokens * price_per_token
+    
+    return f"消耗tokens数: {num_tokens}, 消耗费用: ${total_cost:.6f}"
+
+# 示例用法
+result = calculate_cost_and_tokens("Example text for processing.", "Model A")
+print(result)
 
 
+# 示例：不同模型的价格和token数量
+models_info = {
+    'gpt-3.5-turbo': (6.000, 1000000),  # 价格和token数量
+    'gpt-4o-mini': (0.6, 1000000),
+    'gpt-4o': (15.000, 1000000),
+    'gpt-4o-2024-08-06': (10.000, 1000000),
+}
+
+# 计算每个模型的token成本
+token_costs = calculate_token_cost(models_info)
+
+# 打印结果
+for model, cost in token_costs.items():
+    print(f"{model} token cost: ${cost:.6f}")
+    
+
 def update_textbox(endpoint, transmit_api_key, official_api_key ):
     
     if endpoint == "https://lmzh.top/v1":
@@ -171,7 +213,7 @@ with app:
             output_textbox2 = gr.Textbox(label="显示tokens数", lines=1, max_lines=1, elem_classes="custom-textbox",interactive=False,show_label=True)
         message_btn.click(generate_response,inputs=[question,model_name,endpoint,transmit_api_key, official_api_key],outputs=output_textbox,show_progress=False) 
         clear_btn.click(clear_output,outputs=[question,output_textbox],show_progress=False)
-        output_textbox.change(fn=num_tokens_from_message,inputs=output_textbox,outputs=output_textbox2)
+        output_textbox.change(fn=num_tokens_from_message,inputs=[output_textbox,model_name],outputs=output_textbox2)
 #添加页面底部
     gr.HTML('''
         <div class="footer">