Browse Source

chore(quota): Update deduct quota (#14337)

Signed-off-by: -LAN- <laipz8200@outlook.com>
tags/1.0.0
-LAN- 8 months ago
parent
commit
7a71498a3e
No account linked to committer's email address

+ 10
- 2
api/core/workflow/nodes/llm/node.py View File

import json import json
import logging import logging
from collections.abc import Generator, Mapping, Sequence from collections.abc import Generator, Mapping, Sequence
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any, Optional, cast from typing import TYPE_CHECKING, Any, Optional, cast


from configs import dify_config from configs import dify_config
from core.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey, ModelType from core.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey, ModelType
from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
from core.model_runtime.utils.encoders import jsonable_encoder from core.model_runtime.utils.encoders import jsonable_encoder
from core.plugin.entities.plugin import ModelProviderID
from core.prompt.entities.advanced_prompt_entities import CompletionModelPromptTemplate, MemoryConfig from core.prompt.entities.advanced_prompt_entities import CompletionModelPromptTemplate, MemoryConfig
from core.prompt.utils.prompt_message_util import PromptMessageUtil from core.prompt.utils.prompt_message_util import PromptMessageUtil
from core.variables import ( from core.variables import (
if used_quota is not None and system_configuration.current_quota_type is not None: if used_quota is not None and system_configuration.current_quota_type is not None:
db.session.query(Provider).filter( db.session.query(Provider).filter(
Provider.tenant_id == tenant_id, Provider.tenant_id == tenant_id,
Provider.provider_name == model_instance.provider,
# TODO: Use provider name with prefix after the data migration.
Provider.provider_name == ModelProviderID(model_instance.provider).provider_name,
Provider.provider_type == ProviderType.SYSTEM.value, Provider.provider_type == ProviderType.SYSTEM.value,
Provider.quota_type == system_configuration.current_quota_type.value, Provider.quota_type == system_configuration.current_quota_type.value,
Provider.quota_limit > Provider.quota_used, Provider.quota_limit > Provider.quota_used,
).update({"quota_used": Provider.quota_used + used_quota})
).update(
{
"quota_used": Provider.quota_used + used_quota,
"last_used": datetime.now(tz=UTC).replace(tzinfo=None),
}
)
db.session.commit() db.session.commit()


@classmethod @classmethod

+ 11
- 2
api/events/event_handlers/deduct_quota_when_message_created.py View File

from datetime import UTC, datetime

from configs import dify_config from configs import dify_config
from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, ChatAppGenerateEntity from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, ChatAppGenerateEntity
from core.entities.provider_entities import QuotaUnit from core.entities.provider_entities import QuotaUnit
from core.plugin.entities.plugin import ModelProviderID
from events.message_event import message_was_created from events.message_event import message_was_created
from extensions.ext_database import db from extensions.ext_database import db
from models.provider import Provider, ProviderType from models.provider import Provider, ProviderType
if used_quota is not None and system_configuration.current_quota_type is not None: if used_quota is not None and system_configuration.current_quota_type is not None:
db.session.query(Provider).filter( db.session.query(Provider).filter(
Provider.tenant_id == application_generate_entity.app_config.tenant_id, Provider.tenant_id == application_generate_entity.app_config.tenant_id,
Provider.provider_name == model_config.provider,
# TODO: Use provider name with prefix after the data migration.
Provider.provider_name == ModelProviderID(model_config.provider).provider_name,
Provider.provider_type == ProviderType.SYSTEM.value, Provider.provider_type == ProviderType.SYSTEM.value,
Provider.quota_type == system_configuration.current_quota_type.value, Provider.quota_type == system_configuration.current_quota_type.value,
Provider.quota_limit > Provider.quota_used, Provider.quota_limit > Provider.quota_used,
).update({"quota_used": Provider.quota_used + used_quota})
).update(
{
"quota_used": Provider.quota_used + used_quota,
"last_used": datetime.now(tz=UTC).replace(tzinfo=None),
}
)
db.session.commit() db.session.commit()

Loading…
Cancel
Save