dify/api/models/provider.py
2023-05-15 08:51:32 +08:00

78 lines
2.6 KiB
Python

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"<Provider(id={self.id}, tenant_id={self.tenant_id}, provider_name='{self.provider_name}', provider_type='{self.provider_type}')>"
@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