Model Runtime (#1858)

Co-authored-by: StyleZhang <jasonapring2015@outlook.com>
Co-authored-by: Garfield Dai <dai.hai@foxmail.com>
Co-authored-by: chenhe <guchenhe@gmail.com>
Co-authored-by: jyong <jyong@dify.ai>
Co-authored-by: Joel <iamjoel007@gmail.com>
Co-authored-by: Yeuoly <admin@srmxy.cn>

api/events/event_handlers/__init__.py

@@ -5,3 +5,5 @@ from .clean_when_dataset_deleted import handle
 from .update_app_dataset_join_when_app_model_config_updated import handle
 from .generate_conversation_name_when_first_message_created import handle
 from .create_document_index import handle
+from .deduct_quota_when_messaeg_created import handle
+from .update_provider_last_used_at_when_messaeg_created import handle

api/events/event_handlers/deduct_quota_when_messaeg_created.py

@@ -0,0 +1,43 @@
+from core.entities.application_entities import ApplicationGenerateEntity
+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 ProviderType, Provider
+
+
+@message_was_created.connect
+def handle(sender, **kwargs):
+    message = sender
+    application_generate_entity: ApplicationGenerateEntity = kwargs.get('application_generate_entity')
+
+    model_config = application_generate_entity.app_orchestration_config_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
+            break
+
+    used_quota = None
+    if quota_unit:
+        if quota_unit == QuotaUnit.TOKENS.value:
+            used_quota = message.message_tokens + message.prompt_tokens
+        else:
+            used_quota = 1
+
+    if used_quota is not None:
+        db.session.query(Provider).filter(
+            Provider.tenant_id == application_generate_entity.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()

api/events/event_handlers/generate_conversation_name_when_first_message_created.py

@@ -8,9 +8,13 @@ def handle(sender, **kwargs):
     message = sender
     conversation = kwargs.get('conversation')
     is_first_message = kwargs.get('is_first_message')
-    auto_generate_name = kwargs.get('auto_generate_name', True)
+    extras = kwargs.get('extras', {})
 
-    if auto_generate_name and is_first_message:
+    auto_generate_conversation_name = True
+    if extras:
+        auto_generate_conversation_name = extras.get('auto_generate_conversation_name', True)
+
+    if auto_generate_conversation_name and is_first_message:
         if conversation.mode == 'chat':
             app_model = conversation.app
             if not app_model:

api/events/event_handlers/update_provider_last_used_at_when_messaeg_created.py

@@ -0,0 +1,18 @@
+from datetime import datetime
+
+from core.entities.application_entities import ApplicationGenerateEntity
+from events.message_event import message_was_created
+from extensions.ext_database import db
+from models.provider import Provider
+
+
+@message_was_created.connect
+def handle(sender, **kwargs):
+    message = sender
+    application_generate_entity: ApplicationGenerateEntity = kwargs.get('application_generate_entity')
+
+    db.session.query(Provider).filter(
+        Provider.tenant_id == application_generate_entity.tenant_id,
+        Provider.provider_name == application_generate_entity.app_orchestration_config_entity.model_config.provider
+    ).update({'last_used': datetime.utcnow()})
+    db.session.commit()