celery/extra.py

import time
from config import simplemodel
from menchuangfallback import menchuangfallback


example1='''
1、包水管 做法详见图纸设计及相关图集规范
2、5#镀锌角钢基架
3、水管外包吸音材料
4、水泥纤维板
'''
answer_1='''
{'面层': '水泥纤维板'}
'''


example2='''
1、一层大厅消火栓处，做法详见图纸设计及相关图集规范
2、石膏板、阻燃基层板、40*20*3镀锌方通龙骨、金属转轴等
'''

answer_2='''
{'面层': '石膏板'}
'''

example3='''
1、不锈钢饰面 做法详见图纸设计及相关图集规范
2、1.5mm厚不锈钢面层
3、阻燃板基层，木龙骨找平阻燃处理
'''

answer_3='''
{'面层': '1.5mm厚不锈钢面层'}
'''

def extra(
       data, #data
       aiclient,
       qwclient,
       sfclient,
       menchuang_collection,
       model,
       qita_collection,
       ):
    if '高强螺栓' in data['mc']:
        return '高强螺栓'
    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
    if data['bianma'].startswith("0115"):
        sentence=["特征描述：" + data['mc'] + "\n" + data['tz']]
        embeddings = model.encode(sentence)
        result = qita_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
    if data['bianma'].startswith("011207") or data['bianma'].startswith('011208'):
        completion = aiclient.chat.completions.create(
            model="glm-4.5-air",
            messages=[
                {"role": "system", "content": "You are a helpful assistant.请将最终答案以JSON格式输出"},
                {"role": "user", "content": "墙柱面装饰板工程往往由多道工序组成，包括底层龙骨，附着于龙骨上的基层板、吸音棉、阻燃板等，以及作为饰面的面层板。现在要求你从给定的工作内容描述中抽取出面层的描述。举个例子，给定工作内容："+example1+"\n你应该返回："+answer_1+"\n再举个例子，给定工作内容："+example2+"\n你应该返回:"+answer_2+"\n再举个例子，给定工作内容："+example3+"\n你应该返回："+answer_3+"\n现在给定工作内容："+data['mc']+" "+data['tz']+"\n请给出你的答案"},
            ],
            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-air",
        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,
       sfclient,
       result, name_label):
    if data['bianma'].startswith('011207') or data['bianma'].startswith('011208'):
        result = [x[0].replace('换', '') for x in result]
        result = [name_label[x] for x in result]
        if len([x for x in result if '铝板' in x or '铝单板' in x]) > 0:
            return False
        left = [x.replace('在木龙骨', '')  for x in result ]
        left = [x.replace('在龙骨', '')  for x in left ]
        left = [x for x in left if '龙骨' not in x]
        left = [x for x in left if '细木工板基层' not in x]
        left = [x for x in left if '吸音' not in x]
        if len(left) == 0:
            return True
        return False

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