General-AI-Platform-Backend/website/db/alg_model/alg_model.py

# -*- coding: utf-8 -*-

import logging
from typing import Union

from sqlalchemy import Column, Integer, String, DateTime, func
from sqlalchemy.ext.declarative import declarative_base

from website.db_mysql import get_session

Base = declarative_base()

"""
CREATE TABLE `model` (
  `id` int NOT NULL AUTO_INCREMENT,
  `suid` varchar(10) DEFAULT NULL COMMENT 'short uuid',
  `name` varchar(255) NOT NULL DEFAULT '',
  `model_type` int DEFAULT '1002' COMMENT '模型类型，1001/经典算法，1002/深度学习',
  `classification` int DEFAULT '0' COMMENT '模型分类的id',
  `comment` varchar(255) DEFAULT '' COMMENT '备注',
  `default_version` varchar(100) DEFAULT '',
  `del` tinyint(1) DEFAULT '0' COMMENT '删除状态，1/删除，0/正常',
  `create_time` datetime DEFAULT CURRENT_TIMESTAMP,
  `update_time` datetime DEFAULT NULL,
  PRIMARY KEY (`id`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci COMMENT='模型表';
"""


class Model(Base):
    __tablename__ = "model"

    id = Column(Integer, primary_key=True, autoincrement=True)
    suid = Column(String(10), comment="short uuid")
    name = Column(String(255), nullable=False, default="")
    model_type = Column(Integer, default=1002, comment="模型类型，1001/经典算法，1002/深度学习")
    classification = Column(Integer, default=0, comment="模型分类的id")
    comment = Column(String(255), default="", comment="备注")
    default_version = Column(String(100), default="")
    delete = Column("del", Integer, default=0, comment="删除状态，1/删除，0/正常")
    create_time = Column(DateTime, default=func.now())
    update_time = Column(DateTime, onupdate=func.now())

    def __repr__(self):
        return f"Model(id={self.id}, name='{self.name}', model_type={self.model_type})"


class ModelRepositry(object):

    def get_suid(self, model_id: int) -> str:
        with get_session() as session:
            model = session.query(Model).filter(Model.id == model_id).first()
            if not model or not model.suid:
                return ""

            return model.suid

    def get_model_by_id(self, model_id: int) -> Union[Model, None]:

        with get_session() as session:
            model = session.query(Model).filter(Model.id == model_id).first()
            if not model:
                return None

            return model

    def get_model_dict_by_id(self, model_id: int) -> dict:
        with get_session() as session:
            logging.info(f"model id is : {model_id}")
            model = session.query(Model).filter(Model.id == model_id).first()
            if not model:
                return {}

            model_dict = {
                'id': model.id,
                'suid': model.suid,
                'name': model.name,
                'model_type': model.model_type,
                'classification': model.classification,
                'comment': model.comment,
                'default_version': model.default_version,
                'delete': model.delete,
                'create_time': model.create_time,
                'update_time': model.update_time
            }

            logging.info(f"model dict is : {model_dict}")
            return model_dict

    def get_model_by_ids(self, model_ids: list) -> list:
        with get_session() as session:
            models = session.query(Model).filter(Model.id.in_(model_ids)).all()
            if not models:
                return []

            return models

    def get_model_count(self) -> int:
        with get_session() as session:
            return session.query(Model).filter(Model.delete == 0).count()