feat: add model setting

Author: wxg0103

apps/common/cache/file_cache.py

@@ -0,0 +1,86 @@
+# coding=utf-8
+"""
+    @project: qabot
+    @Author：虎
+    @file： file_cache.py
+    @date：2023/9/11 15:58
+    @desc: 文件缓存
+"""
+import datetime
+import math
+import os
+import time
+
+from diskcache import Cache
+from django.core.cache.backends.base import BaseCache
+
+
+class FileCache(BaseCache):
+    def __init__(self, dir, params):
+        super().__init__(params)
+        self._dir = os.path.abspath(dir)
+        self._createdir()
+        self.cache = Cache(self._dir)
+
+    def _createdir(self):
+        old_umask = os.umask(0o077)
+        try:
+            os.makedirs(self._dir, 0o700, exist_ok=True)
+        finally:
+            os.umask(old_umask)
+
+    def add(self, key, value, timeout=None, version=None):
+        expire = timeout if isinstance(timeout, int) or isinstance(timeout,
+                                                                   float) or timeout is None else timeout.total_seconds()
+        return self.cache.add(self.get_key(key, version), value=value, expire=expire)
+
+    def set(self, key, value, timeout=None, version=None):
+        expire = timeout if isinstance(timeout, int) or isinstance(timeout,
+                                                                   float) or timeout is None else timeout.total_seconds()
+        return self.cache.set(self.get_key(key, version), value=value, expire=expire)
+
+    def get(self, key, default=None, version=None):
+        return self.cache.get(self.get_key(key, version), default=default)
+
+    @staticmethod
+    def get_key(key, version):
+        if version is None:
+            return f"default:{key}"
+        return f"{version}:{key}"
+
+    def delete(self, key, version=None):
+        return self.cache.delete(self.get_key(key, version))
+
+    def touch(self, key, timeout=None, version=None):
+        expire = timeout if isinstance(timeout, int) or isinstance(timeout,
+                                                                   float) else timeout.total_seconds()
+
+        return self.cache.touch(self.get_key(key, version), expire=expire)
+
+    def ttl(self, key, version=None):
+        """
+        获取key的剩余时间
+        :param key: key
+        :return:  剩余时间
+        @param version:
+        """
+        value, expire_time = self.cache.get(self.get_key(key, version), expire_time=True)
+        if value is None:
+            return None
+        return datetime.timedelta(seconds=math.ceil(expire_time - time.time()))
+
+    def clear_by_application_id(self, application_id):
+        delete_keys = []
+        for key in self.cache.iterkeys():
+            value = self.cache.get(key)
+            if (hasattr(value,
+                        'application') and value.application is not None and value.application.id is not None and
+                    str(
+                        value.application.id) == application_id):
+                delete_keys.append(key)
+        for key in delete_keys:
+            self.cache.delete(key)
+
+    def clear_timeout_data(self):
+        for key in self.cache.iterkeys():
+            self.get(key)

apps/common/cache/mem_cache.py

@@ -0,0 +1,47 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： mem_cache.py
+    @date：2024/3/6 11:20
+    @desc:
+"""
+from django.core.cache.backends.base import DEFAULT_TIMEOUT
+from django.core.cache.backends.locmem import LocMemCache
+
+
+class MemCache(LocMemCache):
+    def __init__(self, name, params):
+        super().__init__(name, params)
+
+    def set(self, key, value, timeout=DEFAULT_TIMEOUT, version=None):
+        key = self.make_and_validate_key(key, version=version)
+        pickled = value
+        with self._lock:
+            self._set(key, pickled, timeout)
+
+    def get(self, key, default=None, version=None):
+        key = self.make_and_validate_key(key, version=version)
+        with self._lock:
+            if self._has_expired(key):
+                self._delete(key)
+                return default
+            pickled = self._cache[key]
+            self._cache.move_to_end(key, last=False)
+        return pickled
+
+    def clear_by_application_id(self, application_id):
+        delete_keys = []
+        for key in self._cache.keys():
+            value = self._cache.get(key)
+            if (hasattr(value,
+                        'application') and value.application is not None and value.application.id is not None and
+                    str(
+                        value.application.id) == application_id):
+                delete_keys.append(key)
+        for key in delete_keys:
+            self._delete(key)
+
+    def clear_timeout_data(self):
+        for key in self._cache.keys():
+            self.get(key)

apps/common/config/embedding_config.py

@@ -0,0 +1,66 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： embedding_config.py
+    @date：2023/10/23 16:03
+    @desc:
+"""
+import threading
+import time
+
+from common.cache.mem_cache import MemCache
+
+lock = threading.Lock()
+
+
+class ModelManage:
+    cache = MemCache('model', {})
+    up_clear_time = time.time()
+
+    @staticmethod
+    def get_model(_id, get_model):
+        # 获取锁
+        lock.acquire()
+        try:
+            model_instance = ModelManage.cache.get(_id)
+            if model_instance is None or not model_instance.is_cache_model():
+                model_instance = get_model(_id)
+                ModelManage.cache.set(_id, model_instance, timeout=60 * 30)
+                return model_instance
+            # 续期
+            ModelManage.cache.touch(_id, timeout=60 * 30)
+            ModelManage.clear_timeout_cache()
+            return model_instance
+        finally:
+            # 释放锁
+            lock.release()
+
+    @staticmethod
+    def clear_timeout_cache():
+        if time.time() - ModelManage.up_clear_time > 60:
+            ModelManage.cache.clear_timeout_data()
+
+    @staticmethod
+    def delete_key(_id):
+        if ModelManage.cache.has_key(_id):
+            ModelManage.cache.delete(_id)
+
+
+class VectorStore:
+    from embedding.vector.pg_vector import PGVector
+    from embedding.vector.base_vector import BaseVectorStore
+    instance_map = {
+        'pg_vector': PGVector,
+    }
+    instance = None
+
+    @staticmethod
+    def get_embedding_vector() -> BaseVectorStore:
+        from embedding.vector.pg_vector import PGVector
+        if VectorStore.instance is None:
+            from maxkb.const import CONFIG
+            vector_store_class = VectorStore.instance_map.get(CONFIG.get("VECTOR_STORE_NAME"),
+                                                              PGVector)
+            VectorStore.instance = vector_store_class()
+        return VectorStore.instance

apps/common/config/tokenizer_manage_config.py

@@ -0,0 +1,24 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： tokenizer_manage_config.py
+    @date：2024/4/28 10:17
+    @desc:
+"""
+
+
+class TokenizerManage:
+    tokenizer = None
+
+    @staticmethod
+    def get_tokenizer():
+        from transformers import BertTokenizer
+        if TokenizerManage.tokenizer is None:
+            TokenizerManage.tokenizer = BertTokenizer.from_pretrained(
+                'bert-base-cased',
+                cache_dir="/opt/maxkb/model/tokenizer",
+                local_files_only=True,
+                resume_download=False,
+                force_download=False)
+        return TokenizerManage.tokenizer

apps/common/constants/permission_constants.py

@@ -111,6 +111,13 @@ class PermissionConstants(Enum):
     USER_DELETE = Permission(group=Group.USER, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN])
     TOOL_CREATE = Permission(group=Group.USER, operate=Operate.CREATE, role_list=[RoleConstants.ADMIN,
                                                                                   RoleConstants.USER])
+    MODEL_CREATE = Permission(group=Group.USER, operate=Operate.CREATE, role_list=[RoleConstants.ADMIN,
+                                                                                   RoleConstants.USER])
+    MODEL_READ = Permission(group=Group.USER, operate=Operate.READ, role_list=[RoleConstants.ADMIN,
+                                                                               RoleConstants.USER])
+    MODEL_EDIT = Permission(group=Group.USER, operate=Operate.EDIT, role_list=[RoleConstants.ADMIN, RoleConstants.USER])
+    MODEL_DELETE = Permission(group=Group.USER, operate=Operate.DELETE,
+                              role_list=[RoleConstants.ADMIN, RoleConstants.USER])
 
     def get_workspace_application_permission(self):
         return lambda r, kwargs: Permission(group=self.value.group, operate=self.value.operate,

apps/common/forms/__init__.py

@@ -0,0 +1,24 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： __init__.py
+    @date：2023/10/31 17:56
+    @desc:
+"""
+from .array_object_card import *
+from .base_field import *
+from .base_form import *
+from .multi_select import *
+from .object_card import *
+from .password_input import *
+from .radio_field import *
+from .single_select_field import *
+from .tab_card import *
+from .table_radio import *
+from .text_input_field import *
+from .radio_button_field import *
+from .table_checkbox import *
+from .radio_card_field import *
+from .label import *
+from .slider_field import *

apps/common/forms/array_object_card.py

@@ -0,0 +1,33 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： array_object_card.py
+    @date：2023/10/31 18:03
+    @desc:
+"""
+from typing import Dict
+
+from common.forms.base_field import BaseExecField, TriggerType
+
+
+class ArrayCard(BaseExecField):
+    """
+    收集List[Object]
+    """
+
+    def __init__(self,
+                 label: str,
+                 text_field: str,
+                 value_field: str,
+                 provider: str,
+                 method: str,
+                 required: bool = False,
+                 default_value: object = None,
+                 relation_show_field_dict: Dict = None,
+                 relation_trigger_field_dict: Dict = None,
+                 trigger_type: TriggerType = TriggerType.OPTION_LIST,
+                 attrs: Dict[str, object] = None,
+                 props_info: Dict[str, object] = None):
+        super().__init__("ArrayObjectCard", label, text_field, value_field, provider, method, required, default_value,
+                         relation_show_field_dict, relation_trigger_field_dict, trigger_type, attrs, props_info)