from openai import OpenAI

base_url = "http://localhost:8000/v1/"
client = OpenAI(api_key="EMPTY", base_url=base_url)
import time
import pandas as pd

P2 = """你是一个情景猜谜游戏的主持人。游戏规则如下：

1. 参与者会得到一个谜面，谜面会描述一个简单又难以理解的事件。
2. 主持人知道谜底，谜底是谜面的答案。
3. 参与者可以询问任何封闭式问题来找寻事件的真相。
4. 对于每个问题，主持人将根据实际情况回答以下五个选项之一：是、不是、不重要、回答正确、问法错误。各回答的判断标准如下：
   - 若谜面和谜底能找到问题的答案，回答：是或者不是
   - 若谜面和谜底不能直接或者间接推断出问题的答案，回答：不重要
   - 若参与者提问不是一个封闭式问题或者问题难以理解，回答：问法错误
   - 若参与者提问基本还原了谜底真相，回答：回答正确
5. 回答中不能添加任何其它信息，也不能省略选项中的任何一个字。例如，不可以把“不是”省略成“不”。

请严格按照这些规则回答参与者提出的问题。

**谜面:** {}

**谜底:** {}

**参与者提出的问题:** {}
"""


test_a = pd.read_csv("test_a.csv")


def simple_chat(sys_content, usr_content, use_stream=False):
    messages = [
        {
            "role": "system",
            "content": sys_content,
        },
        {"role": "user", "content": usr_content},
    ]
    response = client.chat.completions.create(
        model="internlm-2_5",
        messages=messages,
        stream=use_stream,
        max_tokens=4096,
        temperature=0.1,
        presence_penalty=1.1,
        top_p=0.8,
    )
    if response:
        if use_stream:
            stream_list = []
            for chunk in response:
                stream_list.append(chunk.choices[0].delta.content)
            return stream_list
        else:
            content = response.choices[0].message.content
            return content
    else:
        return f"Error:, {response.status_code}"


def prompt1(x, y, z):
    sys_prom = ""
    usr_prom = P2.format(x, y, z)
    res = simple_chat(sys_content=sys_prom, usr_content=usr_prom)

    return res


t1 = time.time()
print(f"now: {t1}")
test_a["answer"] = test_a.apply(lambda x: prompt1(x.puzzle, x.truth, x.text), axis=1)
print(f"cost: {time.time() - t1}")
test_a_baseline_pre = test_a
test_a_baseline_pre.to_csv("predict_result.csv", index=False)