|
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_config
|
|
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()
|
|
|
