From 2c6e00174ba55c46297ec684bf826179a7e13e9e Mon Sep 17 00:00:00 2001
From: Jyong <76649700+JohnJyong@users.noreply.github.com>
Date: Sun, 16 Jul 2023 13:21:56 +0800
Subject: [PATCH] add document limit check (#570)

---
 api/config.py                   |  5 ++++-
 api/services/dataset_service.py | 23 +++++++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/api/config.py b/api/config.py
index d7cacaf001..99ecc67656 100644
--- a/api/config.py
+++ b/api/config.py
@@ -50,7 +50,8 @@ DEFAULTS = {
     'PDF_PREVIEW': 'True',
     'LOG_LEVEL': 'INFO',
     'DISABLE_PROVIDER_CONFIG_VALIDATION': 'False',
-    'DEFAULT_LLM_PROVIDER': 'openai'
+    'DEFAULT_LLM_PROVIDER': 'openai',
+    'TENANT_DOCUMENT_COUNT': 100
 }
 
 
@@ -207,6 +208,8 @@ class Config:
         self.NOTION_INTERNAL_SECRET = get_env('NOTION_INTERNAL_SECRET')
         self.NOTION_INTEGRATION_TOKEN = get_env('NOTION_INTEGRATION_TOKEN')
 
+        self.TENANT_DOCUMENT_COUNT = get_env('TENANT_DOCUMENT_COUNT')
+
 
 class CloudEditionConfig(Config):
 
diff --git a/api/services/dataset_service.py b/api/services/dataset_service.py
index 381c4d843d..182adc6c78 100644
--- a/api/services/dataset_service.py
+++ b/api/services/dataset_service.py
@@ -4,6 +4,9 @@ import datetime
 import time
 import random
 from typing import Optional, List
+
+from flask import current_app
+
 from extensions.ext_redis import redis_client
 from flask_login import current_user
 
@@ -374,6 +377,12 @@ class DocumentService:
     def save_document_with_dataset_id(dataset: Dataset, document_data: dict,
                                       account: Account, dataset_process_rule: Optional[DatasetProcessRule] = None,
                                       created_from: str = 'web'):
+        # check document limit
+        if current_app.config['EDITION'] == 'CLOUD':
+            documents_count = DocumentService.get_tenant_documents_count()
+            tenant_document_count = int(current_app.config['TENANT_DOCUMENT_COUNT'])
+            if documents_count > tenant_document_count:
+                raise ValueError(f"over document limit {tenant_document_count}.")
         # if dataset is empty, update dataset data_source_type
         if not dataset.data_source_type:
             dataset.data_source_type = document_data["data_source"]["type"]
@@ -521,6 +530,14 @@ class DocumentService:
         )
         return document
 
+    @staticmethod
+    def get_tenant_documents_count():
+        documents_count = Document.query.filter(Document.completed_at.isnot(None),
+                                                Document.enabled == True,
+                                                Document.archived == False,
+                                                Document.tenant_id == current_user.current_tenant_id).count()
+        return documents_count
+
     @staticmethod
     def update_document_with_dataset_id(dataset: Dataset, document_data: dict,
                                         account: Account, dataset_process_rule: Optional[DatasetProcessRule] = None,
@@ -616,6 +633,12 @@ class DocumentService:
 
     @staticmethod
     def save_document_without_dataset_id(tenant_id: str, document_data: dict, account: Account):
+        # check document limit
+        if current_app.config['EDITION'] == 'CLOUD':
+            documents_count = DocumentService.get_tenant_documents_count()
+            tenant_document_count = int(current_app.config['TENANT_DOCUMENT_COUNT'])
+            if documents_count > tenant_document_count:
+                raise ValueError(f"over document limit {tenant_document_count}.")
         # save dataset
         dataset = Dataset(
             tenant_id=tenant_id,