import time
from menchuangfallback import menchuangfallback
def extra(
       data, #data
       aiclient,
       qwclient,
       menchuang_collection,
       model):
    if data['bianma'].startswith("0108"):
        sentence=["特征描述：" + data['mc'] + "\n" + data['tz']]
        embeddings = model.encode(sentence)
        result = menchuang_collection.query(query_embeddings=embeddings, n_results=10)
        print(result['documents'][0])
        l = len([x for x in result['distances'][0] if x < 0.5])
        if l < 2:
            l = 2
        completion = aiclient.chat.completions.create(
            model="glm-4.5-flash",
            messages=[
                {"role": "system", "content": "You are a helpful assistant.请将最终答案以JSON格式输出"},
                {"role": "user", "content": "特征描述往往比较具体，工作内容是对特征描述的主要关键的总结提炼。以下是一些特征描述以及对应的提炼的工作内容的例子。" + '\n\n'.join(result['documents'][0][:l]) + "给定一段特征描述，内容为" + data['mc'] +data['tz'] + "。请参照示例，给出提炼的工作内容. 注意，不需要输出特征描述，仅输出工作内容"},
            ],
            extra_body={"thinking": {"type": "disabled"}},
        )
        json_string = completion.choices[0].message.content
        print(json_string)
        answers = json_string.split("\n")
        answers = [x for x in answers if ':' in x ]
        answer2 = answers[0].split(":")[1].replace(" ", "")
        return answer2
    completion = aiclient.chat.completions.create(
        model="glm-4.5-flash",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": " 背景知识：已知预应力高强混凝土管桩（PHC）代号定义为PHC-AAA(BB)CC-DDD-E1,E2,E3,E4,其中AAA代表管桩外径，BB代表管桩壁厚，CC表示型号，DDD表示混凝土强度等级，E1/E2/E3/E4表示分段桩长。例如，PHC-500(125)-AB-C80-9,7 表示外径500mm,壁厚125mm,型号AB，混凝土强度C80, 分段桩长分别为9米和7米，总桩长16米,施工时需要将两个分段接桩"},
        {"role": "user", "content": "问题描述：  给定一段工作内容描述，内容为" + data['mc'] +data['tz'] + "。请判断内容中是否包含桩的代号，如果没有，请输出“无”,如果有，请输出代号"},
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    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 + ",其中给出了一个代号作为答案，请将该最终答案输出"},                       
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    
    answers = json_string.split("\n")
    answers = [x for x in answers if ':' in x ]
    answers = [x for x in answers if not 'true' in x]
    answers = [x for x in answers if not '是' in x]
    print(answers)
    if len(answers) == 0:
        return "无"
    answer2 = answers[0].split(":")[1].replace(" ", "")
    return answer2

def need_extra(
       data, #data
       aiclient,
       qwclient,
       result):
    if data['bianma'].startswith("0108") and len(result) == 0:
        return True
    time.sleep(1)
    completion = qwclient.chat.completions.create(
        model="ZhipuAI/GLM-4.6",
        #model="glm-4.5-flash",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "问题描述：  给定一段工作内容描述，内容为" + data['mc'] +data['tz'] + "。请判断内容是否属于打桩、压桩。请回答是或者否"},
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    if "是" in json_string:
        return True
    else:
        return False