graphQa/actions/api/answer_search.py

from py2neo import Graph


class AnswerSearcher:
    def __init__(self):
        self.g = Graph("bolt://localhost:7687", auth=("neo4j", "fdx3081475970"))
        self.num_limit = 20

    '''执行cypher查询，并返回相应结果'''
    def search_main(self, sqls):
        final_answers = []
        for sql_ in sqls:
            question_type = sql_['question_type']
            queries = sql_['sql']
            answers = []
            for query in queries:
                ress = self.g.run(query).data()
                answers += ress
            final_answer = self.answer_prettify(question_type, answers)
            if final_answer:
                final_answers.append(final_answer)
        return final_answers

    def get_entity(self, sqls):
        entity_lists = []
        for sql_ in sqls:
            question_type = sql_['question_type']
            queries = sql_['sql']
            answers = []
            for query in queries:
                ress = self.g.run(query).data()
                answers += ress
            entity_list = self.get_entity_list(question_type, answers)
            print("entity_list", entity_list)
            if entity_list:
                entity_lists.append(entity_list)
        return entity_lists

    def get_entity_list(self, question_type, answers):
        entity_list = []
        if not answers:
            return ''

        if question_type == 'query_des':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = str(i['m.name'])
                obj = str(i['n.name'])
                rel = str(i['r.name'])
                ossid = i['n.doc_id']
                file_name = i['n.doc_name']
                filePage = i['n.filePage']
                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "故障描述"})
                    seen_subjects.add(sub)
                entity_list["data"].append({"name": obj, "category": rel})
                entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'query_hmc':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = str(i['m.name'])
                obj = str(i['n.name'])
                rel = str(i['r.name'])
                ossid = i['n.doc_id']
                file_name = i['n.doc_name']
                filePage = i['n.filePage']
                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "HMC"})
                    seen_subjects.add(sub)
                entity_list["data"].append({"name": obj, "category": rel})
                entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'query_obj':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = str(i['m.name'])
                obj = str(i['n.name'])
                rel = str(i['r.name'])
                ossid = i['n.doc_id']
                file_name = i['n.doc_name']
                filePage = i['n.filePage']
                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "成品"})
                    seen_subjects.add(sub)
                entity_list["data"].append({"name": obj, "category": rel})
                entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'query_x_obj':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = str(i['m.name'])
                obj = str(i['n.name'])
                rel = str(i['r.name'])
                ossid = i['n.doc_id']
                file_name = i['n.doc_name']
                filePage = i['n.filePage']
                if sub not in seen_subjects:
                    entity_list["data"].append({"name": obj, "category": "型号"})
                    seen_subjects.add(obj)
                entity_list["data"].append({"name": sub, "category": rel})
                entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'query_s_sys':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = str(i['m.name'])
                obj = str(i['n.name'])
                rel = str(i['r.name'])
                ossid = i['n.doc_id']
                file_name = i['n.doc_name']
                filePage = i['n.filePage']
                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "子系统"})
                    seen_subjects.add(sub)
                entity_list["data"].append({"name": obj, "category": rel})
                entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'query_system':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = str(i['m.name'])
                obj = str(i['n.name'])
                rel = str(i['r.name'])
                ossid = i['n.doc_id']
                file_name = i['n.doc_name']
                filePage = i['n.filePage']
                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "系统"})
                    seen_subjects.add(sub)
                entity_list["data"].append({"name": obj, "category": rel})
                entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'fault_des':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = str(i['m.name'])
                obj = str(i['n.name'])
                rel = str(i['r.name'])
                ossid = i['n.doc_id']
                file_name = i['n.doc_name']
                filePage = i['n.filePage']
                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "HMC"})
                    seen_subjects.add(sub)
                entity_list["data"].append({"name": obj, "category": rel})
                entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'fault_excluds':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = str(i['m.name'])
                obj = str(i['n.name'])
                rel = str(i['r.name'])
                ossid = i['n.doc_id']
                file_name = i['n.doc_name']
                filePage = i['n.filePage']

                a = i['a.name'] if 'a.name' in i else None
                b = i['b.name'] if 'b.name' in i else None
                c = i['c.name'] if 'c.name' in i else None

                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "HMC"})
                    seen_subjects.add(sub)

                if a and b and c:
                    entity_list["data"].append({"name": c, "category": b})
                    entity_list["links"].append({"source": a, "target": c, "value": b})
                    entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})
                else:
                    entity_list["data"].append({"name": obj, "category": rel})
                    entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                    entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'des_excluds':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = str(i['m.name'])
                obj = str(i['n.name'])
                rel = str(i['r.name'])
                ossid = i['n.doc_id']
                file_name = i['n.doc_name']
                filePage = i['n.filePage']

                a = i['a.name'] if 'a.name' in i else None
                b = i['b.name'] if 'b.name' in i else None
                c = i['c.name'] if 'c.name' in i else None

                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "故障名称"})
                    seen_subjects.add(sub)

                if a and b and c:
                    entity_list["data"].append({"name": c, "category": b})
                    entity_list["links"].append({"source": a, "target": c, "value": b})
                    entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})
                else:
                    entity_list["data"].append({"name": obj, "category": rel})
                    entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                    entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'fault_obj':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = str(i['m.name'])
                obj = str(i['n.name'])
                rel = str(i['r.name'])
                ossid = i['n.doc_id']
                file_name = i['n.doc_name']
                filePage = i['n.filePage']
                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "HMC"})
                    seen_subjects.add(sub)
                entity_list["data"].append({"name": obj, "category": rel})
                entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'hmc_system':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = str(i['m.name'])
                obj = str(i['n.name'])
                rel = str(i['r.name'])
                ossid = i['n.doc_id']
                file_name = i['n.doc_name']
                filePage = i['n.filePage']
                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "HMC"})
                    seen_subjects.add(sub)
                entity_list["data"].append({"name": obj, "category": rel})
                entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'sys_s_sys':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = str(i['m.name'])
                obj = str(i['n.name'])
                rel = str(i['r.name'])
                ossid = i['n.doc_id']
                file_name = str(i['n.doc_name'])
                filePage = i['n.filePage']
                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "系统"})
                    seen_subjects.add(sub)
                entity_list["data"].append({"name": obj, "category": rel})
                entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'sys_fault_obj':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = i['m.name']
                obj = i['n.name']
                rel = i['r.name']
                ossid = i['n.doc_id']
                file_name = i['n.doc_name']
                filePage = i['n.filePage']
                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "系统"})
                    seen_subjects.add(sub)
                entity_list["data"].append({"name": obj, "category": rel})
                entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'fault_obj_sys':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = i['m.name']
                obj = i['n.name']
                rel = i['r.name']
                ossid = i['m.doc_id']
                file_name = i['m.doc_name']
                filePage = i['n.filePage']
                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "成品"})
                    seen_subjects.add(sub)
                entity_list["data"].append({"name": obj, "category": rel})
                entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        elif question_type == 'fault_obj_fault_x_obj':
            entity_list = {
                "data": [],
                "links": [],
                "docs": []
            }
            seen_subjects = set()
            for i in answers:
                sub = i['m.name']
                obj = i['n.name']
                rel = i['r.name']
                ossid = i['n.doc_id']
                file_name = i['n.doc_name']
                filePage = i['n.filePage']
                if sub not in seen_subjects:
                    entity_list["data"].append({"name": sub, "category": "成品"})
                    seen_subjects.add(sub)
                entity_list["data"].append({"name": obj, "category": rel})
                entity_list["links"].append({"source": sub, "target": obj, "value": rel})
                entity_list['docs'].append({"ossId": ossid, "file_name": file_name, "filePage": filePage})

        return entity_list

    '''根据对应的qustion_type，调用相应的回复模板'''
    def answer_prettify(self, question_type, answers):
        final_answer = []
        if not answers:
            return ''

        if question_type == 'query_des':
            # 初始化集合用于存储去重后的尾实体及其关系
            desc_with_rel_set = set()

                # 遍历答案，同时获取头实体、关系和尾实体
            for answer in answers:
                subject = answer['m.name']
                rel = answer['r.name']
                desc = answer['n.name']

                # 将关系和尾实体组合成字符串并加入集合
                desc_with_rel_set.add(f"{rel}为{desc}")

            # 将集合转换为列表，并按需限制数量
            desc_with_rel = list(desc_with_rel_set)[:self.num_limit]

            # 将头实体和与其有关联的尾实体及关系组合成最终的回答
            final_answer = '故障描述{0}{1}。'.format(subject, '; '.join(desc_with_rel))

        elif question_type == 'query_hmc':
            # 初始化集合用于存储去重后的尾实体及其关系
            desc_with_rel_set = set()

            # 遍历答案，同时获取头实体、关系和尾实体
            for answer in answers:
                subject = answer['m.name']
                rel = answer['r.name']
                desc = answer['n.name']

                # 将关系和尾实体组合成字符串并加入集合
                desc_with_rel_set.add(f"{rel}为{desc}")

            # 将集合转换为列表，并按需限制数量
            desc_with_rel = list(desc_with_rel_set)[:self.num_limit]

            # 将头实体和与其有关联的尾实体及关系组合成最终的回答
            final_answer = 'HMC码{0}{1}。'.format(subject, '; '.join(desc_with_rel))

        elif question_type == 'query_obj':
            # 初始化集合用于存储去重后的尾实体及其关系
            desc_with_rel_set = set()

            # 遍历答案，同时获取头实体、关系和尾实体
            for answer in answers:
                subject = answer['m.name']
                rel = answer['r.name']
                desc = answer['n.name']

                # 将关系和尾实体组合成字符串并加入集合
                desc_with_rel_set.add(f"{rel}为{desc}")

            # 将集合转换为列表，并按需限制数量
            desc_with_rel = list(desc_with_rel_set)[:self.num_limit]

            # 将头实体和与其有关联的尾实体及关系组合成最终的回答
            final_answer = '成品{0}{1}。'.format(subject, '; '.join(desc_with_rel))

        elif question_type == 'query_x_obj':
            # 初始化集合用于存储去重后的尾实体及其关系
            desc_with_rel_set = set()

            # 遍历答案，同时获取头实体、关系和尾实体
            for answer in answers:
                subject = answer['m.name']
                rel = answer['r.name']
                desc = answer['n.name']

                # 将关系和尾实体组合成字符串并加入集合
                desc_with_rel_set.add(f"{rel}为{desc}")

            # 将集合转换为列表，并按需限制数量
            desc_with_rel = list(desc_with_rel_set)[:self.num_limit]

            # 将头实体和与其有关联的尾实体及关系组合成最终的回答
            final_answer = '型号{0}{1}。'.format(subject, '; '.join(desc_with_rel))

        elif question_type == 'query_s_sys':
            # 初始化集合用于存储去重后的尾实体及其关系
            desc_with_rel_set = set()

            # 遍历答案，同时获取头实体、关系和尾实体
            for answer in answers:
                subject = answer['m.name']
                rel = answer['r.name']
                desc = answer['n.name']

                # 将关系和尾实体组合成字符串并加入集合
                desc_with_rel_set.add(f"{rel}为{desc}")

            # 将集合转换为列表，并按需限制数量
            desc_with_rel = list(desc_with_rel_set)[:self.num_limit]

            # 将头实体和与其有关联的尾实体及关系组合成最终的回答
            final_answer = '子系统{0}{1}。'.format(subject, '; '.join(desc_with_rel))

        elif question_type == 'query_system':
            # 初始化集合用于存储去重后的尾实体及其关系
            desc_with_rel_set = set()

            # 遍历答案，同时获取头实体、关系和尾实体
            for answer in answers:
                subject = answer['m.name']
                rel = answer['r.name']
                desc = answer['n.name']

                # 将关系和尾实体组合成字符串并加入集合
                desc_with_rel_set.add(f"{rel}为{desc}")

            # 将集合转换为列表，并按需限制数量
            desc_with_rel = list(desc_with_rel_set)[:self.num_limit]

            # 将头实体和与其有关联的尾实体及关系组合成最终的回答
            final_answer = '系统{0}{1}。'.format(subject, '; '.join(desc_with_rel))

        elif question_type == 'fault_des':
            desc = [i['n.name'] for i in answers]
            subject = answers[0]['m.name']
            final_answer = '{0}的故障现象是: {1}。'.format(subject, ';'.join(list(set(desc))[:self.num_limit]))

        elif question_type == 'fault_excluds':
            desc = []
            for i in answers:
                if 'c.name' in i and i['c.name']:
                    desc.append(i['c.name'])
                elif 'n.name' in i:
                    desc.append(i['n.name'])

            subject = answers[0]['m.name']
            final_answer = '发生{0}时应该{1}。'.format(subject, ';'.join(list(set(desc))[:self.num_limit]))

        elif question_type == 'des_excluds':
            desc = []
            for i in answers:
                if 'c.name' in i and i['c.name']:
                    desc.append(i['c.name'])
                elif 'n.name' in i:
                    desc.append(i['n.name'])

            subject = answers[0]['m.name']
            final_answer = '发生{0}时应该{1}。'.format(subject, ';'.join(list(set(desc))[:self.num_limit]))

        elif question_type == 'fault_obj':
            desc = [i['n.name'] for i in answers]
            subject = answers[0]['m.name']
            final_answer = 'HMC码{0}对应的部件位：{1}。'.format(subject, ';'.join(list(set(desc))[:self.num_limit]))

        elif question_type == 'hmc_system':
            desc = [i['n.name'] for i in answers]
            subject = answers[0]['m.name']
            final_answer = 'HMC码{0}对应的系统为{1}。'.format(subject, ';'.join(list(set(desc))[:self.num_limit]))

        elif question_type == "sys_s_sys":
            desc = [i['n.name'] for i in answers]
            subject = answers[0]['m.name']
            final_answer = '{0}包含的子系统有：{1}。'.format(subject, ';'.join(list(set(desc))[:self.num_limit]))

        elif question_type == "sys_fault_obj":
            desc = [i['n.name'] for i in answers]
            subject = answers[0]['m.name']
            final_answer = '{0}包含的成品有：{1}。'.format(subject, ';'.join(list(set(desc))[:self.num_limit]))

        elif question_type == "fault_obj_sys":
            desc = [i['n.name'] for i in answers]
            subject = answers[0]['m.name']
            final_answer = '{0}属于{1}。'.format(subject, ';'.join(list(set(desc))[:self.num_limit]))

        elif question_type == "fault_obj_fault_x_obj":
            desc = [i['n.name'] for i in answers]
            subject = answers[0]['m.name']
            final_answer = '{0}的型号为:{1}。'.format(subject, ';'.join(list(set(desc))[:self.num_limit]))

        return final_answer


if __name__ == '__main__':
    searcher = AnswerSearcher()