api/models/provider.py
Rama main · réplica de lectura del repositorio original.
from __future__ import annotations
from datetime import datetime
from enum import StrEnum, auto
from typing import override
from uuid import uuid4
import sqlalchemy as sa
from sqlalchemy import DateTime, String, func, text
from sqlalchemy.orm import Mapped, foreign, mapped_column, relationship
from graphon.model_runtime.entities.model_entities import ModelType
from libs.uuid_utils import uuidv7
from .base import TypeBase
from .enums import CredentialSourceType, PaymentStatus, PermissionEnum, ProviderQuotaType
from .types import EnumText, LongText, StringUUID
class ProviderType(StrEnum):
CUSTOM = auto()
SYSTEM = auto()
@staticmethod
def value_of(value: str) -> ProviderType:
for member in ProviderType:
if member.value == value:
return member
raise ValueError(f"No matching enum found for value '{value}'")
class Provider(TypeBase):
"""
Provider model representing the API providers and their configurations.
"""
__tablename__ = "providers"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="provider_pkey"),
sa.Index("provider_tenant_id_provider_idx", "tenant_id", "provider_name"),
sa.UniqueConstraint(
"tenant_id", "provider_name", "provider_type", "quota_type", name="unique_provider_name_type_quota"
),
)
id: Mapped[str] = mapped_column(
StringUUID,
primary_key=True,
insert_default=lambda: str(uuidv7()),
default_factory=lambda: str(uuidv7()),
init=False,
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
provider_type: Mapped[ProviderType] = mapped_column(
EnumText(ProviderType, length=40), nullable=False, default=ProviderType.CUSTOM
)
is_valid: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=False)
last_used: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, init=False)
credential_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
quota_type: Mapped[ProviderQuotaType | None] = mapped_column(
EnumText(ProviderQuotaType, length=40), nullable=True, server_default=text("''"), default=None
)
quota_limit: Mapped[int | None] = mapped_column(sa.BigInteger, nullable=True, default=None)
quota_used: Mapped[int | None] = mapped_column(sa.BigInteger, nullable=True, default=0)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)
# The active credential row, or None when credential_id is unset or the row was deleted.
# lazy="raise" keeps reads from opening a session: callers load it with selectinload in their own session.
credential: Mapped[ProviderCredential | None] = relationship(
lambda: ProviderCredential,
primaryjoin=lambda: foreign(Provider.credential_id) == ProviderCredential.id,
lazy="raise",
viewonly=True,
init=False,
compare=False,
)
@override
def __repr__(self):
return (
f"<Provider(id={self.id}, tenant_id={self.tenant_id}, provider_name='{self.provider_name}',"
f" provider_type='{self.provider_type}')>"
)
@property
def credential_name(self):
# Records without credential_id, such as newly created system records, never need the relationship loaded.
credential = self.credential if self.credential_id else None
return credential.credential_name if credential else None
@property
def encrypted_config(self):
credential = self.credential if self.credential_id else None
return credential.encrypted_config if credential else None
@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:
return self.is_valid
else:
return self.is_valid and self.token_is_set
class ProviderModel(TypeBase):
"""
Provider model representing the API provider_models and their configurations.
"""
__tablename__ = "provider_models"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="provider_model_pkey"),
sa.Index("provider_model_tenant_id_provider_idx", "tenant_id", "provider_name"),
sa.UniqueConstraint(
"tenant_id", "provider_name", "model_name", "model_type", name="unique_provider_model_name"
),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_type: Mapped[ModelType] = mapped_column(EnumText(ModelType, length=40), nullable=False)
credential_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
is_valid: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=False)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)
class TenantDefaultModel(TypeBase):
__tablename__ = "tenant_default_models"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="tenant_default_model_pkey"),
sa.Index("tenant_default_model_tenant_id_provider_type_idx", "tenant_id", "provider_name", "model_type"),
sa.UniqueConstraint("tenant_id", "model_type", name="unique_tenant_default_model_type"),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_type: Mapped[ModelType] = mapped_column(EnumText(ModelType, length=40), nullable=False)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)
class TenantPreferredModelProvider(TypeBase):
__tablename__ = "tenant_preferred_model_providers"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="tenant_preferred_model_provider_pkey"),
sa.Index("tenant_preferred_model_provider_tenant_provider_idx", "tenant_id", "provider_name"),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
preferred_provider_type: Mapped[ProviderType] = mapped_column(EnumText(ProviderType, length=40), nullable=False)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)
class ProviderOrder(TypeBase):
__tablename__ = "provider_orders"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="provider_order_pkey"),
sa.Index("provider_order_tenant_provider_idx", "tenant_id", "provider_name"),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
account_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
payment_product_id: Mapped[str] = mapped_column(String(191), nullable=False)
payment_id: Mapped[str | None] = mapped_column(String(191))
transaction_id: Mapped[str | None] = mapped_column(String(191))
quantity: Mapped[int] = mapped_column(sa.Integer, nullable=False, server_default=text("1"))
currency: Mapped[str | None] = mapped_column(String(40))
total_amount: Mapped[int | None] = mapped_column(sa.Integer)
payment_status: Mapped[PaymentStatus] = mapped_column(
EnumText(PaymentStatus, length=40), nullable=False, server_default=text("'wait_pay'")
)
paid_at: Mapped[datetime | None] = mapped_column(DateTime)
pay_failed_at: Mapped[datetime | None] = mapped_column(DateTime)
refunded_at: Mapped[datetime | None] = mapped_column(DateTime)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)
class ProviderModelSetting(TypeBase):
"""
Provider model settings for record the model enabled status and load balancing status.
"""
__tablename__ = "provider_model_settings"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="provider_model_setting_pkey"),
sa.Index("provider_model_setting_tenant_provider_model_idx", "tenant_id", "provider_name", "model_type"),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_type: Mapped[ModelType] = mapped_column(EnumText(ModelType, length=40), nullable=False)
enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=True)
load_balancing_enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=False)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)
class LoadBalancingModelConfig(TypeBase):
"""
Configurations for load balancing models.
"""
__tablename__ = "load_balancing_model_configs"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="load_balancing_model_config_pkey"),
sa.Index("load_balancing_model_config_tenant_provider_model_idx", "tenant_id", "provider_name", "model_type"),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_type: Mapped[ModelType] = mapped_column(EnumText(ModelType, length=40), nullable=False)
name: Mapped[str] = mapped_column(String(255), nullable=False)
encrypted_config: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None)
credential_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
credential_source_type: Mapped[CredentialSourceType | None] = mapped_column(
EnumText(CredentialSourceType, length=40), nullable=True, default=None
)
enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=True)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)
class ProviderCredential(TypeBase):
"""
Provider credential - stores multiple named credentials for each provider
"""
__tablename__ = "provider_credentials"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="provider_credential_pkey"),
sa.Index("provider_credential_tenant_provider_idx", "tenant_id", "provider_name"),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
credential_name: Mapped[str] = mapped_column(String(255), nullable=False)
encrypted_config: Mapped[str] = mapped_column(LongText, nullable=False)
user_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
visibility: Mapped[PermissionEnum] = mapped_column(
EnumText(PermissionEnum, length=40), nullable=False, default=PermissionEnum.ALL_TEAM
)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)
class ProviderModelCredential(TypeBase):
"""
Provider model credential - stores multiple named credentials for each provider model
"""
__tablename__ = "provider_model_credentials"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="provider_model_credential_pkey"),
sa.Index(
"provider_model_credential_tenant_provider_model_idx",
"tenant_id",
"provider_name",
"model_name",
"model_type",
),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_type: Mapped[ModelType] = mapped_column(EnumText(ModelType, length=40), nullable=False)
credential_name: Mapped[str] = mapped_column(String(255), nullable=False)
encrypted_config: Mapped[str] = mapped_column(LongText, nullable=False)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
)