suhh-开启多租户功能

api/constants/__init__.py

@@ -12,7 +12,7 @@ VIDEO_EXTENSIONS.extend([ext.upper() for ext in VIDEO_EXTENSIONS])
 AUDIO_EXTENSIONS = ["mp3", "m4a", "wav", "webm", "amr"]
 AUDIO_EXTENSIONS.extend([ext.upper() for ext in AUDIO_EXTENSIONS])
 
-TEMPLATE_EXTENSIONS=["doc", "docx", "csv", "xlsx", "xls","ppt"]
+TEMPLATE_EXTENSIONS = ["doc", "docx", "csv", "xlsx", "xls", "ppt"]
 
 if dify_config.ETL_TYPE == "Unstructured":
     DOCUMENT_EXTENSIONS = ["txt", "markdown", "md", "mdx", "pdf", "html", "htm", "xlsx", "xls"]

api/controllers/console/datasets/datasets_templates.py

@@ -42,7 +42,6 @@ class DatasetTemplateListApi(Resource):
     @setup_required
     @login_required
     @account_initialization_required
-
     def get(self, dataset_id):
         dataset_id = str(dataset_id)
         page = request.args.get("page", default=1, type=int)
@@ -130,7 +129,7 @@ class DatasetTemplateListApi(Resource):
             raise UnsupportedFileTypeError()
 
         try:
-            templates,batch = TemplateService.save_template_with_dataset_id(upload_file,dataset, current_user)
+            templates, batch = TemplateService.save_template_with_dataset_id(upload_file, dataset, current_user)
         except ProviderTokenNotInitError as ex:
             raise ProviderNotInitializeError(ex.description)
         except QuotaExceededError:
@@ -156,6 +155,7 @@ class DatasetTemplateListApi(Resource):
         TemplateService.delete_templates(dataset, template_ids)
         return {"result": "success"}, 204
 
+
 class DatasetTemplateApi(Resource):
     @setup_required
     @login_required
@@ -167,7 +167,7 @@ class DatasetTemplateApi(Resource):
         template_id = str(template_id)
         template = Template.query.filter_by(id=template_id).first()
 
-        #as_attachment下载作为附件下载
+        # as_attachment下载作为附件下载
         return send_file(template.file_url, as_attachment=True)
 
     @setup_required
@@ -183,5 +183,6 @@ class DatasetTemplateApi(Resource):
         TemplateService.delete_template(template)
         return {"result": "success"}, 204
 
+
 api.add_resource(DatasetTemplateListApi, "/datasets/<uuid:dataset_id>/templates")
 api.add_resource(DatasetTemplateApi, "/datasets/template/<uuid:template_id>")

api/controllers/console/datasets/error.py

@@ -66,11 +66,13 @@ class DocumentIndexingError(BaseHTTPException):
     description = "The document is being processed and cannot be edited."
     code = 400
 
+
 class TemplateIndexingError(BaseHTTPException):
     error_code = "template_indexing"
     description = "The template is being processed and cannot be edited."
     code = 400
 
+
 class InvalidMetadataError(BaseHTTPException):
     error_code = "invalid_metadata"
     description = "The metadata content is incorrect. Please check and verify."

api/controllers/console/files.py

@@ -101,6 +101,7 @@ class FileSupportTypeApi(Resource):
     def get(self):
         return {"allowed_extensions": DOCUMENT_EXTENSIONS}
 
+
 class TemplateFileSupportTypeApi(Resource):
     @setup_required
     @login_required

api/core/rag/index_processor/constant/built_in_field.py

@@ -8,6 +8,7 @@ class BuiltInField(str, Enum):
     last_update_date = "last_update_date"
     source = "source"
 
+
 class BuiltIntemplateField(str, Enum):
     template_name = "template_name"
     uploader = "uploader"
@@ -15,6 +16,7 @@ class BuiltIntemplateField(str, Enum):
     last_update_date = "last_update_date"
     source = "source"
 
+
 class MetadataDataSource(Enum):
     upload_file = "file_upload"
     website_crawl = "website"

api/fields/template_fields.py

@@ -5,8 +5,8 @@ from libs.helper import TimestampField
 
 template_fields = {
     "id": fields.String,
-    "tenant_id":fields.String,
-    "dataset_id":fields.String,
+    "tenant_id": fields.String,
+    "dataset_id": fields.String,
     "position": fields.Integer,
     "data_source_type": fields.String,
     "data_source_info": fields.Raw(attribute="data_source_info_dict"),
@@ -25,7 +25,7 @@ template_fields = {
     "display_status": fields.String,
     "word_count": fields.Integer,
     "hit_count": fields.Integer,
-    "doc_form": fields.String
+    "doc_form": fields.String,
 }
 dataset_and_template_fields = {
     "dataset": fields.Nested(dataset_fields),
@@ -45,4 +45,4 @@ template_status_fields = {
     "stopped_at": TimestampField,
     "completed_segments": fields.Integer,
     "total_segments": fields.Integer,
-}
+}

api/models/dataset.py

@@ -634,6 +634,7 @@ class Document(db.Model):  # type: ignore[name-defined]
             doc_language=data.get("doc_language"),
         )
 
+
 class Template(db.Model):  # type: ignore[name-defined]
     __tablename__ = "template"
     __table_args__ = (
@@ -664,7 +665,7 @@ class Template(db.Model):  # type: ignore[name-defined]
 
     # parsing
     file_id = db.Column(db.Text, nullable=True)
-    file_url=db.Column(db.Text,nullable=True)
+    file_url = db.Column(db.Text, nullable=True)
     word_count = db.Column(db.Integer, nullable=True)
     parsing_completed_at = db.Column(db.DateTime, nullable=True)
 
@@ -973,6 +974,8 @@ class Template(db.Model):  # type: ignore[name-defined]
             doc_form=data.get("doc_form"),
             doc_language=data.get("doc_language"),
         )
+
+
 class DocumentSegment(db.Model):  # type: ignore[name-defined]
     __tablename__ = "document_segments"
     __table_args__ = (

api/services/account_service.py

@@ -113,7 +113,11 @@ class AccountService:
             account.current_tenant_id = current_tenant.tenant_id
         else:
             available_ta = (
-                TenantAccountJoin.query.filter_by(account_id=account.id).order_by(TenantAccountJoin.id.asc()).first()
+                # TenantAccountJoin.query.filter_by(account_id=account.id).order_by(TenantAccountJoin.id.asc()).first()
+                TenantAccountJoin.query.filter_by(account_id=account.id)
+                .filter_by(role="owner")
+                .order_by(TenantAccountJoin.id.asc())
+                .first()
             )
             if not available_ta:
                 return None

api/services/dataset_service.py

@@ -521,6 +521,8 @@ class DatasetService:
             "document_ids": [],
             "count": 0,
         }
+
+
 class TemplateService:
     DEFAULT_RULES: dict[str, Any] = {
         "mode": "custom",
@@ -536,9 +538,9 @@ class TemplateService:
         },
     }
 
-    #批量删除
+    # 批量删除
     @staticmethod
-    def delete_templates(dataset: Dataset,template_ids: list[str]):
+    def delete_templates(dataset: Dataset, template_ids: list[str]):
         templates = db.session.query(Template).filter(Template.id.in_(template_ids)).all()
         file_ids = [
             template.data_source_info_dict["upload_file_id"]
@@ -558,14 +560,10 @@ class TemplateService:
                 data_source_info = template.data_source_info_dict
                 if data_source_info and "upload_file_id" in data_source_info:
                     file_id = data_source_info["upload_file_id"]
-        template_was_deleted.send(
-
-            template.id, file_id=file_id
-        )
+        template_was_deleted.send(template.id, file_id=file_id)
         db.session.delete(template)
         db.session.commit()
 
-
     @staticmethod
     def save_template_with_dataset_id(
         upload_file: UploadFile,
@@ -573,41 +571,40 @@ class TemplateService:
         account: Account | Any,
         dataset_process_rule: Optional[DatasetProcessRule] = None,
         created_from: str = "web",
-        ):
+    ):
         batch = time.strftime("%Y%m%d%H%M%S") + str(random.randint(100000, 999999))
         # save process rule
         lock_name = "add_template_lock_dataset_id_{}".format(dataset.id)
         with redis_client.lock(lock_name, timeout=600):
-          position = TemplateService.get_templates_position(dataset.id)
+            position = TemplateService.get_templates_position(dataset.id)
         if not upload_file:
-         raise FileNotExistsError()
+            raise FileNotExistsError()
 
         file_name = upload_file.name
-        print("文件名称"+file_name)
-        file_id=upload_file.id
-        print("文件id"+file_id)
-        data_source_info = {
-        "upload_file_id": upload_file.id
-         }
+        print("文件名称" + file_name)
+        file_id = upload_file.id
+        print("文件id" + file_id)
+        data_source_info = {"upload_file_id": upload_file.id}
         template = TemplateService.build_template(
-                            dataset,
-                            #dataset_process_rule.id,  # type: ignore
-                            #data_source_type,
-                            data_source_info,
-                            created_from,
-                            position,
-                            account,
-                            file_id,
-                            file_name,
-                            batch,
-                            )
+            dataset,
+            # dataset_process_rule.id,  # type: ignore
+            # data_source_type,
+            data_source_info,
+            created_from,
+            position,
+            account,
+            file_id,
+            file_name,
+            batch,
+        )
         db.session.add(template)
         db.session.flush()
         db.session.commit()
         templates = []
         templates.append(template)
         position += 1
-        return templates,batch
+        return templates, batch
+
     @staticmethod
     def get_templates_position(dataset_id):
         template = Template.query.filter_by(dataset_id=dataset_id).order_by(Template.position.desc()).first()
@@ -618,16 +615,16 @@ class TemplateService:
 
     @staticmethod
     def build_template(
-            dataset: Dataset,
-            #process_rule_id: str,
-            #data_source_type: 'upload_file',
-            data_source_info: dict,
-            created_from: str,
-            position: int,
-            account: Account,
-            file_id: str,
-            file_name: str,
-            batch: str,
+        dataset: Dataset,
+        # process_rule_id: str,
+        # data_source_type: 'upload_file',
+        data_source_info: dict,
+        created_from: str,
+        position: int,
+        account: Account,
+        file_id: str,
+        file_name: str,
+        batch: str,
     ):
         template = Template(
             tenant_id=dataset.tenant_id,
@@ -635,24 +632,25 @@ class TemplateService:
             position=position,
             data_source_type="upload_file",
             data_source_info=json.dumps(data_source_info),
-            #dataset_process_rule_id=process_rule_id,
+            # dataset_process_rule_id=process_rule_id,
             batch=batch,
             name=file_name,
             created_from=created_from,
             created_by=account.id,
-            file_id=file_id
+            file_id=file_id,
         )
         return template
 
     @staticmethod
-    def get_templates(template_id)-> Optional[Template]:
+    def get_templates(template_id) -> Optional[Template]:
         if template_id:
             print("模版id" + template_id)
-            template:Optional[Template] = Template.query.filter_by(id=template_id).first()
+            template: Optional[Template] = Template.query.filter_by(id=template_id).first()
             return template
         else:
             return None
 
+
 class DocumentService:
     DEFAULT_RULES: dict[str, Any] = {
         "mode": "custom",
@@ -2462,11 +2460,13 @@ class DatasetPermissionService:
         except Exception as e:
             db.session.rollback()
             raise e
-    @classmethod
+
+    @staticmethod
     def save_template_with_dataset_id(
         dataset: Dataset,
         knowledge_config: KnowledgeConfig,
         account: Account | Any,
         dataset_process_rule: Optional[DatasetProcessRule] = None,
-        created_from: str = "web"):
+        created_from: str = "web",
+    ):
         return 1

api/services/feature_service.py

@@ -120,7 +120,8 @@ class FeatureService:
         system_features.enable_email_password_login = dify_config.ENABLE_EMAIL_PASSWORD_LOGIN
         system_features.enable_social_oauth_login = dify_config.ENABLE_SOCIAL_OAUTH_LOGIN
         system_features.is_allow_register = dify_config.ALLOW_REGISTER
-        system_features.is_allow_create_workspace = dify_config.ALLOW_CREATE_WORKSPACE
+        # system_features.is_allow_create_workspace = dify_config.ALLOW_CREATE_WORKSPACE
+        system_features.is_allow_create_workspace = True
         system_features.is_email_setup = dify_config.MAIL_TYPE is not None and dify_config.MAIL_TYPE != ""
 
     @classmethod

api/tasks/batch_clean_template_task.py

@@ -30,7 +30,6 @@ def batch_clean_template_task(template_ids: list[str], dataset_id: str, doc_form
         if not dataset:
             raise Exception("template has no dataset")
 
-
         if file_ids:
             files = db.session.query(UploadFile).filter(UploadFile.id.in_(file_ids)).all()
             for file in files: