app.py

from openai import OpenAI
import gradio as gr
import time
import tiktoken

def calculate_cost_and_tokens(string: str, model_name: str, encoding_name: str = "cl100k_base", models_info=None):
    """
    计算字符串的token数量和消耗费用。
    :参数 string: 需要计算的字符串
    :参数 model_name: 模型名称，用于获取token价格
    :参数 encoding_name: 编码名称，默认为 "cl100k_base"
    :参数 models_info: 包含模型价格和token数量的字典
    :return: 格式化的字符串，包含token数量和消耗费用
    """
    if models_info is None:
        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价格
    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))
    
    # 计算总费用
    total_cost = num_tokens * price_per_token
    
    return f"使用模型:{model_name},消耗tokens数: {num_tokens}, 消耗费用: ${total_cost:.6f}"

   

def update_textbox(endpoint, transmit_api_key, official_api_key ):
    
    if endpoint == "https://lmzh.top/v1":
        if not transmit_api_key:
            raise gr.Error('你选择的是中转接口，请在 key 验证中先验证中转 API Key')
        return transmit_api_key
    elif endpoint == "https://api.openai.com/v1":
        if not official_api_key:
            raise gr.Error('你选择的是官方接口请在 key 验证中先验证官方 API Key')
        return official_api_key   
    else:
        raise gr.Error('未选择端点')

    
def generate_response(question, model_name, endpoint, transmit_api_key, official_api_key):
    
    api_key = update_textbox(endpoint, transmit_api_key, official_api_key )

    # 创建 OpenAI 客户端
    #判断接口是否有值
    if endpoint == '':
        raise gr.Error('请根据key类型选择不同的接口地址')
    client = OpenAI(api_key=api_key, base_url=endpoint)
    messages = [
        {"role": "user", "content": question}
    ]

    try:
        response = client.chat.completions.create(
            model=model_name,
            messages=messages,
            max_tokens=1000,
            temperature=0.1,
            stream=True,
        )
        result=""
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                content = chunk.choices[0].delta.content
                for char in content:
                    result += char
                    yield result
                    if char in "，。！？":
                        time.sleep(0.2)  # 较短的延迟时间
                    else:
                        time.sleep(0.05)  # 更短的延迟时间
    except Exception as e:
        return f"响应生成失败: {str(e)}"

base_url_options=["https://lmzh.top/v1","https://api.openai.com/v1"]
model_options = ["gpt-3.5-turbo", "gpt-4o-mini", "gpt-4o", "gpt-4o-2024-08-06"]

def clear_output():
    return "","",""

def validate_tran_api_key(api_key):
    if api_key == '':
        raise gr.Error('请输入您的中转API Key')
    try:
        client = OpenAI(api_key=api_key, base_url='https://lmzh.top/v1')
        response = client.models.list()
        models_list = [model.id for model in response.data]

        # 进行简单的聊天测试
        try:
            test_response = client.chat.completions.create(
                model="gpt-4o",
                messages=[
                    {"role": "user", "content": "This is a test"}
                ],
                max_tokens=5
            )
            quota_status = "Quota 容量经检测可用."
        except Exception as e:
            quota_status = "Quota 容量已经耗尽或发生未知错误"

        return f"经检查API-key有效.\n{quota_status}\n可访问Models列表:\n{', '.join(models_list)}"
    except Exception as e:
        return f"Error: {str(e)}"

def validate_offi_api_key(api_key):
    if api_key == '':
        raise gr.Error('请输入您的官方API Key')
    try:
        client = OpenAI(api_key=api_key)
        response = client.models.list()
        models_list = [model.id for model in response.data]

        # 进行简单的聊天测试
        try:
            test_response = client.chat.completions.create(
                model="gpt-4o",  # 使用一个合适的模型名称进行测试，如'gpt-4o-mini'等
                messages=[
                    {"role": "user", "content": "This is a test"}
                ],
                max_tokens=5
            )
            quota_status = "Quota 容量经检测可用."
        except Exception as e:
            quota_status = "Quota 容量已经耗尽或发生未知错误."

        return f"经检查官方API-key有效 .\n{quota_status}\n可访问Models列表:\n{', '.join(models_list)}"
    except Exception as e:
        return f"Error: {str(e)}"

# def update_api_key(endpoint, transmit_api_key, official_api_key):
#     if  transmit_api_key == '' or official_api_key =='':
#         raise gr.Error('请在key验证中先验证api key')
#     if "transmit" in endpoint:  # 检查 endpoint 是否包含 "transmit"
#         return transmit_api_key if transmit_api_key else "请填写您的中转API Key"
#     else:  # 否则认为是官方端点
#         return official_api_key if official_api_key else "请填写您的官方API Key"
    
# Gradio 前端设计
app = gr.Blocks(css="style.css",title="sipole_verify_key")

with app:
    gr.HTML('''
        <div class="header">
            <div class="header-title">API-KEY验证</div>
            <div class="header-subtitle">
                key购买地址 <a href="https://buy.sipola.cn" target="_blank">here</a>，超强ai工具请访问 <a href="https://ai.sipola.cn" target="_blank">here</a>
            </div>
        </div>
    ''')
    
    with gr.Row(variant='panel'):
        with gr.Column():
            transmit_api_key = gr.Textbox(type='password', label='中转API-Key', placeholder='请在此填写您的中转API Key', elem_classes="custom-textbox")
            with gr.Row(variant='panel'):
                transmit_btn_text = gr.Button("点击验证中转key", variant="primary", elem_classes="custom-button")
                clear_btn = gr.Button("清除", variant="primary", elem_classes="custom-button")

        with gr.Column():
            transmit_output_textbox = gr.Textbox(label="显示key状态及可访问的模型列表", lines=5,max_lines=5, elem_classes="custom-textbox")
        transmit_btn_text.click(validate_tran_api_key,transmit_api_key,transmit_output_textbox)
        clear_btn.click(clear_output,outputs=[transmit_api_key,transmit_output_textbox],show_progress=False)
    with gr.Row(variant='panel'):
        with gr.Column():
            official_api_key = gr.Textbox(type='password', label='官方API Key', placeholder='请在此输入你的官方API Key', elem_classes="custom-textbox")
            with gr.Row(variant='panel'):
                official_btn_text = gr.Button("点击验证官方key", variant="primary", elem_classes="custom-button")
                clear_btn = gr.Button("清除", variant="primary", elem_classes="custom-button")
        with gr.Column():
            official_output_textbox = gr.Textbox(label="显示key状态及可访问的模型列表", lines=5, max_lines=5, elem_classes="custom-textbox" )
        official_btn_text.click(validate_offi_api_key, official_api_key, official_output_textbox)
        clear_btn.click(clear_output,outputs=[official_api_key,official_output_textbox],show_progress=False)
    with gr.Row(variant='panel'):
        with gr.Column():
            question=gr.Textbox(label="请输入你的问题: ",lines=2,max_lines=3)
            model_name=gr.Dropdown(label="选择模型", choices=model_options, value="gpt-3.5-turbo")
            endpoint=gr.Dropdown(label="选择官方或者接口地址", choices=base_url_options, value='https://lmzh.top/v1')
            # endpoint.change(update_api_key, inputs=[endpoint, transmit_api_key, official_api_key], outputs=api_key)
            with gr.Row(variant='panel'):
                message_btn = gr.Button("发送", variant="primary", elem_classes="custom-button")
                clear_btn = gr.Button("清除", variant="primary", elem_classes="custom-button")
        with gr.Column():
            output_textbox = gr.Textbox(label="显示问题答案", lines=12, max_lines=12, elem_classes="custom-textbox",interactive=True,show_label=True,show_copy_button=True)
            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,output_textbox2],show_progress=False)
        output_textbox.change(fn=calculate_cost_and_tokens,inputs=[output_textbox,model_name],outputs=output_textbox2,show_progress=False)
#添加页面底部
    gr.HTML('''
        <div class="footer">
         <center><p>Power by sipola </p></center>
        </div>
    ''')
app.queue()
app.launch(show_error=True)