hz/ebiz-dify-ai
1
0
Fork 0
mirror of http://112.124.100.131/huang.ze/ebiz-dify-ai.git synced 2025-12-09 02:46:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

add document limit check (#570)

Browse Source
This commit is contained in:
Jyong
2023-07-16 13:21:56 +08:00
committed by GitHub
parent 24f3456990
commit 2c6e00174b
2 changed files with 27 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
api/services/dataset_service.py
View File

@@ -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,