from enum import Enum from sqlalchemy.dialects.postgresql import UUID from extensions.ext_database import db class ProviderType(Enum): CUSTOM = 'custom' SYSTEM = 'system' class ProviderName(Enum): OPENAI = 'openai' AZURE_OPENAI = 'azure_openai' ANTHROPIC = 'anthropic' COHERE = 'cohere' HUGGINGFACEHUB = 'huggingfacehub' @staticmethod def value_of(value): for member in ProviderName: if member.value == value: return member raise ValueError(f"No matching enum found for value '{value}'") class ProviderQuotaType(Enum): MONTHLY = 'monthly' TRIAL = 'trial' class Provider(db.Model): """ Provider model representing the API providers and their configurations. """ __tablename__ = 'providers' __table_args__ = ( db.PrimaryKeyConstraint('id', name='provider_pkey'), db.Index('provider_tenant_id_provider_idx', 'tenant_id', 'provider_name'), db.UniqueConstraint('tenant_id', 'provider_name', 'provider_type', 'quota_type', name='unique_provider_name_type_quota') ) id = db.Column(UUID, server_default=db.text('uuid_generate_v4()')) tenant_id = db.Column(UUID, nullable=False) provider_name = db.Column(db.String(40), nullable=False) provider_type = db.Column(db.String(40), nullable=False, server_default=db.text("'custom'::character varying")) encrypted_config = db.Column(db.Text, nullable=True) is_valid = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) last_used = db.Column(db.DateTime, nullable=True) quota_type = db.Column(db.String(40), nullable=True, server_default=db.text("''::character varying")) quota_limit = db.Column(db.Integer, nullable=True) quota_used = db.Column(db.Integer, default=0) created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) def __repr__(self): return f"" @property def token_is_set(self): """ Returns True if the encrypted_config is not None, indicating that the token is set. """ return self.encrypted_config is not None @property def is_enabled(self): """ Returns True if the provider is enabled. """ if self.provider_type == ProviderType.SYSTEM.value: return self.is_valid else: return self.is_valid and self.token_is_set