import json
with open("cuoshi_corpus.txt", "r") as f:
    content = f.read()
pairs = content.split("\n")
pair_tuple = []
i = 0
while i < len(pairs):
    pair_tuple.append([pairs[i], pairs[i+1]])
    i = i + 2
def aifilter(A, #options
       B, #data
       aiclient):
    options=[]
    letters = "ABCDEFGHIJKLMN"
    for i in range(len(A)):
        options.append("给定选项" + letters[i]+"，内容为"+A[i])

    completion = aiclient.chat.completions.create(
        model="glm-4.5-flash",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": " 背景知识：如果工作内容是土石方工程、土方工程等，那么不要能选用塔式起重机"},
            {"role": "user", "content": "问题描述： 给定一段工作内容: " + B['label'] + " " +  B['mc'] + " " + B['tz'] +   "," + "，".join(options) + "。请做出筛选，并返回结果。例如，如果处理完后剩余A,B,C三个选项，请返回[A,B,C]"},
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    completion = aiclient.chat.completions.create(
        model="glm-4.5-flash",
        messages=[
            {"role": "system", "content": "You are a helpful assistant.请将最终答案以JSON格式输出"},
            {"role": "user", "content": " 给你一段文字如下， " + json_string + ",其中给出了一个类似于[A,B,C]的数组作为答案，请将该最终答案输出"},                       
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    answer=[]
    if 'A' in json_string:
        answer.append(A[0])
    if 'B' in json_string:
        answer.append(A[1])
    if 'C' in json_string:
        answer.append(A[2])
    if 'D' in json_string:
        answer.append(A[3])
    if 'E' in json_string:
        answer.append(A[4])
    if 'F' in json_string:
        answer.append(A[5])
    if 'G' in json_string:
        answer.append(A[6])
    if 'H' in json_string:
        answer.append(A[7])
    return answer
def postprocess0117(selected, data, aiclient):
    correct=[]
    for entry in selected:
        correct.append(entry)
        for item in pair_tuple:
            if entry in item:
                correct = correct + item
    correct = list(set(correct))
    return aifilter(correct, data, aiclient)