refactor: email setting

Author: wxg0103

apps/common/constants/permission_constants.py

@@ -217,23 +217,28 @@ class PermissionConstants(Enum):
     KNOWLEDGE_CREATE = Permission(group=Group.KNOWLEDGE, operate=Operate.CREATE, role_list=[RoleConstants.ADMIN,
                                                                                             RoleConstants.USER])
     KNOWLEDGE_EDIT = Permission(group=Group.KNOWLEDGE, operate=Operate.EDIT, role_list=[RoleConstants.ADMIN,
-                                                                                            RoleConstants.USER])
+                                                                                        RoleConstants.USER])
     KNOWLEDGE_DELETE = Permission(group=Group.KNOWLEDGE, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN,
                                                                                             RoleConstants.USER])
     DOCUMENT_READ = Permission(group=Group.KNOWLEDGE, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN,
-                                                                                            RoleConstants.USER])
+                                                                                         RoleConstants.USER])
     DOCUMENT_CREATE = Permission(group=Group.KNOWLEDGE, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN,
-                                                                                            RoleConstants.USER])
+                                                                                           RoleConstants.USER])
     DOCUMENT_EDIT = Permission(group=Group.KNOWLEDGE, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN,
-                                                                                            RoleConstants.USER])
+                                                                                         RoleConstants.USER])
     DOCUMENT_DELETE = Permission(group=Group.KNOWLEDGE, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN,
-                                                                                            RoleConstants.USER])
+                                                                                           RoleConstants.USER])
 
     WORKSPACE_USER_RESOURCE_PERMISSION_READ = Permission(group=Group.WORKSPACE_USER_RESOURCE_PERMISSION,
                                                          operate=Operate.READ,
                                                          role_list=[RoleConstants.ADMIN,
                                                                     RoleConstants.WORKSPACE_MANAGE])
 
+    EMAIL_SETTING_READ = Permission(group=Group.USER, operate=Operate.READ,
+                                    role_list=[RoleConstants.ADMIN])
+    EMAIL_SETTING_EDIT = Permission(group=Group.USER, operate=Operate.EDIT,
+                                    role_list=[RoleConstants.ADMIN])
+
     def get_workspace_application_permission(self):
         return lambda r, kwargs: Permission(group=self.value.group, operate=self.value.operate,
                                             resource_path=

apps/models_provider/serializers/model_apply_serializers.py

@@ -0,0 +1,76 @@
+# coding=utf-8
+"""
+    @project: MaxKB
+    @Author：虎
+    @file： model_apply_serializers.py
+    @date：2024/8/20 20:39
+    @desc:
+"""
+from django.db import connection
+from django.db.models import QuerySet
+from langchain_core.documents import Document
+from rest_framework import serializers
+
+from common.config.embedding_config import ModelManage
+from django.utils.translation import gettext_lazy as _
+
+from models_provider.models import Model
+from models_provider.tools import get_model
+
+
+def get_embedding_model(model_id):
+    model = QuerySet(Model).filter(id=model_id).first()
+    # 手动关闭数据库连接
+    connection.close()
+    embedding_model = ModelManage.get_model(model_id,
+                                            lambda _id: get_model(model, use_local=True))
+    return embedding_model
+
+
+class EmbedDocuments(serializers.Serializer):
+    texts = serializers.ListField(required=True, child=serializers.CharField(required=True,
+                                                                             label=_('vector text')),
+                                  label=_('vector text list')),
+
+
+class EmbedQuery(serializers.Serializer):
+    text = serializers.CharField(required=True, label=_('vector text'))
+
+
+class CompressDocument(serializers.Serializer):
+    page_content = serializers.CharField(required=True, label=_('text'))
+    metadata = serializers.DictField(required=False, label=_('metadata'))
+
+
+class CompressDocuments(serializers.Serializer):
+    documents = CompressDocument(required=True, many=True)
+    query = serializers.CharField(required=True, label=_('query'))
+
+
+class ModelApplySerializers(serializers.Serializer):
+    model_id = serializers.UUIDField(required=True, label=_('model id'))
+
+    def embed_documents(self, instance, with_valid=True):
+        if with_valid:
+            self.is_valid(raise_exception=True)
+            EmbedDocuments(data=instance).is_valid(raise_exception=True)
+
+        model = get_embedding_model(self.data.get('model_id'))
+        return model.embed_documents(instance.getlist('texts'))
+
+    def embed_query(self, instance, with_valid=True):
+        if with_valid:
+            self.is_valid(raise_exception=True)
+            EmbedQuery(data=instance).is_valid(raise_exception=True)
+
+        model = get_embedding_model(self.data.get('model_id'))
+        return model.embed_query(instance.get('text'))
+
+    def compress_documents(self, instance, with_valid=True):
+        if with_valid:
+            self.is_valid(raise_exception=True)
+            CompressDocuments(data=instance).is_valid(raise_exception=True)
+        model = get_embedding_model(self.data.get('model_id'))
+        return [{'page_content': d.page_content, 'metadata': d.metadata} for d in model.compress_documents(
+            [Document(page_content=document.get('page_content'), metadata=document.get('metadata')) for document in
+             instance.get('documents')], instance.get('query'))]

apps/models_provider/urls.py

@@ -1,3 +1,5 @@
+import os
+
 from django.urls import path
 
 from . import views
@@ -15,3 +17,10 @@
     path('workspace/<str:workspace_id>/model/<str:model_id>/pause_download', views.Model.PauseDownload.as_view()),
     path('workspace/<str:workspace_id>/model/<str:model_id>/meta', views.Model.ModelMeta.as_view()),
 ]
+
+if os.environ.get('SERVER_NAME', 'web') == 'local_model':
+    urlpatterns += [
+        path('workspace/<str:workspace_id>/model/<str:model_id>/embed_documents', views.ModelApply.EmbedDocuments.as_view()),
+        path('workspace/<str:workspace_id>model/<str:model_id>/embed_query', views.ModelApply.EmbedQuery.as_view()),
+        path('workspace/<str:workspace_id>model/<str:model_id>/compress_documents', views.ModelApply.CompressDocuments.as_view()),
+    ]

apps/models_provider/views/__init__.py

@@ -1,4 +1,5 @@
 # coding=utf-8
 
 from .model import *
-from .provide import *
+from .provide import *
+from .model_apply import *

apps/models_provider/views/model_apply.py

@@ -0,0 +1,60 @@
+# coding=utf-8
+"""
+    @project: MaxKB
+    @Author：虎
+    @file： model_apply.py
+    @date：2024/8/20 20:38
+    @desc:
+"""
+from urllib.request import Request
+
+from django.utils.translation import gettext_lazy as _
+from drf_spectacular.utils import extend_schema
+from rest_framework.views import APIView
+
+from common.auth.authentication import has_permissions
+from common.constants.permission_constants import PermissionConstants
+from common.result import result
+from models_provider.api.model import DefaultModelResponse
+from models_provider.serializers.model_apply_serializers import ModelApplySerializers
+
+
+class ModelApply(APIView):
+    class EmbedDocuments(APIView):
+        @extend_schema(methods=['POST'],
+                       summary=_('Vectorization documentation'),
+                       description=_('Vectorization documentation'),
+                       operation_id=_('Vectorization documentation'),
+                       responses=DefaultModelResponse.get_response(),
+                       tags=[_('Model')]
+                       )
+        @has_permissions(PermissionConstants.MODEL_READ.get_workspace_permission())
+        def post(self, request: Request, workspace_id, model_id):
+            return result.success(
+                ModelApplySerializers(data={'model_id': model_id}).embed_documents(request.data))
+
+    class EmbedQuery(APIView):
+        @extend_schema(methods=['POST'],
+                       summary=_('Vectorization documentation'),
+                       description=_('Vectorization documentation'),
+                       operation_id=_('Vectorization documentation'),
+                       responses=DefaultModelResponse.get_response(),
+                       tags=[_('Model')]
+                       )
+        @has_permissions(PermissionConstants.MODEL_READ.get_workspace_permission())
+        def post(self, request: Request, workspace_id, model_id):
+            return result.success(
+                ModelApplySerializers(data={'model_id': model_id}).embed_query(request.data))
+
+    class CompressDocuments(APIView):
+        @extend_schema(methods=['POST'],
+                       summary=_('Reorder documents'),
+                       description=_('Reorder documents'),
+                       operation_id=_('Reorder documents'),
+                       responses=DefaultModelResponse.get_response(),
+                       tags=[_('Model')]
+                       )
+        @has_permissions(PermissionConstants.MODEL_READ.get_workspace_permission())
+        def post(self, request: Request, workspace_id, model_id):
+            return result.success(
+                ModelApplySerializers(data={'model_id': model_id}).compress_documents(request.data))

apps/system_manage/api/email_setting.py

@@ -0,0 +1,31 @@
+# coding=utf-8
+"""
+    @project: MaxKB
+    @Author：虎虎
+    @file： workspace_user_resource_permission.py
+    @date：2025/4/28 18:13
+    @desc:
+"""
+from drf_spectacular.types import OpenApiTypes
+from drf_spectacular.utils import OpenApiParameter
+
+from common.mixins.api_mixin import APIMixin
+from common.result import ResultSerializer
+from system_manage.serializers.email_setting import EmailSettingSerializer
+from system_manage.serializers.user_resource_permission import UserResourcePermissionResponse, \
+    UpdateUserResourcePermissionRequest
+
+
+class EmailResponse(ResultSerializer):
+    def get_data(self):
+        return EmailSettingSerializer.Create()
+
+
+class EmailSettingAPI(APIMixin):
+    @staticmethod
+    def get_request():
+        return EmailSettingSerializer.Create()
+
+    @staticmethod
+    def get_response():
+        return EmailResponse

apps/system_manage/serializers/email_setting.py

@@ -0,0 +1,68 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： system_setting.py
+    @date：2024/3/19 16:29
+    @desc:
+"""
+from django.core.mail.backends.smtp import EmailBackend
+from django.db.models import QuerySet
+from rest_framework import serializers
+
+from common.exception.app_exception import AppApiException
+from django.utils.translation import gettext_lazy as _
+
+from system_manage.models import SystemSetting, SettingType
+
+
+class EmailSettingSerializer(serializers.Serializer):
+    @staticmethod
+    def one():
+        system_setting = QuerySet(SystemSetting).filter(type=SettingType.EMAIL.value).first()
+        if system_setting is None:
+            return {}
+        return system_setting.meta
+
+    class Create(serializers.Serializer):
+        email_host = serializers.CharField(required=True, label=_('SMTP host'))
+        email_port = serializers.IntegerField(required=True, label=_('SMTP port'))
+        email_host_user = serializers.CharField(required=True, label=_('Sender\'s email'))
+        email_host_password = serializers.CharField(required=True, label=_('Password'))
+        email_use_tls = serializers.BooleanField(required=True, label=_('Whether to enable TLS'))
+        email_use_ssl = serializers.BooleanField(required=True, label=_('Whether to enable SSL'))
+        from_email = serializers.EmailField(required=True, label=_('Sender\'s email'))
+
+        def is_valid(self, *, raise_exception=False):
+            super().is_valid(raise_exception=True)
+            try:
+                EmailBackend(self.data.get("email_host"),
+                             self.data.get("email_port"),
+                             self.data.get("email_host_user"),
+                             self.data.get("email_host_password"),
+                             self.data.get("email_use_tls"),
+                             False,
+                             self.data.get("email_use_ssl")
+                             ).open()
+            except Exception as e:
+                print(e)
+                raise AppApiException(1004, _('Email verification failed'))
+
+        def update_or_save(self):
+            self.is_valid(raise_exception=True)
+            system_setting = QuerySet(SystemSetting).filter(type=SettingType.EMAIL.value).first()
+            if system_setting is None:
+                system_setting = SystemSetting(type=SettingType.EMAIL.value)
+            system_setting.meta = self.to_email_meta()
+            system_setting.save()
+            return system_setting.meta
+
+        def to_email_meta(self):
+            return {'email_host': self.data.get('email_host'),
+                    'email_port': self.data.get('email_port'),
+                    'email_host_user': self.data.get('email_host_user'),
+                    'email_host_password': self.data.get('email_host_password'),
+                    'email_use_tls': self.data.get('email_use_tls'),
+                    'email_use_ssl': self.data.get('email_use_ssl'),
+                    'from_email': self.data.get('from_email')
+                    }

apps/system_manage/urls.py

@@ -4,5 +4,6 @@
 
 app_name = "system_manage"
 urlpatterns = [
-    path('workspace/<str:workspace_id>/user_resource_permission', views.WorkSpaceUserResourcePermissionView.as_view())
+    path('workspace/<str:workspace_id>/user_resource_permission', views.WorkSpaceUserResourcePermissionView.as_view()),
+    path('email_setting', views.SystemSetting.Email.as_view()),
 ]

apps/system_manage/views/__init__.py

@@ -7,3 +7,4 @@
     @desc:
 """
 from .user_resource_permission import *
+from .email_setting import *

apps/system_manage/views/email_setting.py

@@ -0,0 +1,65 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： system_setting.py
+    @date：2024/3/19 16:01
+    @desc:
+"""
+from drf_spectacular.utils import extend_schema
+from rest_framework.request import Request
+from rest_framework.views import APIView
+
+from common.auth import TokenAuth
+from common.auth.authentication import has_permissions
+from common.constants.permission_constants import PermissionConstants
+
+from django.utils.translation import gettext_lazy as _
+
+from common.result import result
+from models_provider.api.model import DefaultModelResponse
+from system_manage.api.email_setting import EmailSettingAPI
+from system_manage.serializers.email_setting import EmailSettingSerializer
+
+
+class SystemSetting(APIView):
+    class Email(APIView):
+        authentication_classes = [TokenAuth]
+
+        @extend_schema(methods=['PUT'],
+                       summary=_('Create or update email settings'),
+                       description=_('Create or update email settings'),
+                       operation_id=_('Create or update email settings'),
+                       request=EmailSettingAPI.get_request(),
+                       responses=EmailSettingAPI.get_response(),
+                       tags=[_('Email settings')])
+        @has_permissions(PermissionConstants.EMAIL_SETTING_EDIT)
+        def put(self, request: Request):
+            return result.success(
+                EmailSettingSerializer.Create(
+                    data=request.data).update_or_save())
+
+        @extend_schema(
+            methods=['POST'],
+            summary=_('Test email settings'),
+            operation_id=_('Test email settings'),
+            request=EmailSettingAPI.get_request(),
+            responses=DefaultModelResponse.get_response(),
+            tags=[_('Email settings')]
+        )
+        @has_permissions(PermissionConstants.EMAIL_SETTING_EDIT)
+        def post(self, request: Request):
+            return result.success(
+                EmailSettingSerializer.Create(
+                    data=request.data).is_valid())
+
+        @extend_schema(methods=['GET'],
+                       summary=_('Get email settings'),
+                       description=_('Get email settings'),
+                       operation_id=_('Get email settings'),
+                       responses=DefaultModelResponse.get_response(),
+                       tags=[_('Email settings')])
+        @has_permissions(PermissionConstants.EMAIL_SETTING_READ)
+        def get(self, request: Request):
+            return result.success(
+                EmailSettingSerializer.one())