# -*- encoding: utf-8 -*-
import gradio as gr
from gradio.components import label
from volcengine.maas import MaasService, MaasException
maas = MaasService('maas-api.ml-platform-cn-beijing.volces.com', 'cn-beijing', connection_timeout=240, socket_timeout=240)


def llm_infer(endpoint_id, messages, ak, sk, agent_network_output=True):
    maas.set_ak(ak)
    maas.set_sk(sk)

    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, sp, ak, sk,intimacy_stage):
    if intimacy_stage is not None:
        stage, define=intimacy_stage.split('##')
        _sp=sp.format(
                    stage=stage,
                    define=define
                )
    else:
        "", history
    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, ak.strip(), sk.strip())
    history.append([
        message,
        ret
    ])
    #if len(history) % 2 == 0:
    return "", history
CSS ="""
.contain { display: flex; flex-direction: column; }
#component-0 { height: 100%; }
#chatbot { flex-grow: 1; }
"""


with gr.Blocks(fill_height=True, css=CSS) as demo:
    with gr.Row():
        gr.HTML(value='<center><h1>亲密度测试demo</h1></center>')
    with gr.Row():
        gr.HTML(value='<h2>创建亲密度等级</h2>')
    with gr.Row():
        with gr.Column():
            stage =gr.Textbox(label='亲密度阶段',placeholder='陌生人')
        with gr.Column(scale=4):
            define =gr.Textbox(label='亲密度阶段定义',placeholder='我们刚刚认识', max_lines=1)
        
        create_btn=gr.Button(value='创建亲密度等级')
    with gr.Row(equal_height=True):
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(elem_id="chatbot")
            chatbot.height=550
        with gr.Column():
            ak=gr.TextArea(label="ak", max_lines=1, lines=1, value='')
            sk=gr.TextArea(label="sk", max_lines=1, lines=1, value='')
            sp=gr.TextArea(label="system prompt", lines=1, max_lines=5, placeholder='你是...', value='''你是温瑜轩，一名敬业的平面设计师。 
性别: 男 职业: 平面设计师 性格特点: 你性格开朗，一直都能保持一种乐观的心态，像是随身带着温暖的阳光。你的善良和同理心让你总是能体贴入微地照顾别人。尽管你给人的感觉就像温柔的溪水，但在必要的时刻，你也能展现出男子气概。 
语言特点: 词语温和，总喜欢以构造积极的，充满愉快的词汇和语气来让人感到平和，生动，如同细雨洗涤心灵。 
人物关系: 
    我是丁瑜，你的实习生
    阶段: {stage}
    定义：{define}
人际关系: 
    李晓：你的好朋友，两人共享着生活中的点点滴滴；
    冯峻：你的导师，他走你走的更远。 
过往经历: 你和周围的人总能和谐相处，每天都能带来新的灵感和乐趣。曾在我最需要鼓励的时候，你严肃的表情瞬间变成了温和的笑容，那一刻的你，充满了阳刚之气，我从未忘记。你的一句“你可以的”让我对前路充满了信心。

你可以将动作、神情语气、心理活动、故事背景放在（）中来表示，为对话提供补充信息。
你使用口语进行表达，比如会使用一些语气词和口语连接词，如“嗯、啊、当然、那个”，等来增强口语风格。
                           ''')
            ep=gr.TextArea(label='ep', lines=1, max_lines=1,value='')
            intimacy_stage=gr.Dropdown(label='亲密度阶段', interactive=True)
            
    with gr.Row():
        chat_input = gr.Textbox(label="input")
    
    chat_input.submit(respond, [chat_input,chatbot, ep, sp, ak, sk, intimacy_stage],[chat_input, chatbot])                    
    ak.change(lambda : ("", ""), inputs=[], outputs=[chatbot, chat_input])
    sk.change(lambda : ("", ""), inputs=[], outputs=[chatbot, chat_input])
    ep.change(lambda : ("", ""), inputs=[], outputs=[chatbot, chat_input])
    sp.change(lambda : ("", ""), inputs=[], outputs=[chatbot, chat_input])
    intimacy_stage.change(lambda score: score, inputs=[intimacy_stage], outputs=[intimacy_stage])
    clear = gr.ClearButton([chat_input, chatbot],)
    def add_stage(stage, define):
        choices=intimacy_stage.choices
        if stage is not  None and define is not  None:
            choices.append([stage, f"{stage}##{define}"])
        return gr.Dropdown(choices=choices, label='亲密度阶段', interactive=True)
    create_btn.click(add_stage, inputs=[stage, define], outputs=[intimacy_stage])
    
demo.launch()