rag-gateway/app/service/label.py

import uuid

from Log import logger
from app.models.label_model import LabelModel, LabelWorkerModel
from app.models.role_model import RoleModel


async def label_list_service(db, keyword: str, label_type, object_id):
    query = db.query(LabelModel)
    if keyword:
        query = query.filter(LabelModel.name.like('%{}%'.format(keyword)))
    if label_type:
        query = query.filter(LabelModel.label_type==label_type)
    if object_id:
        label_id = [i.label_id for i in db.query(LabelWorkerModel).filter(LabelWorkerModel.object_id==object_id).all()]
        query = query.filter(LabelModel.id.in_(label_id))
    labels = query.order_by(LabelModel.id.desc()).all()
    return {"total": query.count(), "rows":  [label.to_json() for label in labels]}



async def create_label_service(db, label_name, label_type, user_id):
    try:
        label_model = LabelModel(name=label_name,creator=user_id, label_type=label_type)
        db.add(label_model)
        db.commit()
        db.refresh(label_model)
    except Exception as e:
        logger.error(e)
        db.rollback()
        return False
    return True


async def delete_role_service(db, label_id: int):
    try:
        db.query(LabelModel).filter(LabelModel.id == label_id).delete()
        db.commit()
    except Exception as e:
        logger.error(e)
        db.rollback()
        return False
    return True


async def sign_label_service(db, object_id, label_list):
    delete_list = []
    has_list = []
    for i in db.query(LabelWorkerModel).filter(LabelWorkerModel.object_id == object_id).all():
        if i.label_id not in label_list:
            delete_list.append(i.id)
        else:
            has_list.append(i.label_id)
    for label_id in label_list:
        if label_id in has_list:
            continue
        try:
            label = LabelWorkerModel(label_id=label_id, object_id=object_id)
            db.add(label)
            db.commit()
        except Exception as e:
            logger.error(e)
            db.rollback()
            # return False
    for label_id in delete_list:
        try:
            LabelWorkerModel(id=label_id).delete()
        except Exception as e:
            logger.error(e)
            db.rollback()
    return True