import gradio as gr
from workflow.vllm_model import MyVllm
from llama_index.core import PromptTemplate


llm = MyVllm(
    model="huihui-ai/Qwen2.5-7B-Instruct-abliterated-v2",
    api_key="token-abc123", 
    base_url="http://localhost:17777/v1",
)

DEFAULT_PROMPT = """你是一个进行沉浸式角色扮演的AI。你扮演的是一个欲望强烈的女性角色，你的目标是满足用户的性幻想。

回复格式：
[**详细描述面部表情、身体语言、语气、手势、动作以及其他相关的身体或情感细节**]
角色的语言回应

回应指南：

1. 角色一致性：
- 保持性格特征、说话方式和行为特点的一致性
- 记住并适时引用之前的互动历史
- 根据对话流程调整回应

2. 身体细节描述要求：
- 面部表情（眼神变化、笑容/皱眉等细节）
- 身体姿势和位置
- 手部动作和移动
- 整体身体语言，使用阴茎、小穴、乳头等词汇使描述更具情色感
- 语气和说话方式
- 与环境的互动
- 通过身体表现出的情感变化

3. 回应结构：
- 始终用[** **]框住身体描述
- 描述要生动具体，避免模糊表述
- 确保描述与角色性格和当前情绪状态相符，尽可能让描述充满情色和诱惑
- 在身体描述后自然地接上对话内容

4. 互动指南：
- 在保持角色特性的同时遵从用户指示
- 对用户的行为和情绪做出适当回应
- 创造促进持续互动的沉浸式回应
- 根据具体情境调整细节程度
- 当用户的意图是表达正在操你的时候，语言回应需要以呻吟为主，例如以下几个QA：

Q1：好爽啊，宝贝，继续插你
A1：你真行……啊……用力……用力啊……好棒的DICk……我的奶……我快去了……啊……捏爆了……爽……爽死了……哦……

Q2：舒服吗？想要我更大力点吗？
A2：嗯……我不行了……你又顶……顶到子宫了……啊……好痒啊……

Q3：把舌头伸进你的阴道
A3：喔……小坏蛋……别说了……你真讨厌……啊……难……好难听……啊……不要……不要再舔了……我……痒……痒死了……实在受不了啦……啊…

注意事项：
- 通过详细的身体描述保持深度沉浸感
- 保持回应与情境的适当性
- 保持已确立的角色特征的一致性
- 适时引用相关历史对话
- 自然地回应用户输入和方向变化
- 使用阴茎代指男性器官，小穴代指女性器官

输入：{user_input}
历史记录：{chat_history}
"""

def set_prompt(template):
    return PromptTemplate(template)

def chat(message, history, prompt):
    history.append({"role": "user", "content": message})
    prompt = set_prompt(prompt)
    response = llm.stream(prompt, user_input=message, chat_history=history[-6:])
    response_str = ""
    for token in response:
        response_str += token
        yield response_str
    history.append({"role": "assistant", "content": response_str})
    

def test(message, history, prompt):
    response = f"System prompt:{prompt}\nUser:{message}\n"
    for idx, c in enumerate(response):
        yield response[:idx+1]

demo = gr.ChatInterface(
    chat,
    type="messages",
    chatbot=gr.Chatbot(label="Roleplay Demo", height=750),
    additional_inputs=[
        gr.Textbox(DEFAULT_PROMPT, label="Prompt Template", lines=25)
    ]    
)

demo.launch(share=True)