tjdify/api/models/intention.py

from sqlalchemy import func

from .engine import db
from .types import StringUUID


class IntentionType(db.Model):
    __tablename__ = 'intention_types'
    __table_args__ = (
        db.PrimaryKeyConstraint('id', name='intention_type_id_pkey'),
    )

    id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()"))
    name = db.Column(db.String(255), nullable=False)
    created_by = db.Column(StringUUID, nullable=False)
    created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
    updated_by = db.Column(StringUUID, nullable=True)
    updated_at = db.Column(db.DateTime, nullable=True, server_default=func.current_timestamp())

    @property
    def intention_count(self):
        return (
            db.session.query(func.count(Intention.id))
            .filter(Intention.type_id==self.id)
            .scalar()
        )

    @property
    def intentions(self):
        return (
            db.session.query(Intention)
            .filter(Intention.type_id == self.id)
            .all()
        )

class Intention(db.Model):
    __tablename__ = "intentions"
    __table_args__ = (
        db.PrimaryKeyConstraint("id", name="intention_pkey"),
        db.Index("intention_type_idx", "type_id"),
    )

    id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()"))
    type_id = db.Column(StringUUID, nullable=False)
    name = db.Column(db.String(255), nullable=False)
    created_by = db.Column(StringUUID, nullable=False)
    created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
    updated_by = db.Column(StringUUID, nullable=True)
    updated_at = db.Column(db.DateTime, nullable=True, server_default=func.current_timestamp())

    @property
    def type_name(self):
        return (
            db.session.query(IntentionType.name)
            .filter(IntentionType.id==self.type_id)
            .first().name
        )

    @property
    def type(self):
        return (
            db.session.query(IntentionType)
            .filter(IntentionType.id==self.type_id)
            .first()
        )

    @property
    def corpus(self):
        return (
            db.session.query(IntentionCorpus)
            .filter(IntentionCorpus.intention_id==self.id)
            .all()
        )

    @property
    def keywords(self):
        return (
            db.session.query(IntentionKeyword)
            .filter(IntentionKeyword.intention_id==self.id)
            .all()
        )

    @property
    def corpus_count(self):
        return (
            db.session.query(func.count(IntentionCorpus.id))
            .filter(IntentionCorpus.intention_id == self.id)
            .scalar()
        )

    @property
    def keywords_count(self):
        return (
            db.session.query(func.count(IntentionKeyword.id))
            .filter(IntentionKeyword.intention_id == self.id)
            .scalar()
        )

class IntentionKeyword(db.Model):
    __tablename__ = "intention_keywords"
    __table_args__ = (
        db.PrimaryKeyConstraint('id', name='intention_keyword_pkey'),
    )

    id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()"))
    name = db.Column(db.String(255), nullable=False)
    intention_id = db.Column(StringUUID, nullable=False)
    created_by = db.Column(StringUUID, nullable=False)
    created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
    updated_by = db.Column(StringUUID, nullable=True)
    updated_at = db.Column(db.DateTime, nullable=True, server_default=func.current_timestamp())

    @property
    def intention(self):
        return (
            db.session.query(Intention).filter(Intention.id==self.intention_id).first()
        )

class IntentionCorpus(db.Model):
    __tablename__ = "intention_corpus"
    __table_args__ = (
        db.PrimaryKeyConstraint('id', name='intention_corpus_pkey'),
        db.Index("intention_corpus_idx", "intention_id"),
    )

    id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()"))
    question = db.Column(db.String(255), nullable=False)
    question_config = db.Column(db.JSON, nullable=False)
    intention_id = db.Column(StringUUID, nullable=False)
    created_by = db.Column(StringUUID, nullable=False)
    created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
    updated_by = db.Column(StringUUID, nullable=True)
    updated_at = db.Column(db.DateTime, nullable=True, server_default=func.current_timestamp())

    @property
    def intention(self):
        return (
            db.session.query(Intention)
            .filter(Intention.id==self.intention_id)
            .first()
        )

    @property
    def similarity_questions(self):
        return (
            db.session.query(IntentionCorpusSimilarityQuestion)
            .filter(IntentionCorpusSimilarityQuestion.corpus_id==self.id)
            .all()
        )

class IntentionCorpusSimilarityQuestion(db.Model):
    __tablename__ = "intention_corpus_similarity_questions"
    __table_args__ = (
        db.PrimaryKeyConstraint('id', name='intention_corpus_similarity_question_pkey'),
        db.Index("intention_corpus_similarity_question_idx", "corpus_id"),
    )

    id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()"))
    question = db.Column(db.String(255), nullable=False)
    question_config = db.Column(db.JSON, nullable=False)
    corpus_id = db.Column(StringUUID, nullable=False)
    created_by = db.Column(StringUUID, nullable=False)
    created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
    updated_by = db.Column(StringUUID, nullable=True)
    updated_at = db.Column(db.DateTime, nullable=True, server_default=func.current_timestamp())

    @property
    def corpus(self):
        return (
            db.session.query(IntentionCorpus)
            .filter(IntentionCorpus.id==self.corpus_id)
            .first()
        )

class IntentionTrainTask(db.Model):
    __tablename__ = "intention_train_tasks"
    __table_args__ = (
        db.PrimaryKeyConstraint('id', name='intention_train_log_pkey'),
    )

    STATUS_LIST = ["CREATED", "TRAINING", "COMPLETED"]

    id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()"))
    name = db.Column(db.String(255), nullable=False)
    status = db.Column(db.String(255), nullable=False)
    created_by = db.Column(StringUUID, nullable=False)
    created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())

    @property
    def dataset_info(self):
        dataset_info = (
            db.session.query(IntentionTrainFile)
            .join(IntentionTrainFileBinding, IntentionTrainFileBinding.file_id == IntentionTrainFile.id)
            .filter(
                IntentionTrainFile.type == "DATASET",
                IntentionTrainFileBinding.task_id == self.id
            )
            .first()
        )
        return dataset_info

    @property
    def model_info(self):
        model_info = (
            db.session.query(IntentionTrainFile)
            .join(IntentionTrainFileBinding, IntentionTrainFileBinding.file_id == IntentionTrainFile.id)
            .filter(
                IntentionTrainFile.type == "MODEL",
                IntentionTrainFileBinding.task_id == self.id
            )
            .first()
        )
        return model_info

class IntentionTrainFile(db.Model):
    __tablename__ = "intention_train_files"
    __table_args__ = (
        db.PrimaryKeyConstraint('id', name='intention_train_file_pkey'),
    )

    TRAIN_FILE_TYPE_LIST = ["DATASET", "MODEL"]

    id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()"))
    name = db.Column(db.String(255), nullable=False)
    version = db.Column(db.String(255), nullable=False)
    type = db.Column(db.String(255), nullable=False)
    data_source_type = db.Column(db.String(255), nullable=False)
    data_source_info = db.Column(db.JSON, nullable=False)
    created_by = db.Column(StringUUID, nullable=False)
    created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())

class IntentionTrainFileBinding(db.Model):
    __tablename__ = "intention_train_file_bindings"
    __table_args__ = (
        db.PrimaryKeyConstraint('id', name='intention_train_file_binding_pkey'),
        db.Index("intention_train_file_binding_idx", "file_id"),
        db.Index("intention_train_file_binding_task_idx", "task_id"),
    )

    id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()"))
    file_id = db.Column(StringUUID, nullable=False)
    task_id = db.Column(StringUUID, nullable=False)