tjdify/api/core/hosting_configuration.py

import os
from typing import Optional

from core.entities.provider_entities import QuotaUnit, RestrictModel
from core.model_runtime.entities.model_entities import ModelType
from flask import Flask
from models.provider import ProviderQuotaType
from pydantic import BaseModel


class HostingQuota(BaseModel):
    quota_type: ProviderQuotaType
    restrict_models: list[RestrictModel] = []


class TrialHostingQuota(HostingQuota):
    quota_type: ProviderQuotaType = ProviderQuotaType.TRIAL
    quota_limit: int = 0
    """Quota limit for the hosting provider models. -1 means unlimited."""


class PaidHostingQuota(HostingQuota):
    quota_type: ProviderQuotaType = ProviderQuotaType.PAID
    stripe_price_id: str = None
    increase_quota: int = 1
    min_quantity: int = 20
    max_quantity: int = 100


class FreeHostingQuota(HostingQuota):
    quota_type: ProviderQuotaType = ProviderQuotaType.FREE


class HostingProvider(BaseModel):
    enabled: bool = False
    credentials: Optional[dict] = None
    quota_unit: Optional[QuotaUnit] = None
    quotas: list[HostingQuota] = []


class HostedModerationConfig(BaseModel):
    enabled: bool = False
    providers: list[str] = []


class HostingConfiguration:
    provider_map: dict[str, HostingProvider] = {}
    moderation_config: HostedModerationConfig = None

    def init_app(self, app: Flask) -> None:
        if app.config.get('EDITION') != 'CLOUD':
            return

        self.provider_map["azure_openai"] = self.init_azure_openai()
        self.provider_map["openai"] = self.init_openai()
        self.provider_map["anthropic"] = self.init_anthropic()
        self.provider_map["minimax"] = self.init_minimax()
        self.provider_map["spark"] = self.init_spark()
        self.provider_map["zhipuai"] = self.init_zhipuai()

        self.moderation_config = self.init_moderation_config()

    def init_azure_openai(self) -> HostingProvider:
        quota_unit = QuotaUnit.TIMES
        if os.environ.get("HOSTED_AZURE_OPENAI_ENABLED") and os.environ.get("HOSTED_AZURE_OPENAI_ENABLED").lower() == 'true':
            credentials = {
                "openai_api_key": os.environ.get("HOSTED_AZURE_OPENAI_API_KEY"),
                "openai_api_base": os.environ.get("HOSTED_AZURE_OPENAI_API_BASE"),
                "base_model_name": "gpt-35-turbo"
            }

            quotas = []
            hosted_quota_limit = int(os.environ.get("HOSTED_AZURE_OPENAI_QUOTA_LIMIT", "1000"))
            if hosted_quota_limit != -1 or hosted_quota_limit > 0:
                trial_quota = TrialHostingQuota(
                    quota_limit=hosted_quota_limit,
                    restrict_models=[
                        RestrictModel(model="gpt-4", base_model_name="gpt-4", model_type=ModelType.LLM),
                        RestrictModel(model="gpt-4-32k", base_model_name="gpt-4-32k", model_type=ModelType.LLM),
                        RestrictModel(model="gpt-4-1106-preview", base_model_name="gpt-4-1106-preview", model_type=ModelType.LLM),
                        RestrictModel(model="gpt-4-vision-preview", base_model_name="gpt-4-vision-preview", model_type=ModelType.LLM),
                        RestrictModel(model="gpt-35-turbo", base_model_name="gpt-35-turbo", model_type=ModelType.LLM),
                        RestrictModel(model="gpt-35-turbo-1106", base_model_name="gpt-35-turbo-1106", model_type=ModelType.LLM),
                        RestrictModel(model="gpt-35-turbo-instruct", base_model_name="gpt-35-turbo-instruct", model_type=ModelType.LLM),
                        RestrictModel(model="gpt-35-turbo-16k", base_model_name="gpt-35-turbo-16k", model_type=ModelType.LLM),
                        RestrictModel(model="text-davinci-003", base_model_name="text-davinci-003", model_type=ModelType.LLM),
                        RestrictModel(model="text-embedding-ada-002", base_model_name="text-embedding-ada-002", model_type=ModelType.TEXT_EMBEDDING),
                    ]
                )
                quotas.append(trial_quota)

            return HostingProvider(
                enabled=True,
                credentials=credentials,
                quota_unit=quota_unit,
                quotas=quotas
            )

        return HostingProvider(
            enabled=False,
            quota_unit=quota_unit,
        )

    def init_openai(self) -> HostingProvider:
        quota_unit = QuotaUnit.TIMES
        if os.environ.get("HOSTED_OPENAI_ENABLED") and os.environ.get("HOSTED_OPENAI_ENABLED").lower() == 'true':
            credentials = {
                "openai_api_key": os.environ.get("HOSTED_OPENAI_API_KEY"),
            }

            if os.environ.get("HOSTED_OPENAI_API_BASE"):
                credentials["openai_api_base"] = os.environ.get("HOSTED_OPENAI_API_BASE")

            if os.environ.get("HOSTED_OPENAI_API_ORGANIZATION"):
                credentials["openai_organization"] = os.environ.get("HOSTED_OPENAI_API_ORGANIZATION")

            quotas = []
            hosted_quota_limit = int(os.environ.get("HOSTED_OPENAI_QUOTA_LIMIT", "200"))
            if hosted_quota_limit != -1 or hosted_quota_limit > 0:
                trial_quota = TrialHostingQuota(
                    quota_limit=hosted_quota_limit,
                    restrict_models=[
                        RestrictModel(model="gpt-3.5-turbo", model_type=ModelType.LLM),
                        RestrictModel(model="gpt-3.5-turbo-1106", model_type=ModelType.LLM),
                        RestrictModel(model="gpt-3.5-turbo-instruct", model_type=ModelType.LLM),
                        RestrictModel(model="gpt-3.5-turbo-16k", model_type=ModelType.LLM),
                        RestrictModel(model="text-davinci-003", model_type=ModelType.LLM),
                    ]
                )
                quotas.append(trial_quota)

            if os.environ.get("HOSTED_OPENAI_PAID_ENABLED") and os.environ.get(
                    "HOSTED_OPENAI_PAID_ENABLED").lower() == 'true':
                paid_quota = PaidHostingQuota(
                    stripe_price_id=os.environ.get("HOSTED_OPENAI_PAID_STRIPE_PRICE_ID"),
                    increase_quota=int(os.environ.get("HOSTED_OPENAI_PAID_INCREASE_QUOTA", "1")),
                    min_quantity=int(os.environ.get("HOSTED_OPENAI_PAID_MIN_QUANTITY", "1")),
                    max_quantity=int(os.environ.get("HOSTED_OPENAI_PAID_MAX_QUANTITY", "1"))
                )
                quotas.append(paid_quota)

            return HostingProvider(
                enabled=True,
                credentials=credentials,
                quota_unit=quota_unit,
                quotas=quotas
            )

        return HostingProvider(
            enabled=False,
            quota_unit=quota_unit,
        )

    def init_anthropic(self) -> HostingProvider:
        quota_unit = QuotaUnit.TOKENS
        if os.environ.get("HOSTED_ANTHROPIC_ENABLED") and os.environ.get("HOSTED_ANTHROPIC_ENABLED").lower() == 'true':
            credentials = {
                "anthropic_api_key": os.environ.get("HOSTED_ANTHROPIC_API_KEY"),
            }

            if os.environ.get("HOSTED_ANTHROPIC_API_BASE"):
                credentials["anthropic_api_url"] = os.environ.get("HOSTED_ANTHROPIC_API_BASE")

            quotas = []
            hosted_quota_limit = int(os.environ.get("HOSTED_ANTHROPIC_QUOTA_LIMIT", "0"))
            if hosted_quota_limit != -1 or hosted_quota_limit > 0:
                trial_quota = TrialHostingQuota(
                    quota_limit=hosted_quota_limit
                )
                quotas.append(trial_quota)

            if os.environ.get("HOSTED_ANTHROPIC_PAID_ENABLED") and os.environ.get(
                    "HOSTED_ANTHROPIC_PAID_ENABLED").lower() == 'true':
                paid_quota = PaidHostingQuota(
                    stripe_price_id=os.environ.get("HOSTED_ANTHROPIC_PAID_STRIPE_PRICE_ID"),
                    increase_quota=int(os.environ.get("HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA", "1000000")),
                    min_quantity=int(os.environ.get("HOSTED_ANTHROPIC_PAID_MIN_QUANTITY", "20")),
                    max_quantity=int(os.environ.get("HOSTED_ANTHROPIC_PAID_MAX_QUANTITY", "100"))
                )
                quotas.append(paid_quota)

            return HostingProvider(
                enabled=True,
                credentials=credentials,
                quota_unit=quota_unit,
                quotas=quotas
            )

        return HostingProvider(
            enabled=False,
            quota_unit=quota_unit,
        )

    def init_minimax(self) -> HostingProvider:
        quota_unit = QuotaUnit.TOKENS
        if os.environ.get("HOSTED_MINIMAX_ENABLED") and os.environ.get("HOSTED_MINIMAX_ENABLED").lower() == 'true':
            quotas = [FreeHostingQuota()]

            return HostingProvider(
                enabled=True,
                credentials=None,  # use credentials from the provider
                quota_unit=quota_unit,
                quotas=quotas
            )

        return HostingProvider(
            enabled=False,
            quota_unit=quota_unit,
        )

    def init_spark(self) -> HostingProvider:
        quota_unit = QuotaUnit.TOKENS
        if os.environ.get("HOSTED_SPARK_ENABLED") and os.environ.get("HOSTED_SPARK_ENABLED").lower() == 'true':
            quotas = [FreeHostingQuota()]

            return HostingProvider(
                enabled=True,
                credentials=None,  # use credentials from the provider
                quota_unit=quota_unit,
                quotas=quotas
            )

        return HostingProvider(
            enabled=False,
            quota_unit=quota_unit,
        )

    def init_zhipuai(self) -> HostingProvider:
        quota_unit = QuotaUnit.TOKENS
        if os.environ.get("HOSTED_ZHIPUAI_ENABLED") and os.environ.get("HOSTED_ZHIPUAI_ENABLED").lower() == 'true':
            quotas = [FreeHostingQuota()]

            return HostingProvider(
                enabled=True,
                credentials=None,  # use credentials from the provider
                quota_unit=quota_unit,
                quotas=quotas
            )

        return HostingProvider(
            enabled=False,
            quota_unit=quota_unit,
        )

    def init_moderation_config(self) -> HostedModerationConfig:
        if os.environ.get("HOSTED_MODERATION_ENABLED") and os.environ.get("HOSTED_MODERATION_ENABLED").lower() == 'true' \
                and os.environ.get("HOSTED_MODERATION_PROVIDERS"):
            return HostedModerationConfig(
                enabled=True,
                providers=os.environ.get("HOSTED_MODERATION_PROVIDERS").split(',')
            )

        return HostedModerationConfig(
            enabled=False
        )