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/core/helper/encrypter.py

@@ -18,3 +18,17 @@ def encrypt_token(tenant_id: str, token: str):
 
 def decrypt_token(tenant_id: str, token: str):
     return rsa.decrypt(base64.b64decode(token), tenant_id)
+
+
+def batch_decrypt_token(tenant_id: str, tokens: list[str]):
+    rsa_key, cipher_rsa = rsa.get_decrypt_decoding(tenant_id)
+
+    return [rsa.decrypt_token_with_decoding(base64.b64decode(token), rsa_key, cipher_rsa) for token in tokens]
+
+
+def get_decrypt_decoding(tenant_id: str):
+    return rsa.get_decrypt_decoding(tenant_id)
+
+
+def decrypt_token_with_decoding(token: str, rsa_key, cipher_rsa):
+    return rsa.decrypt_token_with_decoding(base64.b64decode(token), rsa_key, cipher_rsa)

api/core/helper/lru_cache.py

@@ -0,0 +1,22 @@
+from collections import OrderedDict
+from typing import Any
+
+
+class LRUCache:
+    def __init__(self, capacity: int):
+        self.cache = OrderedDict()
+        self.capacity = capacity
+
+    def get(self, key: Any) -> Any:
+        if key not in self.cache:
+            return None
+        else:
+            self.cache.move_to_end(key)  # move the key to the end of the OrderedDict
+            return self.cache[key]
+
+    def put(self, key: Any, value: Any) -> None:
+        if key in self.cache:
+            self.cache.move_to_end(key)
+        self.cache[key] = value
+        if len(self.cache) > self.capacity:
+            self.cache.popitem(last=False)  # pop the first item

api/core/helper/moderation.py

@@ -1,18 +1,27 @@
 import logging
 import random
 
-import openai
-
-from core.model_providers.error import LLMBadRequestError
-from core.model_providers.providers.base import BaseModelProvider
-from core.model_providers.providers.hosted import hosted_config, hosted_model_providers
+from core.entities.application_entities import ModelConfigEntity
+from core.model_runtime.errors.invoke import InvokeBadRequestError
+from core.model_runtime.model_providers.openai.moderation.moderation import OpenAIModerationModel
+from extensions.ext_hosting_provider import hosting_configuration
 from models.provider import ProviderType
 
+logger = logging.getLogger(__name__)
+
+
+def check_moderation(model_config: ModelConfigEntity, text: str) -> bool:
+    moderation_config = hosting_configuration.moderation_config
+    if (moderation_config and moderation_config.enabled is True
+            and 'openai' in hosting_configuration.provider_map
+            and hosting_configuration.provider_map['openai'].enabled is True
+    ):
+        using_provider_type = model_config.provider_model_bundle.configuration.using_provider_type
+        provider_name = model_config.provider
+        if using_provider_type == ProviderType.SYSTEM \
+                and provider_name in moderation_config.providers:
+            hosting_openai_config = hosting_configuration.provider_map['openai']
 
-def check_moderation(model_provider: BaseModelProvider, text: str) -> bool:
-    if hosted_config.moderation.enabled is True and hosted_model_providers.openai:
-        if model_provider.provider.provider_type == ProviderType.SYSTEM.value \
-                and model_provider.provider_name in hosted_config.moderation.providers:
             # 2000 text per chunk
             length = 2000
             text_chunks = [text[i:i + length] for i in range(0, len(text), length)]
@@ -23,14 +32,17 @@ def check_moderation(model_provider: BaseModelProvider, text: str) -> bool:
             text_chunk = random.choice(text_chunks)
 
             try:
-                moderation_result = openai.Moderation.create(input=text_chunk,
-                                                             api_key=hosted_model_providers.openai.api_key)
+                model_type_instance = OpenAIModerationModel()
+                moderation_result = model_type_instance.invoke(
+                    model='text-moderation-stable',
+                    credentials=hosting_openai_config.credentials,
+                    text=text_chunk
+                )
+
+                if moderation_result is True:
+                    return True
             except Exception as ex:
-                logging.exception(ex)
-                raise LLMBadRequestError('Rate limit exceeded, please try again later.')
+                logger.exception(ex)
+                raise InvokeBadRequestError('Rate limit exceeded, please try again later.')
 
-            for result in moderation_result.results:
-                if result['flagged'] is True:
-                    return False
-
-    return True
+    return False