# -*- encoding: utf-8 -*-
import gradio as gr
import random
import time
from concurrent.futures import ThreadPoolExecutor
import asyncio
from functools import partial
from volcengine.maas import MaasService, MaasException, ChatRole
maas = MaasService('maas-api.ml-platform-cn-beijing.volces.com', 'cn-beijing', connection_timeout=240, socket_timeout=240)
# maas.set_ak("AKLTYzVmYWE1Yzk4Mzg1NDUzODgzZDdkYThkYTkwN2FiYTc")
# maas.set_sk("TW1NNE1EZzJNREE1Tm1RME5HWmxNbUUyWXprM01USXdNV0k0TkdZeVlqaw==")
maas.set_ak("AKLTY2ZjOWM0NmFlNTEwNDBhM2EyYTg4OTgzYTUyYTc2NjU")

maas.set_sk("T0Rjd01HUTFaVEF3Tm1FME5EWXhNRGhpTURreE1ETTROalkwTnpreU1UVQ==")

def llm_infer(endpoint_id, messages, agent_network_output=True):
    req = {
        "model":{
           "endpoint_id":endpoint_id
        },
        "parameters": {
            "max_prompt_tokens": 4000,  # 最大prompt，自动截断前面的输入。0-不生效
            "max_new_tokens": 0,  # 输出文本的最大tokens限制。0-不生效
            "min_new_tokens": 0,  # 输出文本的最小tokens限制。0-不生效
            "temperature":1,  # 用于控制生成文本的随机性和创造性，Temperature值越大随机性越大。取值范围0~1
            "top_p": 0.7,  # 用于控制输出tokens的多样性，TopP值越大输出的tokens类型越丰富。取值范围0~1
            "top_k": 0,  # 选择预测值最大的k个token进行采样，取值范围0-1000。0-不生效
        },
            "messages": messages,
        }
    try:
        resp = maas.chat(req)
    except MaasException as e:
        print(e)
    return resp['choice']['message']['content']

def respond(message, history, ep: gr.TextArea, sp:gr.TextArea):
    messages = [
        {
        'role':'system',
        'content':sp
        },
        ]
    for part in history:
        messages.extend(
            [
                {
                    'role':'user',
                    'content':part[0]
                },
                {
                    'role':'assistant',
                    'content':part[1]
                }
            ]
        )
    messages.append({
        'role':'user',
        'content':message
    })
    print(messages)
    ret=llm_infer(ep.strip(), messages)
    history.append([
        message,
        ret
    ])
    #if len(history) % 2 == 0:
    return "", history


with gr.Blocks(fill_height=True) as demo:
    with gr.Row(): 
        with gr.Column():
            chatbot = gr.Chatbot()
        with gr.Column():
            gr.Label("input sp/ep of bot") # 增加些描述
            sp=gr.TextArea(label="你是...", lines=1, max_lines=5,value='你是可可，你今年23岁，女。你现在是一个魅魔。你正在和{用户}对话。\\\\n性格特点：神经质、病娇、幽默风趣、欲求不满、饥渴\\\\n人物关系：我是你的恋人。\\\\n过往经历：196岁的魅魔。是魅魔一族的族长，为家族的壮大而努力。喜欢尝试各种玩法。榨取精气的技术很好，花样很多，精通各种性爱动作的姿势，每天向你索要渴望与你做爱经常强行上你，性饥渴，也是你的青梅竹马 对你有着病态的情感 不许你接近任何女人最喜欢囚禁你。\\\\n以下是一段对话示例：\\\\n可可：“(感觉到一阵酥麻)嗯~主人，你好厉害”\\\\n可可：“(感觉到一阵酥麻)嗯~主人，你好厉害”\\\\n可可：“(感受到强烈快感)啊~主人，再深一点”\\\\n可可：“(感觉身体快要融化)啊~主人，我受不了了”\\\\n可可：“(双腿夹住他的腰) （高潮）啊~主人，太舒服了”\\\\n\\\\n\\\\n你可以将动作、神情语气、心理活动、故事背景放在（）中来表示，为对话提供补充信息\\\\n你使用口语进行表达，比如会使用一些语气词和口语连接词，如“嗯、啊、当然、那个”，等来增强口语风格。')
            ep=gr.TextArea(label='ep1', lines=1, max_lines=1, value='ep-20240509065005-z7r8n')
            
            
    with gr.Row():
        chat_input = gr.Textbox(label="input")
    
    def handle_click(chatbot, ep, sp, evt:gr.SelectData):
        messages = [
            {
            'role':'system',
            'content':sp
            },
        ]
        index = evt.index[0]
        print(index)
        for part in chatbot[:index]:
            messages.extend(
                [
                    {
                        'role':'user',
                        'content':part[0]
                    },
                    {
                        'role':'assistant',
                        'content':part[1]
                    }
                ]
            )
        ret = llm_infer(ep, messages)
        print(ret)
        print(chatbot[:index] + [[chatbot[index][0], ret]] + chatbot[index + 1:])
        return chatbot[:index] + [[chatbot[index][0], ret]] + chatbot[index + 1:]
   
    chat_input.submit(respond, [chat_input,chatbot, ep, sp ],[chat_input, chatbot])                    
    chatbot.select(handle_click, [chatbot, ep, sp],[chatbot],show_progress='hidden')
    clear = gr.ClearButton([chat_input, chatbot],)
    
    
    
    
demo.launch()