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

Signed-off-by: -LAN- <[email protected]>
langgenius · Feb 25, 2025 · 7a71498 · 7a71498
1 parent 76bcdc2
commit 7a71498
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 4 deletions.
diff --git 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
@@ -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()