diff --git a/api/core/workflow/nodes/llm/node.py b/api/core/workflow/nodes/llm/node.py index 349d1a7fca29a6..aff3d2b7c68464 100644 --- a/api/core/workflow/nodes/llm/node.py +++ b/api/core/workflow/nodes/llm/node.py @@ -1,6 +1,7 @@ import json import logging from collections.abc import Generator, Mapping, Sequence +from datetime import UTC, datetime from typing import TYPE_CHECKING, Any, Optional, cast from configs import dify_config @@ -29,6 +30,7 @@ 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.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.utils.prompt_message_util import PromptMessageUtil from core.variables import ( @@ -758,11 +760,17 @@ def deduct_llm_quota(cls, tenant_id: str, model_instance: ModelInstance, usage: if used_quota is not None and system_configuration.current_quota_type is not None: db.session.query(Provider).filter( 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.quota_type == system_configuration.current_quota_type.value, 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() @classmethod diff --git a/api/events/event_handlers/deduct_quota_when_message_created.py b/api/events/event_handlers/deduct_quota_when_message_created.py index b1b2d2a0864197..b8e7019446cff9 100644 --- a/api/events/event_handlers/deduct_quota_when_message_created.py +++ b/api/events/event_handlers/deduct_quota_when_message_created.py @@ -1,6 +1,9 @@ +from datetime import UTC, datetime + from configs import dify_config from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, ChatAppGenerateEntity from core.entities.provider_entities import QuotaUnit +from core.plugin.entities.plugin import ModelProviderID from events.message_event import message_was_created from extensions.ext_database import db from models.provider import Provider, ProviderType @@ -48,9 +51,15 @@ def handle(sender, **kwargs): if used_quota is not None and system_configuration.current_quota_type is not None: db.session.query(Provider).filter( 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.quota_type == system_configuration.current_quota_type.value, 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()