Underground-Digital/Workflow-Engine

1from enum import Enum2 3from extensions.ext_database import db4 5from .types import StringUUID6 7 8class ProviderType(Enum):9    CUSTOM = "custom"10    SYSTEM = "system"11 12    @staticmethod13    def value_of(value):14        for member in ProviderType:15            if member.value == value:16                return member17        raise ValueError(f"No matching enum found for value '{value}'")18 19 20class ProviderQuotaType(Enum):21    PAID = "paid"22    """hosted paid quota"""23 24    FREE = "free"25    """third-party free quota"""26 27    TRIAL = "trial"28    """hosted trial quota"""29 30    @staticmethod31    def value_of(value):32        for member in ProviderQuotaType:33            if member.value == value:34                return member35        raise ValueError(f"No matching enum found for value '{value}'")36 37 38class Provider(db.Model):39    """40    Provider model representing the API providers and their configurations.41    """42 43    __tablename__ = "providers"44    __table_args__ = (45        db.PrimaryKeyConstraint("id", name="provider_pkey"),46        db.Index("provider_tenant_id_provider_idx", "tenant_id", "provider_name"),47        db.UniqueConstraint(48            "tenant_id", "provider_name", "provider_type", "quota_type", name="unique_provider_name_type_quota"49        ),50    )51 52    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))53    tenant_id = db.Column(StringUUID, nullable=False)54    provider_name = db.Column(db.String(255), nullable=False)55    provider_type = db.Column(db.String(40), nullable=False, server_default=db.text("'custom'::character varying"))56    encrypted_config = db.Column(db.Text, nullable=True)57    is_valid = db.Column(db.Boolean, nullable=False, server_default=db.text("false"))58    last_used = db.Column(db.DateTime, nullable=True)59 60    quota_type = db.Column(db.String(40), nullable=True, server_default=db.text("''::character varying"))61    quota_limit = db.Column(db.BigInteger, nullable=True)62    quota_used = db.Column(db.BigInteger, default=0)63 64    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))65    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))66 67    def __repr__(self):68        return (69            f"<Provider(id={self.id}, tenant_id={self.tenant_id}, provider_name='{self.provider_name}',"70            f" provider_type='{self.provider_type}')>"71        )72 73    @property74    def token_is_set(self):75        """76        Returns True if the encrypted_config is not None, indicating that the token is set.77        """78        return self.encrypted_config is not None79 80    @property81    def is_enabled(self):82        """83        Returns True if the provider is enabled.84        """85        if self.provider_type == ProviderType.SYSTEM.value:86            return self.is_valid87        else:88            return self.is_valid and self.token_is_set89 90 91class ProviderModel(db.Model):92    """93    Provider model representing the API provider_models and their configurations.94    """95 96    __tablename__ = "provider_models"97    __table_args__ = (98        db.PrimaryKeyConstraint("id", name="provider_model_pkey"),99        db.Index("provider_model_tenant_id_provider_idx", "tenant_id", "provider_name"),100        db.UniqueConstraint(101            "tenant_id", "provider_name", "model_name", "model_type", name="unique_provider_model_name"102        ),103    )104 105    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))106    tenant_id = db.Column(StringUUID, nullable=False)107    provider_name = db.Column(db.String(255), nullable=False)108    model_name = db.Column(db.String(255), nullable=False)109    model_type = db.Column(db.String(40), nullable=False)110    encrypted_config = db.Column(db.Text, nullable=True)111    is_valid = db.Column(db.Boolean, nullable=False, server_default=db.text("false"))112    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))113    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))114 115 116class TenantDefaultModel(db.Model):117    __tablename__ = "tenant_default_models"118    __table_args__ = (119        db.PrimaryKeyConstraint("id", name="tenant_default_model_pkey"),120        db.Index("tenant_default_model_tenant_id_provider_type_idx", "tenant_id", "provider_name", "model_type"),121    )122 123    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))124    tenant_id = db.Column(StringUUID, nullable=False)125    provider_name = db.Column(db.String(255), nullable=False)126    model_name = db.Column(db.String(255), nullable=False)127    model_type = db.Column(db.String(40), nullable=False)128    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))129    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))130 131 132class TenantPreferredModelProvider(db.Model):133    __tablename__ = "tenant_preferred_model_providers"134    __table_args__ = (135        db.PrimaryKeyConstraint("id", name="tenant_preferred_model_provider_pkey"),136        db.Index("tenant_preferred_model_provider_tenant_provider_idx", "tenant_id", "provider_name"),137    )138 139    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))140    tenant_id = db.Column(StringUUID, nullable=False)141    provider_name = db.Column(db.String(255), nullable=False)142    preferred_provider_type = db.Column(db.String(40), nullable=False)143    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))144    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))145 146 147class ProviderOrder(db.Model):148    __tablename__ = "provider_orders"149    __table_args__ = (150        db.PrimaryKeyConstraint("id", name="provider_order_pkey"),151        db.Index("provider_order_tenant_provider_idx", "tenant_id", "provider_name"),152    )153 154    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))155    tenant_id = db.Column(StringUUID, nullable=False)156    provider_name = db.Column(db.String(255), nullable=False)157    account_id = db.Column(StringUUID, nullable=False)158    payment_product_id = db.Column(db.String(191), nullable=False)159    payment_id = db.Column(db.String(191))160    transaction_id = db.Column(db.String(191))161    quantity = db.Column(db.Integer, nullable=False, server_default=db.text("1"))162    currency = db.Column(db.String(40))163    total_amount = db.Column(db.Integer)164    payment_status = db.Column(db.String(40), nullable=False, server_default=db.text("'wait_pay'::character varying"))165    paid_at = db.Column(db.DateTime)166    pay_failed_at = db.Column(db.DateTime)167    refunded_at = db.Column(db.DateTime)168    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))169    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))170 171 172class ProviderModelSetting(db.Model):173    """174    Provider model settings for record the model enabled status and load balancing status.175    """176 177    __tablename__ = "provider_model_settings"178    __table_args__ = (179        db.PrimaryKeyConstraint("id", name="provider_model_setting_pkey"),180        db.Index("provider_model_setting_tenant_provider_model_idx", "tenant_id", "provider_name", "model_type"),181    )182 183    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))184    tenant_id = db.Column(StringUUID, nullable=False)185    provider_name = db.Column(db.String(255), nullable=False)186    model_name = db.Column(db.String(255), nullable=False)187    model_type = db.Column(db.String(40), nullable=False)188    enabled = db.Column(db.Boolean, nullable=False, server_default=db.text("true"))189    load_balancing_enabled = db.Column(db.Boolean, nullable=False, server_default=db.text("false"))190    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))191    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))192 193 194class LoadBalancingModelConfig(db.Model):195    """196    Configurations for load balancing models.197    """198 199    __tablename__ = "load_balancing_model_configs"200    __table_args__ = (201        db.PrimaryKeyConstraint("id", name="load_balancing_model_config_pkey"),202        db.Index("load_balancing_model_config_tenant_provider_model_idx", "tenant_id", "provider_name", "model_type"),203    )204 205    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))206    tenant_id = db.Column(StringUUID, nullable=False)207    provider_name = db.Column(db.String(255), nullable=False)208    model_name = db.Column(db.String(255), nullable=False)209    model_type = db.Column(db.String(40), nullable=False)210    name = db.Column(db.String(255), nullable=False)211    encrypted_config = db.Column(db.Text, nullable=True)212    enabled = db.Column(db.Boolean, nullable=False, server_default=db.text("true"))213    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))214    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))215