| from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, ChatAppGenerateEntity |
| from core.entities.provider_entities import QuotaUnit |
| from events.message_event import message_was_created |
| from extensions.ext_database import db |
| from models.provider import Provider, ProviderType |
|
|
|
|
| @message_was_created.connect |
| def handle(sender, **kwargs): |
| message = sender |
| application_generate_entity = kwargs.get("application_generate_entity") |
|
|
| if not isinstance(application_generate_entity, ChatAppGenerateEntity | AgentChatAppGenerateEntity): |
| return |
|
|
| model_config = application_generate_entity.model_conf |
| provider_model_bundle = model_config.provider_model_bundle |
| provider_configuration = provider_model_bundle.configuration |
|
|
| if provider_configuration.using_provider_type != ProviderType.SYSTEM: |
| return |
|
|
| system_configuration = provider_configuration.system_configuration |
|
|
| quota_unit = None |
| for quota_configuration in system_configuration.quota_configurations: |
| if quota_configuration.quota_type == system_configuration.current_quota_type: |
| quota_unit = quota_configuration.quota_unit |
|
|
| if quota_configuration.quota_limit == -1: |
| return |
|
|
| break |
|
|
| used_quota = None |
| if quota_unit: |
| if quota_unit == QuotaUnit.TOKENS: |
| used_quota = message.message_tokens + message.answer_tokens |
| elif quota_unit == QuotaUnit.CREDITS: |
| used_quota = 1 |
|
|
| if "gpt-4" in model_config.model: |
| used_quota = 20 |
| else: |
| used_quota = 1 |
|
|
| if used_quota is not None: |
| db.session.query(Provider).filter( |
| Provider.tenant_id == application_generate_entity.app_config.tenant_id, |
| Provider.provider_name == model_config.provider, |
| Provider.provider_type == ProviderType.SYSTEM.value, |
| Provider.quota_type == system_configuration.current_quota_type.value, |
| Provider.quota_limit > Provider.quota_used, |
| ).update({"quota_used": Provider.quota_used + used_quota}) |
| db.session.commit() |
|
|
