mirror of
https://github.com/langgenius/dify.git
synced 2024-11-16 19:59:50 +08:00
89 lines
2.1 KiB
Python
89 lines
2.1 KiB
Python
from enum import Enum
|
|
from typing import Optional
|
|
|
|
from pydantic import BaseModel, ConfigDict
|
|
|
|
from core.model_runtime.entities.common_entities import I18nObject
|
|
from core.model_runtime.entities.model_entities import ModelType, ProviderModel
|
|
from core.model_runtime.entities.provider_entities import ProviderEntity
|
|
|
|
|
|
class ModelStatus(Enum):
|
|
"""
|
|
Enum class for model status.
|
|
"""
|
|
|
|
ACTIVE = "active"
|
|
NO_CONFIGURE = "no-configure"
|
|
QUOTA_EXCEEDED = "quota-exceeded"
|
|
NO_PERMISSION = "no-permission"
|
|
DISABLED = "disabled"
|
|
|
|
|
|
class SimpleModelProviderEntity(BaseModel):
|
|
"""
|
|
Simple provider.
|
|
"""
|
|
|
|
provider: str
|
|
label: I18nObject
|
|
icon_small: Optional[I18nObject] = None
|
|
icon_large: Optional[I18nObject] = None
|
|
supported_model_types: list[ModelType]
|
|
|
|
def __init__(self, provider_entity: ProviderEntity) -> None:
|
|
"""
|
|
Init simple provider.
|
|
|
|
:param provider_entity: provider entity
|
|
"""
|
|
super().__init__(
|
|
provider=provider_entity.provider,
|
|
label=provider_entity.label,
|
|
icon_small=provider_entity.icon_small,
|
|
icon_large=provider_entity.icon_large,
|
|
supported_model_types=provider_entity.supported_model_types,
|
|
)
|
|
|
|
|
|
class ProviderModelWithStatusEntity(ProviderModel):
|
|
"""
|
|
Model class for model response.
|
|
"""
|
|
|
|
status: ModelStatus
|
|
load_balancing_enabled: bool = False
|
|
|
|
|
|
class ModelWithProviderEntity(ProviderModelWithStatusEntity):
|
|
"""
|
|
Model with provider entity.
|
|
"""
|
|
|
|
provider: SimpleModelProviderEntity
|
|
|
|
|
|
class DefaultModelProviderEntity(BaseModel):
|
|
"""
|
|
Default model provider entity.
|
|
"""
|
|
|
|
provider: str
|
|
label: I18nObject
|
|
icon_small: Optional[I18nObject] = None
|
|
icon_large: Optional[I18nObject] = None
|
|
supported_model_types: list[ModelType]
|
|
|
|
|
|
class DefaultModelEntity(BaseModel):
|
|
"""
|
|
Default model entity.
|
|
"""
|
|
|
|
model: str
|
|
model_type: ModelType
|
|
provider: DefaultModelProviderEntity
|
|
|
|
# pydantic configs
|
|
model_config = ConfigDict(protected_namespaces=())
|