admin/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

perf: optimize save_document_with_dataset_id (#29550)

Browse Source
This commit is contained in:
wangxiaolei
2025-12-12 13:14:45 +08:00
committed by GitHub
parent 12e39365fa
commit ac40309850
1 changed files with 70 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
api/services/dataset_service.py
View File

@@ -1636,6 +1636,20 @@ class DocumentService:
return [], ""
db.session.add(dataset_process_rule)
db.session.flush()
else:
# Fallback when no process_rule provided in knowledge_config:
# 1) reuse dataset.latest_process_rule if present
# 2) otherwise create an automatic rule
dataset_process_rule = getattr(dataset, "latest_process_rule", None)
if not dataset_process_rule:
dataset_process_rule = DatasetProcessRule(
dataset_id=dataset.id,
mode="automatic",
rules=json.dumps(DatasetProcessRule.AUTOMATIC_RULES),
created_by=account.id,
)
db.session.add(dataset_process_rule)
db.session.flush()
lock_name = f"add_document_lock_dataset_id_{dataset.id}"
try:
with redis_client.lock(lock_name, timeout=600):
@@ -1647,35 +1661,36 @@ class DocumentService:
if not knowledge_config.data_source.info_list.file_info_list:
raise ValueError("File source info is required")
upload_file_list = knowledge_config.data_source.info_list.file_info_list.file_ids
for file_id in upload_file_list:
file = (
files = (
db.session.query(UploadFile)
.where(UploadFile.tenant_id == dataset.tenant_id, UploadFile.id == file_id)
.first()
.where(
UploadFile.tenant_id == dataset.tenant_id,
UploadFile.id.in_(upload_file_list),
)
.all()
)
if len(files) != len(set(upload_file_list)):
raise FileNotExistsError("One or more files not found.")
# raise error if file not found
if not file:
raise FileNotExistsError()
file_name = file.name
data_source_info: dict[str, str | bool] = {
"upload_file_id": file_id,
}
# check duplicate
if knowledge_config.duplicate:
document = (
file_names = [file.name for file in files]
db_documents = (
db.session.query(Document)
.filter_by(
dataset_id=dataset.id,
tenant_id=current_user.current_tenant_id,
data_source_type="upload_file",
enabled=True,
name=file_name,
.where(
Document.dataset_id == dataset.id,
Document.tenant_id == current_user.current_tenant_id,
Document.data_source_type == "upload_file",
Document.enabled == True,
Document.name.in_(file_names),
)
.first()
.all()
)
if document:
documents_map = {document.name: document for document in db_documents}
for file in files:
data_source_info: dict[str, str | bool] = {
"upload_file_id": file.id,
}
document = documents_map.get(file.name)
if knowledge_config.duplicate and document:
document.dataset_process_rule_id = dataset_process_rule.id
document.updated_at = naive_utc_now()
document.created_from = created_from
@@ -1688,6 +1703,7 @@ class DocumentService:
documents.append(document)
duplicate_document_ids.append(document.id)
continue
else:
document = DocumentService.build_document(
dataset,
dataset_process_rule.id,
@@ -1698,7 +1714,7 @@ class DocumentService:
created_from,
position,
account,
file_name,
file.name,
batch,
)
db.session.add(document)