opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-24 23:22:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: inject more db.session (#38045)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Asuka Minato
2026-06-30 06:55:54 +00:00
committed by GitHub
co-authored by autofix-ci[bot]
parent 528bf95d1b
commit 102e1ede6e
51 changed files with 1483 additions and 893 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/controllers/console/datasets/data_source.py
+65 -66
View File
@@ -243,72 +243,71 @@ class DataSourceNotionListApi(Resource):
if not credential:
raise NotFound("Credential not found.")
exist_page_ids = []
with sessionmaker(db.engine).begin() as session:
# import notion in the exist dataset
if query.dataset_id:
dataset = DatasetService.get_dataset(query.dataset_id)
if not dataset:
raise NotFound("Dataset not found.")
if dataset.data_source_type != "notion_import":
raise ValueError("Dataset is not notion type.")
# import notion in the exist dataset
if query.dataset_id:
dataset = DatasetService.get_dataset(query.dataset_id, db.session)
if not dataset:
raise NotFound("Dataset not found.")
if dataset.data_source_type != "notion_import":
raise ValueError("Dataset is not notion type.")
documents = session.scalars(
select(Document).where(
Document.dataset_id == query.dataset_id,
Document.tenant_id == current_tenant_id,
Document.data_source_type == "notion_import",
Document.enabled.is_(True),
)
).all()
if documents:
for document in documents:
data_source_info = json.loads(document.data_source_info)
exist_page_ids.append(data_source_info["notion_page_id"])
# get all authorized pages
from core.datasource.datasource_manager import DatasourceManager
datasource_runtime = DatasourceManager.get_datasource_runtime(
provider_id="langgenius/notion_datasource/notion_datasource",
datasource_name="notion_datasource",
tenant_id=current_tenant_id,
datasource_type=DatasourceProviderType.ONLINE_DOCUMENT,
)
datasource_provider_service = DatasourceProviderService()
if credential:
datasource_runtime.runtime.credentials = credential
datasource_runtime = cast(OnlineDocumentDatasourcePlugin, datasource_runtime)
online_document_result: Generator[OnlineDocumentPagesMessage, None, None] = (
datasource_runtime.get_online_document_pages(
user_id=current_user.id,
datasource_parameters={},
provider_type=datasource_runtime.datasource_provider_type(),
documents = db.session.scalars(
select(Document).where(
Document.dataset_id == query.dataset_id,
Document.tenant_id == current_tenant_id,
Document.data_source_type == "notion_import",
Document.enabled.is_(True),
)
).all()
if documents:
for document in documents:
data_source_info = json.loads(document.data_source_info)
exist_page_ids.append(data_source_info["notion_page_id"])
# get all authorized pages
from core.datasource.datasource_manager import DatasourceManager
datasource_runtime = DatasourceManager.get_datasource_runtime(
provider_id="langgenius/notion_datasource/notion_datasource",
datasource_name="notion_datasource",
tenant_id=current_tenant_id,
datasource_type=DatasourceProviderType.ONLINE_DOCUMENT,
)
datasource_provider_service = DatasourceProviderService()
if credential:
datasource_runtime.runtime.credentials = credential
datasource_runtime = cast(OnlineDocumentDatasourcePlugin, datasource_runtime)
online_document_result: Generator[OnlineDocumentPagesMessage, None, None] = (
datasource_runtime.get_online_document_pages(
user_id=current_user.id,
datasource_parameters={},
provider_type=datasource_runtime.datasource_provider_type(),
)
try:
pages = []
workspace_info = {}
for message in online_document_result:
result = message.result
for info in result:
workspace_info = {
"workspace_id": info.workspace_id,
"workspace_name": info.workspace_name,
"workspace_icon": info.workspace_icon,
)
try:
pages = []
workspace_info = {}
for message in online_document_result:
result = message.result
for info in result:
workspace_info = {
"workspace_id": info.workspace_id,
"workspace_name": info.workspace_name,
"workspace_icon": info.workspace_icon,
}
for page in info.pages:
page_info = {
"page_id": page.page_id,
"page_name": page.page_name,
"type": page.type,
"parent_id": page.parent_id,
"is_bound": page.page_id in exist_page_ids,
"page_icon": page.page_icon,
}
for page in info.pages:
page_info = {
"page_id": page.page_id,
"page_name": page.page_name,
"type": page.type,
"parent_id": page.parent_id,
"is_bound": page.page_id in exist_page_ids,
"page_icon": page.page_icon,
}
pages.append(page_info)
except Exception as e:
raise e
notion_info = [{**workspace_info, "pages": pages}] if workspace_info else []
return dump_response(NotionIntegrateInfoListResponse, {"notion_info": notion_info}), 200
pages.append(page_info)
except Exception as e:
raise e
notion_info = [{**workspace_info, "pages": pages}] if workspace_info else []
return dump_response(NotionIntegrateInfoListResponse, {"notion_info": notion_info}), 200
@console_ns.route("/notion/pages/<uuid:page_id>/<string:page_type>/preview")
@@ -401,11 +400,11 @@ class DataSourceNotionDatasetSyncApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, dataset_id: UUID) -> tuple[dict[str, str], int]:
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
documents = DocumentService.get_document_by_dataset_id(dataset_id_str)
documents = DocumentService.get_document_by_dataset_id(dataset_id_str, db.session)
for document in documents:
document_indexing_sync_task.delay(dataset_id_str, document.id)
return {"result": "success"}, 200
@@ -421,11 +420,11 @@ class DataSourceNotionDocumentSyncApi(Resource):
def get(self, dataset_id: UUID, document_id: UUID) -> tuple[dict[str, str], int]:
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if document is None:
raise NotFound("Document not found.")
document_indexing_sync_task.delay(dataset_id_str, document_id_str)
api/controllers/console/datasets/datasets.py
+25 -21
View File
@@ -561,6 +561,7 @@ class DatasetListApi(Resource):
provider=payload.provider,
external_knowledge_api_id=payload.external_knowledge_api_id,
external_knowledge_id=payload.external_knowledge_id,
session=db.session,
)
except services.errors.dataset.DatasetNameDuplicateError:
raise DatasetNameDuplicateError()
@@ -598,7 +599,7 @@ class DatasetApi(Resource):
@with_current_tenant_id
def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
try:
@@ -618,7 +619,7 @@ class DatasetApi(Resource):
provider_id = ModelProviderID(dataset.embedding_model_provider)
data["embedding_model_provider"] = str(provider_id)
if data.get("permission") == "partial_members":
part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session)
data.update({"partial_member_list": part_users_list})
# check embedding setting
@@ -661,7 +662,7 @@ class DatasetApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def patch(self, current_tenant_id: str, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
@@ -680,10 +681,10 @@ class DatasetApi(Resource):
# The role of the current user in the ta table must be admin, owner, editor, or dataset_operator
if not dify_config.RBAC_ENABLED:
DatasetPermissionService.check_permission(
current_user, dataset, payload.permission, payload.partial_member_list
current_user, dataset, payload.permission, payload.partial_member_list, db.session
)
dataset = DatasetService.update_dataset(dataset_id_str, payload_data, current_user)
dataset = DatasetService.update_dataset(dataset_id_str, payload_data, current_user, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
@@ -698,12 +699,14 @@ class DatasetApi(Resource):
tenant_id = current_tenant_id
if payload.partial_member_list is not None and payload.permission == DatasetPermissionEnum.PARTIAL_TEAM:
DatasetPermissionService.update_partial_member_list(tenant_id, dataset_id_str, payload.partial_member_list)
DatasetPermissionService.update_partial_member_list(
tenant_id, dataset_id_str, payload.partial_member_list, db.session
)
# clear partial member list when permission is only_me or all_team_members
elif payload.permission in {DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.ALL_TEAM}:
DatasetPermissionService.clear_partial_member_list(dataset_id_str)
DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session)
partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session)
result_data.update({"partial_member_list": partial_member_list})
return result_data, 200
@@ -722,8 +725,8 @@ class DatasetApi(Resource):
raise Forbidden()
try:
if DatasetService.delete_dataset(dataset_id_str, current_user):
DatasetPermissionService.clear_partial_member_list(dataset_id_str)
if DatasetService.delete_dataset(dataset_id_str, current_user, db.session):
DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session)
return "", 204
else:
raise NotFound("Dataset not found.")
@@ -748,7 +751,7 @@ class DatasetUseCheckApi(Resource):
def get(self, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset_is_using = DatasetService.dataset_use_check(dataset_id_str)
dataset_is_using = DatasetService.dataset_use_check(dataset_id_str, db.session)
return {"is_using": dataset_is_using}, 200
@@ -769,7 +772,7 @@ class DatasetQueryApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
@@ -910,7 +913,7 @@ class DatasetRelatedAppListApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
@@ -919,7 +922,7 @@ class DatasetRelatedAppListApi(Resource):
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
app_dataset_joins = DatasetService.get_related_apps(dataset.id)
app_dataset_joins = DatasetService.get_related_apps(dataset.id, db.session)
related_apps = []
for app_dataset_join in app_dataset_joins:
@@ -1094,7 +1097,7 @@ class DatasetEnableApiApi(Resource):
def post(self, dataset_id: UUID, status: str):
dataset_id_str = str(dataset_id)
DatasetService.update_dataset_api_status(dataset_id_str, status == "enable")
DatasetService.update_dataset_api_status(dataset_id_str, status == "enable", db.session)
return {"result": "success"}, 200
@@ -1163,10 +1166,10 @@ class DatasetErrorDocs(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
results = DocumentService.get_error_documents_by_dataset_id(dataset_id_str)
results = DocumentService.get_error_documents_by_dataset_id(dataset_id_str, db.session)
return dump_response(ErrorDocsResponse, {"data": results, "total": len(results)}), 200
@@ -1190,7 +1193,7 @@ class DatasetPermissionUserListApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
try:
@@ -1198,7 +1201,7 @@ class DatasetPermissionUserListApi(Resource):
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
partial_members_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
partial_members_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session)
return dump_response(PartialMemberListResponse, {"data": partial_members_list}), 200
@@ -1220,7 +1223,8 @@ class DatasetAutoDisableLogApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY)
def get(self, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
return dump_response(AutoDisableLogsResponse, DatasetService.get_dataset_auto_disable_logs(dataset_id_str)), 200
auto_disable_logs = DatasetService.get_dataset_auto_disable_logs(dataset_id_str, db.session)
return dump_response(AutoDisableLogsResponse, auto_disable_logs), 200
api/controllers/console/datasets/datasets_document.py
+44 -38
View File
@@ -181,7 +181,7 @@ class DocumentResource(Resource):
def get_document(
self, dataset_id: str, document_id: str, current_user: Account, current_tenant_id: str
) -> Document:
dataset = DatasetService.get_dataset(dataset_id)
dataset = DatasetService.get_dataset(dataset_id, db.session)
if not dataset:
raise NotFound("Dataset not found.")
@@ -190,7 +190,7 @@ class DocumentResource(Resource):
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
document = DocumentService.get_document(dataset_id, document_id)
document = DocumentService.get_document(dataset_id, document_id, session=db.session)
if not document:
raise NotFound("Document not found.")
@@ -201,7 +201,7 @@ class DocumentResource(Resource):
return document
def get_batch_documents(self, dataset_id: str, batch: str, current_user: Account) -> Sequence[Document]:
dataset = DatasetService.get_dataset(dataset_id)
dataset = DatasetService.get_dataset(dataset_id, db.session)
if not dataset:
raise NotFound("Dataset not found.")
@@ -210,7 +210,7 @@ class DocumentResource(Resource):
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
documents = DocumentService.get_batch_documents(dataset_id, batch)
documents = DocumentService.get_batch_documents(dataset_id, batch, db.session)
if not documents:
raise NotFound("Documents not found.")
@@ -241,7 +241,7 @@ class GetProcessRuleApi(Resource):
# get the latest process rule
document = db.get_or_404(Document, document_id)
dataset = DatasetService.get_dataset(document.dataset_id)
dataset = DatasetService.get_dataset(document.dataset_id, db.session)
if not dataset:
raise NotFound("Dataset not found.")
@@ -317,7 +317,7 @@ class DatasetDocumentListApi(Resource):
)
except (ArgumentTypeError, ValueError, Exception):
fetch = False
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
@@ -421,7 +421,7 @@ class DatasetDocumentListApi(Resource):
def post(self, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
@@ -444,8 +444,10 @@ class DatasetDocumentListApi(Resource):
DocumentService.document_create_args_validate(knowledge_config)
try:
documents, batch = DocumentService.save_document_with_dataset_id(dataset, knowledge_config, current_user)
dataset = DatasetService.get_dataset(dataset_id_str)
documents, batch = DocumentService.save_document_with_dataset_id(
dataset, knowledge_config, current_user, session=db.session
)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
@@ -464,7 +466,7 @@ class DatasetDocumentListApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def delete(self, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
# check user's model setting
@@ -472,7 +474,7 @@ class DatasetDocumentListApi(Resource):
try:
document_ids = request.args.getlist("document_id")
DocumentService.delete_documents(dataset, document_ids)
DocumentService.delete_documents(dataset, document_ids, db.session)
except services.errors.document.DocumentIndexingError:
raise DocumentIndexingError("Cannot delete document during indexing.")
@@ -531,6 +533,7 @@ class DatasetInitApi(Resource):
tenant_id=current_tenant_id,
knowledge_config=knowledge_config,
account=current_user,
session=db.session,
)
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
@@ -867,7 +870,7 @@ class DocumentApi(DocumentResource):
if metadata == "only":
response = {"id": document.id, "doc_type": document.doc_type, "doc_metadata": document.doc_metadata_details}
elif metadata == "without":
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str)
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session)
document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {}
response = {
"id": document.id,
@@ -901,7 +904,7 @@ class DocumentApi(DocumentResource):
"need_summary": document.need_summary if document.need_summary is not None else False,
}
else:
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str)
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session)
document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {}
response = {
"id": document.id,
@@ -950,7 +953,7 @@ class DocumentApi(DocumentResource):
def delete(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
# check user's model setting
@@ -959,7 +962,7 @@ class DocumentApi(DocumentResource):
document = self.get_document(dataset_id_str, document_id_str, current_user, current_tenant_id)
try:
DocumentService.delete_document(document)
DocumentService.delete_document(document, db.session)
except services.errors.document.DocumentIndexingError:
raise DocumentIndexingError("Cannot delete document during indexing.")
@@ -983,7 +986,7 @@ class DocumentDownloadApi(DocumentResource):
def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID) -> dict[str, Any]:
# Reuse the shared permission/tenant checks implemented in DocumentResource.
document = self.get_document(str(dataset_id), str(document_id), current_user, current_tenant_id)
return {"url": DocumentService.get_document_download_url(document)}
return {"url": DocumentService.get_document_download_url(document, db.session)}
@console_ns.route("/datasets/<uuid:dataset_id>/documents/download-zip")
@@ -1013,6 +1016,7 @@ class DocumentBatchDownloadZipApi(DocumentResource):
document_ids=document_ids,
tenant_id=current_tenant_id,
current_user=current_user,
session=db.session,
)
# Delegate ZIP packing to FileService, but keep Flask response+cleanup in the route.
@@ -1161,7 +1165,7 @@ class DocumentStatusApi(DocumentResource):
self, current_user: Account, dataset_id: UUID, action: Literal["enable", "disable", "archive", "un_archive"]
):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
@@ -1178,7 +1182,7 @@ class DocumentStatusApi(DocumentResource):
document_ids = request.args.getlist("document_id")
try:
DocumentService.batch_update_document_status(dataset, document_ids, action, current_user)
DocumentService.batch_update_document_status(dataset, document_ids, action, current_user, db.session)
except services.errors.document.DocumentIndexingError as e:
raise InvalidActionError(str(e))
except ValueError as e:
@@ -1202,11 +1206,11 @@ class DocumentPauseApi(DocumentResource):
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
document = DocumentService.get_document(dataset.id, document_id_str)
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session)
# 404 if document not found
if document is None:
@@ -1218,7 +1222,7 @@ class DocumentPauseApi(DocumentResource):
try:
# pause document
DocumentService.pause_document(document)
DocumentService.pause_document(document, db.session)
except services.errors.document.DocumentIndexingError:
raise DocumentIndexingError("Cannot pause completed document.")
@@ -1237,10 +1241,10 @@ class DocumentRecoverApi(DocumentResource):
"""recover document."""
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
document = DocumentService.get_document(dataset.id, document_id_str)
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session)
# 404 if document not found
if document is None:
@@ -1251,7 +1255,7 @@ class DocumentRecoverApi(DocumentResource):
raise ArchivedDocumentImmutableError()
try:
# pause document
DocumentService.recover_document(document)
DocumentService.recover_document(document, db.session)
except services.errors.document.DocumentIndexingError:
raise DocumentIndexingError("Document is not in paused status.")
@@ -1271,13 +1275,13 @@ class DocumentRetryApi(DocumentResource):
"""retry document."""
payload = DocumentRetryPayload.model_validate(console_ns.payload or {})
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
retry_documents = []
if not dataset:
raise NotFound("Dataset not found.")
for document_id in payload.document_ids:
try:
document = DocumentService.get_document(dataset.id, document_id)
document = DocumentService.get_document(dataset.id, document_id, session=db.session)
# 404 if document not found
if document is None:
@@ -1295,7 +1299,7 @@ class DocumentRetryApi(DocumentResource):
logger.exception("Failed to retry document, document id: %s", document_id)
continue
# retry document
DocumentService.retry_document(dataset_id_str, retry_documents)
DocumentService.retry_document(dataset_id_str, retry_documents, db.session)
return "", 204
@@ -1313,14 +1317,14 @@ class DocumentRenameApi(DocumentResource):
# The role of the current user in the ta table must be admin, owner, editor, or dataset_operator
if not current_user.is_dataset_editor:
raise Forbidden()
dataset = DatasetService.get_dataset(dataset_id)
dataset = DatasetService.get_dataset(dataset_id, db.session)
if not dataset:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_operator_permission(current_user, dataset)
DatasetService.check_dataset_operator_permission(current_user, dataset, session=db.session)
payload = DocumentRenamePayload.model_validate(console_ns.payload or {})
try:
document = DocumentService.rename_document(str(dataset_id), str(document_id), payload.name)
document = DocumentService.rename_document(str(dataset_id), str(document_id), payload.name, db.session)
except services.errors.document.DocumentIndexingError:
raise DocumentIndexingError("Cannot delete document during indexing.")
@@ -1338,11 +1342,11 @@ class WebsiteDocumentSyncApi(DocumentResource):
def get(self, current_tenant_id: str, dataset_id: UUID, document_id: UUID):
"""sync website document."""
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
document_id_str = str(document_id)
document = DocumentService.get_document(dataset.id, document_id_str)
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
if document.tenant_id != current_tenant_id:
@@ -1353,7 +1357,7 @@ class WebsiteDocumentSyncApi(DocumentResource):
if DocumentService.check_archived(document):
raise ArchivedDocumentImmutableError()
# sync document
DocumentService.sync_website_document(dataset_id_str, document)
DocumentService.sync_website_document(dataset_id_str, document, db.session)
return {"result": "success"}, 200
@@ -1373,10 +1377,10 @@ class DocumentPipelineExecutionLogApi(DocumentResource):
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
document = DocumentService.get_document(dataset.id, document_id_str)
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
log = db.session.scalar(
@@ -1431,7 +1435,7 @@ class DocumentGenerateSummaryApi(Resource):
dataset_id_str = str(dataset_id)
# Get dataset
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
@@ -1465,7 +1469,7 @@ class DocumentGenerateSummaryApi(Resource):
raise ValueError("Summary index is not enabled for this dataset. Please enable it in the dataset settings.")
# Verify all documents exist and belong to the dataset
documents = DocumentService.get_documents_by_ids(dataset_id_str, document_list)
documents = DocumentService.get_documents_by_ids(dataset_id_str, document_list, db.session)
if len(documents) != len(document_list):
found_ids = {doc.id for doc in documents}
@@ -1481,6 +1485,7 @@ class DocumentGenerateSummaryApi(Resource):
DocumentService.update_documents_need_summary(
dataset_id=dataset_id_str,
document_ids=document_ids_to_update,
session=db.session,
need_summary=True,
)
@@ -1531,7 +1536,7 @@ class DocumentSummaryStatusApi(DocumentResource):
document_id_str = str(document_id)
# Get dataset
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
@@ -1547,6 +1552,7 @@ class DocumentSummaryStatusApi(DocumentResource):
result = SummaryIndexService.get_document_summary_status_detail(
document_id=document_id_str,
dataset_id=dataset_id_str,
session=db.session,
)
return result, 200
api/controllers/console/datasets/datasets_segments.py
+39 -33
View File
@@ -176,7 +176,7 @@ class DatasetDocumentSegmentListApi(Resource):
def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
dataset_id_str = str(dataset_id)
document_id_str = str(document_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
@@ -185,7 +185,7 @@ class DatasetDocumentSegmentListApi(Resource):
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
@@ -286,14 +286,14 @@ class DatasetDocumentSegmentListApi(Resource):
def delete(self, current_user: Account, dataset_id: UUID, document_id: UUID):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
segment_ids = request.args.getlist("segment_id")
@@ -305,7 +305,7 @@ class DatasetDocumentSegmentListApi(Resource):
DatasetService.check_dataset_permission(dataset, current_user, db.session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
SegmentService.delete_segments(segment_ids, document, dataset)
SegmentService.delete_segments(segment_ids, document, dataset, db.session)
return "", 204
@@ -331,11 +331,11 @@ class DatasetDocumentSegmentApi(Resource):
action: Literal["enable", "disable"],
):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
# check user's model setting
@@ -371,7 +371,7 @@ class DatasetDocumentSegmentApi(Resource):
if cache_result is not None:
raise InvalidActionError("Document is being indexed, please try again later")
try:
SegmentService.update_segments_status(segment_ids, action, dataset, document)
SegmentService.update_segments_status(segment_ids, action, dataset, document, db.session)
except Exception as e:
raise InvalidActionError(str(e))
return dump_response(SimpleResultResponse, {"result": "success"}), 200
@@ -394,12 +394,12 @@ class DatasetDocumentSegmentAddApi(Resource):
def post(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
if not current_user.is_dataset_editor:
@@ -428,7 +428,7 @@ class DatasetDocumentSegmentAddApi(Resource):
payload = SegmentCreatePayload.model_validate(console_ns.payload or {})
payload_dict = payload.model_dump(exclude_none=True)
SegmentService.segment_create_args_validate(payload_dict, document)
segment = type_cast(DocumentSegment, SegmentService.create_segment(payload_dict, document, dataset))
segment = type_cast(DocumentSegment, SegmentService.create_segment(payload_dict, document, dataset, db.session))
summary = SummaryIndexService.get_segment_summary(segment_id=segment.id, dataset_id=dataset_id_str)
response = {
"data": segment_response_with_summary(segment, summary.summary_content if summary else None),
@@ -455,14 +455,14 @@ class DatasetDocumentSegmentUpdateApi(Resource):
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
@@ -504,7 +504,11 @@ class DatasetDocumentSegmentUpdateApi(Resource):
# Update segment (summary update with change detection is handled in SegmentService.update_segment)
segment = SegmentService.update_segment(
SegmentUpdateArgs.model_validate(payload.model_dump(exclude_none=True)), segment, document, dataset
SegmentUpdateArgs.model_validate(payload.model_dump(exclude_none=True)),
segment,
document,
dataset,
db.session,
)
summary = SummaryIndexService.get_segment_summary(segment_id=segment.id, dataset_id=dataset_id_str)
response = {
@@ -527,14 +531,14 @@ class DatasetDocumentSegmentUpdateApi(Resource):
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
# check segment
@@ -553,7 +557,7 @@ class DatasetDocumentSegmentUpdateApi(Resource):
DatasetService.check_dataset_permission(dataset, current_user, db.session)
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
SegmentService.delete_segment(segment, document, dataset)
SegmentService.delete_segment(segment, document, dataset, db.session)
return "", 204
@@ -576,12 +580,12 @@ class DatasetDocumentSegmentBatchImportApi(Resource):
def post(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, document_id: UUID):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
@@ -651,12 +655,12 @@ class ChildChunkAddApi(Resource):
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
# check segment
@@ -693,7 +697,7 @@ class ChildChunkAddApi(Resource):
# validate args
try:
payload = ChildChunkCreatePayload.model_validate(console_ns.payload or {})
child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset)
child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset, db.session)
except ChildChunkIndexingServiceError as e:
raise ChildChunkIndexingError(str(e))
return dump_response(ChildChunkDetailResponse, {"data": child_chunk}), 200
@@ -709,14 +713,14 @@ class ChildChunkAddApi(Resource):
def get(self, current_tenant_id: str, dataset_id: UUID, document_id: UUID, segment_id: UUID):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
# check segment
@@ -766,14 +770,14 @@ class ChildChunkAddApi(Resource):
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
# check segment
@@ -795,7 +799,7 @@ class ChildChunkAddApi(Resource):
# validate args
payload = ChildChunkBatchUpdatePayload.model_validate(console_ns.payload or {})
try:
child_chunks = SegmentService.update_child_chunks(payload.chunks, segment, document, dataset)
child_chunks = SegmentService.update_child_chunks(payload.chunks, segment, document, dataset, db.session)
except ChildChunkIndexingServiceError as e:
raise ChildChunkIndexingError(str(e))
return dump_response(ChildChunkBatchUpdateResponse, {"data": child_chunks}), 200
@@ -825,14 +829,14 @@ class ChildChunkUpdateApi(Resource):
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
# check segment
@@ -866,7 +870,7 @@ class ChildChunkUpdateApi(Resource):
except services.errors.account.NoPermissionError as e:
raise Forbidden(str(e))
try:
SegmentService.delete_child_chunk(child_chunk, dataset)
SegmentService.delete_child_chunk(child_chunk, dataset, db.session)
except ChildChunkDeleteIndexServiceError as e:
raise ChildChunkDeleteIndexError(str(e))
return "", 204
@@ -893,14 +897,14 @@ class ChildChunkUpdateApi(Resource):
):
# check dataset
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if not dataset:
raise NotFound("Dataset not found.")
# check user's model setting
DatasetService.check_dataset_model_setting(dataset)
# check document
document_id_str = str(document_id)
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
# check segment
@@ -936,7 +940,9 @@ class ChildChunkUpdateApi(Resource):
# validate args
try:
payload = ChildChunkUpdatePayload.model_validate(console_ns.payload or {})
child_chunk = SegmentService.update_child_chunk(payload.content, child_chunk, segment, document, dataset)
child_chunk = SegmentService.update_child_chunk(
payload.content, child_chunk, segment, document, dataset, db.session
)
except ChildChunkIndexingServiceError as e:
raise ChildChunkIndexingError(str(e))
return dump_response(ChildChunkDetailResponse, {"data": child_chunk}), 200
api/controllers/console/datasets/external.py
+1 -1
View File
@@ -377,7 +377,7 @@ class ExternalKnowledgeHitTestingApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_PIPELINE_TEST)
def post(self, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
api/controllers/console/datasets/hit_testing_base.py
+1 -1
View File
@@ -85,7 +85,7 @@ class DatasetsHitTestingBase:
dataset_id: str, current_user: Account | None = None, current_tenant_id: str | None = None
) -> Dataset:
current_user, _ = resolve_account_fallback(current_user, current_tenant_id)
dataset = DatasetService.get_dataset(dataset_id)
dataset = DatasetService.get_dataset(dataset_id, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
api/controllers/console/datasets/metadata.py
+6 -6
View File
@@ -61,7 +61,7 @@ class DatasetMetadataCreateApi(Resource):
metadata_args = MetadataArgs.model_validate(console_ns.payload or {})
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session)
@@ -81,7 +81,7 @@ class DatasetMetadataCreateApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_CREATE_AND_MANAGEMENT)
def get(self, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
metadata = MetadataService.get_dataset_metadatas(db.session(), dataset)
@@ -105,7 +105,7 @@ class DatasetMetadataApi(Resource):
dataset_id_str = str(dataset_id)
metadata_id_str = str(metadata_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session)
@@ -125,7 +125,7 @@ class DatasetMetadataApi(Resource):
def delete(self, current_user: Account, dataset_id: UUID, metadata_id: UUID):
dataset_id_str = str(dataset_id)
metadata_id_str = str(metadata_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session)
@@ -162,7 +162,7 @@ class DatasetMetadataBuiltInFieldActionApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def post(self, current_user: Account, dataset_id: UUID, action: Literal["enable", "disable"]):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session)
@@ -191,7 +191,7 @@ class DocumentMetadataEditApi(Resource):
@rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT)
def post(self, current_user: Account, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session)
api/controllers/console/datasets/rag_pipeline/rag_pipeline_datasets.py
+8 -8
View File
@@ -1,6 +1,5 @@
from flask_restx import Resource
from pydantic import BaseModel
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden
import services
@@ -66,19 +65,19 @@ class CreateRagPipelineDatasetApi(Resource):
yaml_content=payload.yaml_content,
)
try:
with Session(db.engine, expire_on_commit=False) as session:
rag_pipeline_dsl_service = RagPipelineDslService(session)
import_info = rag_pipeline_dsl_service.create_rag_pipeline_dataset(
tenant_id=current_tenant_id,
rag_pipeline_dataset_create_entity=rag_pipeline_dataset_create_entity,
)
session.commit()
rag_pipeline_dsl_service = RagPipelineDslService(db.session)
import_info = rag_pipeline_dsl_service.create_rag_pipeline_dataset(
tenant_id=current_tenant_id,
rag_pipeline_dataset_create_entity=rag_pipeline_dataset_create_entity,
)
if rag_pipeline_dataset_create_entity.permission == "partial_members":
DatasetPermissionService.update_partial_member_list(
current_tenant_id,
import_info["dataset_id"],
rag_pipeline_dataset_create_entity.partial_member_list,
db.session,
)
db.session.commit()
except services.errors.dataset.DatasetNameDuplicateError:
raise DatasetNameDuplicateError()
@@ -111,5 +110,6 @@ class CreateEmptyRagPipelineDatasetApi(Resource):
permission=DatasetPermissionEnum.ONLY_ME,
partial_member_list=None,
),
session=db.session,
)
return dump_response(DatasetDetailResponse, dataset), 201
api/controllers/service_api/dataset/dataset.py
+15 -11
View File
@@ -519,6 +519,7 @@ class DatasetListApi(DatasetApiResource):
embedding_model_name=payload.embedding_model,
retrieval_model=payload.retrieval_model,
summary_index_setting=payload.summary_index_setting,
session=db.session,
)
except services.errors.dataset.DatasetNameDuplicateError:
raise DatasetNameDuplicateError()
@@ -561,7 +562,7 @@ class DatasetApi(DatasetApiResource):
)
def get(self, _, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
try:
@@ -597,7 +598,7 @@ class DatasetApi(DatasetApiResource):
retrieval_model_dict["search_method"] = "keyword_search"
if data.get("permission") == "partial_members":
part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session)
data.update({"partial_member_list": part_users_list})
return _dump_service_dataset_with_partial_members(data), 200
@@ -635,7 +636,7 @@ class DatasetApi(DatasetApiResource):
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def patch(self, _, dataset_id: UUID):
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
@@ -676,9 +677,10 @@ class DatasetApi(DatasetApiResource):
dataset,
str(payload.permission) if payload.permission else None,
payload.partial_member_list,
db.session,
)
dataset = DatasetService.update_dataset(dataset_id_str, update_data, current_user)
dataset = DatasetService.update_dataset(dataset_id_str, update_data, current_user, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
@@ -688,12 +690,14 @@ class DatasetApi(DatasetApiResource):
tenant_id = current_user.current_tenant_id
if payload.partial_member_list and payload.permission == DatasetPermissionEnum.PARTIAL_TEAM:
DatasetPermissionService.update_partial_member_list(tenant_id, dataset_id_str, payload.partial_member_list)
DatasetPermissionService.update_partial_member_list(
tenant_id, dataset_id_str, payload.partial_member_list, db.session
)
# clear partial member list when permission is only_me or all_team_members
elif payload.permission in {DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.ALL_TEAM}:
DatasetPermissionService.clear_partial_member_list(dataset_id_str)
DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session)
partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str, db.session)
result_data.update({"partial_member_list": partial_member_list})
return _dump_service_dataset_with_partial_members(result_data), 200
@@ -746,8 +750,8 @@ class DatasetApi(DatasetApiResource):
dataset_id_str = str(dataset_id)
try:
if DatasetService.delete_dataset(dataset_id_str, current_user):
DatasetPermissionService.clear_partial_member_list(dataset_id_str)
if DatasetService.delete_dataset(dataset_id_str, current_user, db.session):
DatasetPermissionService.clear_partial_member_list(dataset_id_str, db.session)
return "", 204
else:
raise NotFound("Dataset not found.")
@@ -812,7 +816,7 @@ class DocumentStatusApi(DatasetApiResource):
InvalidActionError: If the action is invalid or cannot be performed.
"""
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
@@ -831,7 +835,7 @@ class DocumentStatusApi(DatasetApiResource):
document_ids = data.get("document_ids", [])
try:
DocumentService.batch_update_document_status(dataset, document_ids, action, current_user)
DocumentService.batch_update_document_status(dataset, document_ids, action, current_user, db.session)
except services.errors.document.DocumentIndexingError as e:
raise InvalidActionError(str(e))
except ValueError as e:
api/controllers/service_api/dataset/document.py
+13 -8
View File
@@ -400,6 +400,7 @@ def _create_document_by_text(tenant_id: str, dataset_id: UUID) -> tuple[Mapping[
account=current_user,
dataset_process_rule=dataset.latest_process_rule if "process_rule" not in args else None,
created_from="api",
session=db.session,
)
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
@@ -459,6 +460,7 @@ def _update_document_by_text(tenant_id: str, dataset_id: UUID, document_id: UUID
account=current_user,
dataset_process_rule=dataset.latest_process_rule if "process_rule" not in args else None,
created_from="api",
session=db.session,
)
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
@@ -756,6 +758,7 @@ class DocumentAddByFileApi(DatasetApiResource):
account=dataset.created_by_account,
dataset_process_rule=dataset_process_rule,
created_from="api",
session=db.session,
)
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
@@ -832,6 +835,7 @@ def _update_document_by_file(tenant_id: str, dataset_id: UUID, document_id: UUID
account=dataset.created_by_account,
dataset_process_rule=dataset.latest_process_rule if "process_rule" not in args else None,
created_from="api",
session=db.session,
)
except ProviderTokenNotInitError as ex:
raise ProviderNotInitializeError(ex.description)
@@ -1002,6 +1006,7 @@ class DocumentBatchDownloadZipApi(DatasetApiResource):
document_ids=[str(document_id) for document_id in payload.document_ids],
tenant_id=str(tenant_id),
current_user=current_user,
session=db.session,
)
with ExitStack() as stack:
@@ -1058,7 +1063,7 @@ class DocumentIndexingStatusApi(DatasetApiResource):
if not dataset:
raise NotFound("Dataset not found.")
# get documents
documents = DocumentService.get_batch_documents(dataset_id_str, batch)
documents = DocumentService.get_batch_documents(dataset_id_str, batch, db.session)
if not documents:
raise NotFound("Documents not found.")
documents_status = []
@@ -1134,7 +1139,7 @@ class DocumentDownloadApi(DatasetApiResource):
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
def get(self, tenant_id, dataset_id: UUID, document_id: UUID):
dataset = self.get_dataset(str(dataset_id), str(tenant_id))
document = DocumentService.get_document(dataset.id, str(document_id))
document = DocumentService.get_document(dataset.id, str(document_id), session=db.session)
if not document:
raise NotFound("Document not found.")
@@ -1142,7 +1147,7 @@ class DocumentDownloadApi(DatasetApiResource):
if document.tenant_id != str(tenant_id):
raise Forbidden("No permission.")
return {"url": DocumentService.get_document_download_url(document)}
return {"url": DocumentService.get_document_download_url(document, db.session)}
@service_api_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>")
@@ -1190,7 +1195,7 @@ class DocumentApi(DatasetApiResource):
dataset = self.get_dataset(dataset_id_str, tenant_id)
document = DocumentService.get_document(dataset.id, document_id_str)
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
@@ -1215,7 +1220,7 @@ class DocumentApi(DatasetApiResource):
if metadata == "only":
response = {"id": document.id, "doc_type": document.doc_type, "doc_metadata": document.doc_metadata_details}
elif metadata == "without":
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str)
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session)
document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {}
data_source_info = document.data_source_detail_dict
response = {
@@ -1250,7 +1255,7 @@ class DocumentApi(DatasetApiResource):
"need_summary": document.need_summary if document.need_summary is not None else False,
}
else:
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str)
dataset_process_rules = DatasetService.get_process_rules(dataset_id_str, db.session)
document_process_rules = document.dataset_process_rule.to_dict() if document.dataset_process_rule else {}
data_source_info = document.data_source_detail_dict
response = {
@@ -1345,7 +1350,7 @@ class DocumentApi(DatasetApiResource):
if not dataset:
raise ValueError("Dataset does not exist.")
document = DocumentService.get_document(dataset.id, document_id_str)
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session)
# 404 if document not found
if document is None:
@@ -1357,7 +1362,7 @@ class DocumentApi(DatasetApiResource):
try:
# delete document
DocumentService.delete_document(document)
DocumentService.delete_document(document, db.session)
except services.errors.document.DocumentIndexingError:
raise DocumentIndexingError("Cannot delete document during indexing.")
api/controllers/service_api/dataset/metadata.py
+6 -6
View File
@@ -81,7 +81,7 @@ class DatasetMetadataCreateServiceApi(DatasetApiResource):
metadata_args = MetadataArgs.model_validate(service_api_ns.payload or {})
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session)
@@ -116,7 +116,7 @@ class DatasetMetadataCreateServiceApi(DatasetApiResource):
def get(self, tenant_id, dataset_id: UUID):
"""Get all metadata for a dataset."""
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
metadata = MetadataService.get_dataset_metadatas(db.session(), dataset)
@@ -154,7 +154,7 @@ class DatasetMetadataServiceApi(DatasetApiResource):
dataset_id_str = str(dataset_id)
metadata_id_str = str(metadata_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session)
@@ -189,7 +189,7 @@ class DatasetMetadataServiceApi(DatasetApiResource):
"""Delete metadata."""
dataset_id_str = str(dataset_id)
metadata_id_str = str(metadata_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session)
@@ -257,7 +257,7 @@ class DatasetMetadataBuiltInFieldActionServiceApi(DatasetApiResource):
def post(self, tenant_id, dataset_id: UUID, action: Literal["enable", "disable"]):
"""Enable or disable built-in metadata field."""
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session)
@@ -303,7 +303,7 @@ class DocumentMetadataEditServiceApi(DatasetApiResource):
def post(self, tenant_id, dataset_id: UUID):
"""Update metadata for multiple documents."""
dataset_id_str = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id_str)
dataset = DatasetService.get_dataset(dataset_id_str, db.session)
if dataset is None:
raise NotFound("Dataset not found.")
DatasetService.check_dataset_permission(dataset, current_user, db.session)
api/controllers/service_api/dataset/segment.py
+42 -24
View File
@@ -175,7 +175,7 @@ class SegmentApi(DatasetApiResource):
raise NotFound("Dataset not found.")
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset.id, document_id_str)
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
if document.indexing_status != "completed":
@@ -210,7 +210,9 @@ class SegmentApi(DatasetApiResource):
for args_item in segment_items:
SegmentService.segment_create_args_validate(args_item, document)
segments = cast(list[DocumentSegment], SegmentService.multi_create_segment(segment_items, document, dataset))
segments = cast(
list[DocumentSegment], SegmentService.multi_create_segment(segment_items, document, dataset, db.session)
)
segment_ids = [segment.id for segment in segments]
summaries: dict[str, str | None] = {}
if segment_ids:
@@ -267,7 +269,7 @@ class SegmentApi(DatasetApiResource):
raise NotFound("Dataset not found.")
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset.id, document_id_str)
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
# check embedding model setting
@@ -349,15 +351,17 @@ class DatasetSegmentApi(DatasetApiResource):
DatasetService.check_dataset_model_setting(dataset)
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
# check segment
segment = SegmentService.get_segment_by_id(segment_id=segment_id_str, tenant_id=current_tenant_id)
segment = SegmentService.get_segment_by_id(
segment_id=segment_id_str, tenant_id=current_tenant_id, session=db.session
)
if not segment:
raise NotFound("Segment not found.")
SegmentService.delete_segment(segment, document, dataset)
SegmentService.delete_segment(segment, document, dataset, db.session)
return "", 204
@service_api_ns.doc(
@@ -395,7 +399,7 @@ class DatasetSegmentApi(DatasetApiResource):
DatasetService.check_dataset_model_setting(dataset)
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
@@ -416,13 +420,15 @@ class DatasetSegmentApi(DatasetApiResource):
raise ProviderNotInitializeError(ex.description)
segment_id_str = str(segment_id)
# check segment
segment = SegmentService.get_segment_by_id(segment_id=segment_id_str, tenant_id=current_tenant_id)
segment = SegmentService.get_segment_by_id(
segment_id=segment_id_str, tenant_id=current_tenant_id, session=db.session
)
if not segment:
raise NotFound("Segment not found.")
payload = SegmentUpdatePayload.model_validate(service_api_ns.payload or {})
updated_segment = SegmentService.update_segment(payload.segment, segment, document, dataset)
updated_segment = SegmentService.update_segment(payload.segment, segment, document, dataset, db.session)
summary = SummaryIndexService.get_segment_summary(segment_id=updated_segment.id, dataset_id=dataset_id_str)
response = {
"data": segment_response_with_summary(updated_segment, summary.summary_content if summary else None),
@@ -469,12 +475,14 @@ class DatasetSegmentApi(DatasetApiResource):
DatasetService.check_dataset_model_setting(dataset)
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
# check segment
segment = SegmentService.get_segment_by_id(segment_id=segment_id_str, tenant_id=current_tenant_id)
segment = SegmentService.get_segment_by_id(
segment_id=segment_id_str, tenant_id=current_tenant_id, session=db.session
)
if not segment:
raise NotFound("Segment not found.")
@@ -533,13 +541,15 @@ class ChildChunkApi(DatasetApiResource):
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset.id, document_id_str)
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
# check segment
segment = SegmentService.get_segment_by_id(segment_id=segment_id_str, tenant_id=current_tenant_id)
segment = SegmentService.get_segment_by_id(
segment_id=segment_id_str, tenant_id=current_tenant_id, session=db.session
)
if not segment:
raise NotFound("Segment not found.")
@@ -564,7 +574,7 @@ class ChildChunkApi(DatasetApiResource):
payload = ChildChunkCreatePayload.model_validate(service_api_ns.payload or {})
try:
child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset)
child_chunk = SegmentService.create_child_chunk(payload.content, segment, document, dataset, db.session)
except ChildChunkIndexingServiceError as e:
raise ChildChunkIndexingError(str(e))
@@ -607,13 +617,15 @@ class ChildChunkApi(DatasetApiResource):
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset.id, document_id_str)
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
# check segment
segment = SegmentService.get_segment_by_id(segment_id=segment_id_str, tenant_id=current_tenant_id)
segment = SegmentService.get_segment_by_id(
segment_id=segment_id_str, tenant_id=current_tenant_id, session=db.session
)
if not segment:
raise NotFound("Segment not found.")
@@ -677,13 +689,15 @@ class DatasetChildChunkApi(DatasetApiResource):
document_id_str = str(document_id)
# check document
document = DocumentService.get_document(dataset.id, document_id_str)
document = DocumentService.get_document(dataset.id, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
# check segment
segment = SegmentService.get_segment_by_id(segment_id=segment_id_str, tenant_id=current_tenant_id)
segment = SegmentService.get_segment_by_id(
segment_id=segment_id_str, tenant_id=current_tenant_id, session=db.session
)
if not segment:
raise NotFound("Segment not found.")
@@ -694,7 +708,7 @@ class DatasetChildChunkApi(DatasetApiResource):
child_chunk_id_str = str(child_chunk_id)
# check child chunk
child_chunk = SegmentService.get_child_chunk_by_id(
child_chunk_id=child_chunk_id_str, tenant_id=current_tenant_id
child_chunk_id=child_chunk_id_str, tenant_id=current_tenant_id, session=db.session
)
if not child_chunk:
raise NotFound("Child chunk not found.")
@@ -704,7 +718,7 @@ class DatasetChildChunkApi(DatasetApiResource):
raise NotFound("Child chunk not found.")
try:
SegmentService.delete_child_chunk(child_chunk, dataset)
SegmentService.delete_child_chunk(child_chunk, dataset, db.session)
except ChildChunkDeleteIndexServiceError as e:
raise ChildChunkDeleteIndexError(str(e))
@@ -751,13 +765,15 @@ class DatasetChildChunkApi(DatasetApiResource):
document_id_str = str(document_id)
# get document
document = DocumentService.get_document(dataset_id_str, document_id_str)
document = DocumentService.get_document(dataset_id_str, document_id_str, session=db.session)
if not document:
raise NotFound("Document not found.")
segment_id_str = str(segment_id)
# get segment
segment = SegmentService.get_segment_by_id(segment_id=segment_id_str, tenant_id=current_tenant_id)
segment = SegmentService.get_segment_by_id(
segment_id=segment_id_str, tenant_id=current_tenant_id, session=db.session
)
if not segment:
raise NotFound("Segment not found.")
@@ -768,7 +784,7 @@ class DatasetChildChunkApi(DatasetApiResource):
child_chunk_id_str = str(child_chunk_id)
# get child chunk
child_chunk = SegmentService.get_child_chunk_by_id(
child_chunk_id=child_chunk_id_str, tenant_id=current_tenant_id
child_chunk_id=child_chunk_id_str, tenant_id=current_tenant_id, session=db.session
)
if not child_chunk:
raise NotFound("Child chunk not found.")
@@ -781,7 +797,9 @@ class DatasetChildChunkApi(DatasetApiResource):
payload = ChildChunkUpdatePayload.model_validate(service_api_ns.payload or {})
try:
child_chunk = SegmentService.update_child_chunk(payload.content, child_chunk, segment, document, dataset)
child_chunk = SegmentService.update_child_chunk(
payload.content, child_chunk, segment, document, dataset, db.session
)
except ChildChunkIndexingServiceError as e:
raise ChildChunkIndexingError(str(e))