import gradio as gr
from volcengine.maas.v2 import MaasService
from volcengine.maas import MaasException, ChatRole
import json
import os

def maas_chat(input, novel_type):  # 添加小说类型参数
    query = input.split('\n')
    query = [item for item in query if item]
    query = json.dumps(query, ensure_ascii=False)
    maas = MaasService('maas-api.ml-platform-cn-beijing.volces.com', 'cn-beijing',connection_timeout=15)


    # 使用 os 模块来访问变量和秘密值
    ak = os.getenv('AK')
    sk = os.getenv('SK')
    maas.set_ak(ak)
    maas.set_sk(sk)
    prompt = "我将提供一份小说片段并且是使用 JSON 数组进行格式化，JSON 数组内的每个对象的内容就是一个故事的场景，你根据我提供的场景原文生成用来 AI 绘画的提示词，注意提示词应该和小说片段具有强相关性，当小说内容没有场景时，结合 JSON 中其他内容输出提示词，并同样使用 JSON 输出并且数组大小和输入的数组大小一样，下面给出 JSON 格式化后的文字："
    if novel_type == "古文":
        prompt = '我将提供一份古风小说片段并且是使用JSON数组进行格式化，JSON数组内的每个对象的内容就是一个小说内容的一个片段，你根据我提供的小说片段生成中国古代风格的AI绘画的提示词，注意提示词应该和小说片段具有强相关性，当小说内容没有场景时，结合JSON中其他内容输出提示词，并同样使用JSON输出并且数组大小和输入的数组大小一样，下面给出JSON格式化后的文字：';
    # chat
    req = {
        "parameters": {
            "max_new_tokens": 1000,  # 输出文本的最大 tokens 限制
            "min_new_tokens": 1,  # 输出文本的最小 tokens 限制
            "temperature": 0,  # 用于控制生成文本的随机性和创造性，Temperature 值越大随机性越大，取值范围 0~1
            "top_p": 0,  # 用于控制输出 tokens 的多样性，TopP 值越大输出的 tokens 类型越丰富，取值范围 0~1
            "top_k": 0,  # 选择预测值最大的 k 个 token 进行采样，取值范围 0-1000，0 表示不生效
            "max_prompt_tokens": 4096,  # 最大输入 token 数，如果给出的 prompt 的 token 长度超过此限制，取最后 max_prompt_tokens 个 token 输入模型。
        },
        "messages": [
            {
                "role": ChatRole.SYSTEM,
                "content": prompt
            }, {
                "role": ChatRole.USER,
                "content": query
            }
        ]
    }

    ID = os.getenv('ID')
    GUID = os.getenv('GUID')
    # endpoint_id = novel_types =='现代'?ID: GUID
    endpoint_id = ID if novel_type == "现代" else GUID
    resp = maas.chat(endpoint_id, req)
    output = resp.choices[0].message['content']
    #去除 output'符号

    list = json.loads(output)
    return '\n'.join(list)


novel_types = ["现代", "古文"]  # 定义小说类型列表

demo = gr.Interface(
    fn=maas_chat,
    inputs=[gr.Textbox(lines=6, placeholder="输入小说片段（JSON 格式）"), gr.Dropdown(novel_types,label="文章类型")],  # 添加下拉输入框默认现代

    outputs=gr.Textbox(lines=6, placeholder="输出提示词（JSON 格式）"),
    title="小说片段转 AI 绘画提示词",
    description="输入小说片段（JSON 格式），输出用来 AI 绘画的提示词（JSON 格式）。",
)

demo.launch(share=True)