|
|
|
@ -77,11 +77,13 @@ from tasks.sync_website_document_indexing_task import sync_website_document_inde
|
|
|
|
class DatasetService:
|
|
|
|
class DatasetService:
|
|
|
|
@staticmethod
|
|
|
|
@staticmethod
|
|
|
|
def get_datasets(page, per_page, tenant_id=None, user=None, search=None, tag_ids=None, include_all=False):
|
|
|
|
def get_datasets(page, per_page, tenant_id=None, user=None, search=None, tag_ids=None, include_all=False):
|
|
|
|
query = Dataset.query.filter(Dataset.tenant_id == tenant_id).order_by(Dataset.created_at.desc())
|
|
|
|
query = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id).order_by(Dataset.created_at.desc())
|
|
|
|
|
|
|
|
|
|
|
|
if user:
|
|
|
|
if user:
|
|
|
|
# get permitted dataset ids
|
|
|
|
# get permitted dataset ids
|
|
|
|
dataset_permission = DatasetPermission.query.filter_by(account_id=user.id, tenant_id=tenant_id).all()
|
|
|
|
dataset_permission = (
|
|
|
|
|
|
|
|
db.session.query(DatasetPermission).filter_by(account_id=user.id, tenant_id=tenant_id).all()
|
|
|
|
|
|
|
|
)
|
|
|
|
permitted_dataset_ids = {dp.dataset_id for dp in dataset_permission} if dataset_permission else None
|
|
|
|
permitted_dataset_ids = {dp.dataset_id for dp in dataset_permission} if dataset_permission else None
|
|
|
|
|
|
|
|
|
|
|
|
if user.current_role == TenantAccountRole.DATASET_OPERATOR:
|
|
|
|
if user.current_role == TenantAccountRole.DATASET_OPERATOR:
|
|
|
|
@ -153,8 +155,10 @@ class DatasetService:
|
|
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
@staticmethod
|
|
|
|
def get_datasets_by_ids(ids, tenant_id):
|
|
|
|
def get_datasets_by_ids(ids, tenant_id):
|
|
|
|
datasets = Dataset.query.filter(Dataset.id.in_(ids), Dataset.tenant_id == tenant_id).paginate(
|
|
|
|
datasets = (
|
|
|
|
page=1, per_page=len(ids), max_per_page=len(ids), error_out=False
|
|
|
|
db.session.query(Dataset)
|
|
|
|
|
|
|
|
.filter(Dataset.id.in_(ids), Dataset.tenant_id == tenant_id)
|
|
|
|
|
|
|
|
.paginate(page=1, per_page=len(ids), max_per_page=len(ids), error_out=False)
|
|
|
|
)
|
|
|
|
)
|
|
|
|
return datasets.items, datasets.total
|
|
|
|
return datasets.items, datasets.total
|
|
|
|
|
|
|
|
|
|
|
|
@ -174,7 +178,7 @@ class DatasetService:
|
|
|
|
retrieval_model: Optional[RetrievalModel] = None,
|
|
|
|
retrieval_model: Optional[RetrievalModel] = None,
|
|
|
|
):
|
|
|
|
):
|
|
|
|
# check if dataset name already exists
|
|
|
|
# check if dataset name already exists
|
|
|
|
if Dataset.query.filter_by(name=name, tenant_id=tenant_id).first():
|
|
|
|
if db.session.query(Dataset).filter_by(name=name, tenant_id=tenant_id).first():
|
|
|
|
raise DatasetNameDuplicateError(f"Dataset with name {name} already exists.")
|
|
|
|
raise DatasetNameDuplicateError(f"Dataset with name {name} already exists.")
|
|
|
|
embedding_model = None
|
|
|
|
embedding_model = None
|
|
|
|
if indexing_technique == "high_quality":
|
|
|
|
if indexing_technique == "high_quality":
|
|
|
|
@ -235,7 +239,7 @@ class DatasetService:
|
|
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
@staticmethod
|
|
|
|
def get_dataset(dataset_id) -> Optional[Dataset]:
|
|
|
|
def get_dataset(dataset_id) -> Optional[Dataset]:
|
|
|
|
dataset: Optional[Dataset] = Dataset.query.filter_by(id=dataset_id).first()
|
|
|
|
dataset: Optional[Dataset] = db.session.query(Dataset).filter_by(id=dataset_id).first()
|
|
|
|
return dataset
|
|
|
|
return dataset
|
|
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
@staticmethod
|
|
|
|
@ -436,7 +440,7 @@ class DatasetService:
|
|
|
|
# update Retrieval model
|
|
|
|
# update Retrieval model
|
|
|
|
filtered_data["retrieval_model"] = data["retrieval_model"]
|
|
|
|
filtered_data["retrieval_model"] = data["retrieval_model"]
|
|
|
|
|
|
|
|
|
|
|
|
dataset.query.filter_by(id=dataset_id).update(filtered_data)
|
|
|
|
db.session.query(Dataset).filter_by(id=dataset_id).update(filtered_data)
|
|
|
|
|
|
|
|
|
|
|
|
db.session.commit()
|
|
|
|
db.session.commit()
|
|
|
|
if action:
|
|
|
|
if action:
|
|
|
|
@ -460,7 +464,7 @@ class DatasetService:
|
|
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
@staticmethod
|
|
|
|
def dataset_use_check(dataset_id) -> bool:
|
|
|
|
def dataset_use_check(dataset_id) -> bool:
|
|
|
|
count = AppDatasetJoin.query.filter_by(dataset_id=dataset_id).count()
|
|
|
|
count = db.session.query(AppDatasetJoin).filter_by(dataset_id=dataset_id).count()
|
|
|
|
if count > 0:
|
|
|
|
if count > 0:
|
|
|
|
return True
|
|
|
|
return True
|
|
|
|
return False
|
|
|
|
return False
|
|
|
|
@ -475,7 +479,9 @@ class DatasetService:
|
|
|
|
logging.debug(f"User {user.id} does not have permission to access dataset {dataset.id}")
|
|
|
|
logging.debug(f"User {user.id} does not have permission to access dataset {dataset.id}")
|
|
|
|
raise NoPermissionError("You do not have permission to access this dataset.")
|
|
|
|
raise NoPermissionError("You do not have permission to access this dataset.")
|
|
|
|
if dataset.permission == "partial_members":
|
|
|
|
if dataset.permission == "partial_members":
|
|
|
|
user_permission = DatasetPermission.query.filter_by(dataset_id=dataset.id, account_id=user.id).first()
|
|
|
|
user_permission = (
|
|
|
|
|
|
|
|
db.session.query(DatasetPermission).filter_by(dataset_id=dataset.id, account_id=user.id).first()
|
|
|
|
|
|
|
|
)
|
|
|
|
if (
|
|
|
|
if (
|
|
|
|
not user_permission
|
|
|
|
not user_permission
|
|
|
|
and dataset.tenant_id != user.current_tenant_id
|
|
|
|
and dataset.tenant_id != user.current_tenant_id
|
|
|
|
@ -499,14 +505,16 @@ class DatasetService:
|
|
|
|
|
|
|
|
|
|
|
|
elif dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM:
|
|
|
|
elif dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM:
|
|
|
|
if not any(
|
|
|
|
if not any(
|
|
|
|
dp.dataset_id == dataset.id for dp in DatasetPermission.query.filter_by(account_id=user.id).all()
|
|
|
|
dp.dataset_id == dataset.id
|
|
|
|
|
|
|
|
for dp in db.session.query(DatasetPermission).filter_by(account_id=user.id).all()
|
|
|
|
):
|
|
|
|
):
|
|
|
|
raise NoPermissionError("You do not have permission to access this dataset.")
|
|
|
|
raise NoPermissionError("You do not have permission to access this dataset.")
|
|
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
@staticmethod
|
|
|
|
def get_dataset_queries(dataset_id: str, page: int, per_page: int):
|
|
|
|
def get_dataset_queries(dataset_id: str, page: int, per_page: int):
|
|
|
|
dataset_queries = (
|
|
|
|
dataset_queries = (
|
|
|
|
DatasetQuery.query.filter_by(dataset_id=dataset_id)
|
|
|
|
db.session.query(DatasetQuery)
|
|
|
|
|
|
|
|
.filter_by(dataset_id=dataset_id)
|
|
|
|
.order_by(db.desc(DatasetQuery.created_at))
|
|
|
|
.order_by(db.desc(DatasetQuery.created_at))
|
|
|
|
.paginate(page=page, per_page=per_page, max_per_page=100, error_out=False)
|
|
|
|
.paginate(page=page, per_page=per_page, max_per_page=100, error_out=False)
|
|
|
|
)
|
|
|
|
)
|
|
|
|
@ -515,7 +523,8 @@ class DatasetService:
|
|
|
|
@staticmethod
|
|
|
|
@staticmethod
|
|
|
|
def get_related_apps(dataset_id: str):
|
|
|
|
def get_related_apps(dataset_id: str):
|
|
|
|
return (
|
|
|
|
return (
|
|
|
|
AppDatasetJoin.query.filter(AppDatasetJoin.dataset_id == dataset_id)
|
|
|
|
db.session.query(AppDatasetJoin)
|
|
|
|
|
|
|
|
.filter(AppDatasetJoin.dataset_id == dataset_id)
|
|
|
|
.order_by(db.desc(AppDatasetJoin.created_at))
|
|
|
|
.order_by(db.desc(AppDatasetJoin.created_at))
|
|
|
|
.all()
|
|
|
|
.all()
|
|
|
|
)
|
|
|
|
)
|
|
|
|
@ -530,10 +539,14 @@ class DatasetService:
|
|
|
|
}
|
|
|
|
}
|
|
|
|
# get recent 30 days auto disable logs
|
|
|
|
# get recent 30 days auto disable logs
|
|
|
|
start_date = datetime.datetime.now() - datetime.timedelta(days=30)
|
|
|
|
start_date = datetime.datetime.now() - datetime.timedelta(days=30)
|
|
|
|
dataset_auto_disable_logs = DatasetAutoDisableLog.query.filter(
|
|
|
|
dataset_auto_disable_logs = (
|
|
|
|
DatasetAutoDisableLog.dataset_id == dataset_id,
|
|
|
|
db.session.query(DatasetAutoDisableLog)
|
|
|
|
DatasetAutoDisableLog.created_at >= start_date,
|
|
|
|
.filter(
|
|
|
|
).all()
|
|
|
|
DatasetAutoDisableLog.dataset_id == dataset_id,
|
|
|
|
|
|
|
|
DatasetAutoDisableLog.created_at >= start_date,
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
.all()
|
|
|
|
|
|
|
|
)
|
|
|
|
if dataset_auto_disable_logs:
|
|
|
|
if dataset_auto_disable_logs:
|
|
|
|
return {
|
|
|
|
return {
|
|
|
|
"document_ids": [log.document_id for log in dataset_auto_disable_logs],
|
|
|
|
"document_ids": [log.document_id for log in dataset_auto_disable_logs],
|
|
|
|
@ -873,7 +886,9 @@ class DocumentService:
|
|
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
@staticmethod
|
|
|
|
def get_documents_position(dataset_id):
|
|
|
|
def get_documents_position(dataset_id):
|
|
|
|
document = Document.query.filter_by(dataset_id=dataset_id).order_by(Document.position.desc()).first()
|
|
|
|
document = (
|
|
|
|
|
|
|
|
db.session.query(Document).filter_by(dataset_id=dataset_id).order_by(Document.position.desc()).first()
|
|
|
|
|
|
|
|
)
|
|
|
|
if document:
|
|
|
|
if document:
|
|
|
|
return document.position + 1
|
|
|
|
return document.position + 1
|
|
|
|
else:
|
|
|
|
else:
|
|
|
|
@ -1010,13 +1025,17 @@ class DocumentService:
|
|
|
|
}
|
|
|
|
}
|
|
|
|
# check duplicate
|
|
|
|
# check duplicate
|
|
|
|
if knowledge_config.duplicate:
|
|
|
|
if knowledge_config.duplicate:
|
|
|
|
document = Document.query.filter_by(
|
|
|
|
document = (
|
|
|
|
dataset_id=dataset.id,
|
|
|
|
db.session.query(Document)
|
|
|
|
tenant_id=current_user.current_tenant_id,
|
|
|
|
.filter_by(
|
|
|
|
data_source_type="upload_file",
|
|
|
|
dataset_id=dataset.id,
|
|
|
|
enabled=True,
|
|
|
|
tenant_id=current_user.current_tenant_id,
|
|
|
|
name=file_name,
|
|
|
|
data_source_type="upload_file",
|
|
|
|
).first()
|
|
|
|
enabled=True,
|
|
|
|
|
|
|
|
name=file_name,
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
.first()
|
|
|
|
|
|
|
|
)
|
|
|
|
if document:
|
|
|
|
if document:
|
|
|
|
document.dataset_process_rule_id = dataset_process_rule.id # type: ignore
|
|
|
|
document.dataset_process_rule_id = dataset_process_rule.id # type: ignore
|
|
|
|
document.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
|
|
|
|
document.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
|
|
|
|
@ -1054,12 +1073,16 @@ class DocumentService:
|
|
|
|
raise ValueError("No notion info list found.")
|
|
|
|
raise ValueError("No notion info list found.")
|
|
|
|
exist_page_ids = []
|
|
|
|
exist_page_ids = []
|
|
|
|
exist_document = {}
|
|
|
|
exist_document = {}
|
|
|
|
documents = Document.query.filter_by(
|
|
|
|
documents = (
|
|
|
|
dataset_id=dataset.id,
|
|
|
|
db.session.query(Document)
|
|
|
|
tenant_id=current_user.current_tenant_id,
|
|
|
|
.filter_by(
|
|
|
|
data_source_type="notion_import",
|
|
|
|
dataset_id=dataset.id,
|
|
|
|
enabled=True,
|
|
|
|
tenant_id=current_user.current_tenant_id,
|
|
|
|
).all()
|
|
|
|
data_source_type="notion_import",
|
|
|
|
|
|
|
|
enabled=True,
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
.all()
|
|
|
|
|
|
|
|
)
|
|
|
|
if documents:
|
|
|
|
if documents:
|
|
|
|
for document in documents:
|
|
|
|
for document in documents:
|
|
|
|
data_source_info = json.loads(document.data_source_info)
|
|
|
|
data_source_info = json.loads(document.data_source_info)
|
|
|
|
@ -1206,12 +1229,16 @@ class DocumentService:
|
|
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
@staticmethod
|
|
|
|
def get_tenant_documents_count():
|
|
|
|
def get_tenant_documents_count():
|
|
|
|
documents_count = Document.query.filter(
|
|
|
|
documents_count = (
|
|
|
|
Document.completed_at.isnot(None),
|
|
|
|
db.session.query(Document)
|
|
|
|
Document.enabled == True,
|
|
|
|
.filter(
|
|
|
|
Document.archived == False,
|
|
|
|
Document.completed_at.isnot(None),
|
|
|
|
Document.tenant_id == current_user.current_tenant_id,
|
|
|
|
Document.enabled == True,
|
|
|
|
).count()
|
|
|
|
Document.archived == False,
|
|
|
|
|
|
|
|
Document.tenant_id == current_user.current_tenant_id,
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
.count()
|
|
|
|
|
|
|
|
)
|
|
|
|
return documents_count
|
|
|
|
return documents_count
|
|
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
@staticmethod
|
|
|
|
@ -1328,7 +1355,7 @@ class DocumentService:
|
|
|
|
db.session.commit()
|
|
|
|
db.session.commit()
|
|
|
|
# update document segment
|
|
|
|
# update document segment
|
|
|
|
update_params = {DocumentSegment.status: "re_segment"}
|
|
|
|
update_params = {DocumentSegment.status: "re_segment"}
|
|
|
|
DocumentSegment.query.filter_by(document_id=document.id).update(update_params)
|
|
|
|
db.session.query(DocumentSegment).filter_by(document_id=document.id).update(update_params)
|
|
|
|
db.session.commit()
|
|
|
|
db.session.commit()
|
|
|
|
# trigger async task
|
|
|
|
# trigger async task
|
|
|
|
document_indexing_update_task.delay(document.dataset_id, document.id)
|
|
|
|
document_indexing_update_task.delay(document.dataset_id, document.id)
|
|
|
|
@ -1918,7 +1945,8 @@ class SegmentService:
|
|
|
|
@classmethod
|
|
|
|
@classmethod
|
|
|
|
def delete_segments(cls, segment_ids: list, document: Document, dataset: Dataset):
|
|
|
|
def delete_segments(cls, segment_ids: list, document: Document, dataset: Dataset):
|
|
|
|
index_node_ids = (
|
|
|
|
index_node_ids = (
|
|
|
|
DocumentSegment.query.with_entities(DocumentSegment.index_node_id)
|
|
|
|
db.session.query(DocumentSegment)
|
|
|
|
|
|
|
|
.with_entities(DocumentSegment.index_node_id)
|
|
|
|
.filter(
|
|
|
|
.filter(
|
|
|
|
DocumentSegment.id.in_(segment_ids),
|
|
|
|
DocumentSegment.id.in_(segment_ids),
|
|
|
|
DocumentSegment.dataset_id == dataset.id,
|
|
|
|
DocumentSegment.dataset_id == dataset.id,
|
|
|
|
@ -2157,12 +2185,16 @@ class SegmentService:
|
|
|
|
def get_child_chunks(
|
|
|
|
def get_child_chunks(
|
|
|
|
cls, segment_id: str, document_id: str, dataset_id: str, page: int, limit: int, keyword: Optional[str] = None
|
|
|
|
cls, segment_id: str, document_id: str, dataset_id: str, page: int, limit: int, keyword: Optional[str] = None
|
|
|
|
):
|
|
|
|
):
|
|
|
|
query = ChildChunk.query.filter_by(
|
|
|
|
query = (
|
|
|
|
tenant_id=current_user.current_tenant_id,
|
|
|
|
db.session.query(ChildChunk)
|
|
|
|
dataset_id=dataset_id,
|
|
|
|
.filter_by(
|
|
|
|
document_id=document_id,
|
|
|
|
tenant_id=current_user.current_tenant_id,
|
|
|
|
segment_id=segment_id,
|
|
|
|
dataset_id=dataset_id,
|
|
|
|
).order_by(ChildChunk.position.asc())
|
|
|
|
document_id=document_id,
|
|
|
|
|
|
|
|
segment_id=segment_id,
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
.order_by(ChildChunk.position.asc())
|
|
|
|
|
|
|
|
)
|
|
|
|
if keyword:
|
|
|
|
if keyword:
|
|
|
|
query = query.where(ChildChunk.content.ilike(f"%{keyword}%"))
|
|
|
|
query = query.where(ChildChunk.content.ilike(f"%{keyword}%"))
|
|
|
|
return query.paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
|
|
|
|
return query.paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
|
|
|
|
@ -2170,7 +2202,11 @@ class SegmentService:
|
|
|
|
@classmethod
|
|
|
|
@classmethod
|
|
|
|
def get_child_chunk_by_id(cls, child_chunk_id: str, tenant_id: str) -> Optional[ChildChunk]:
|
|
|
|
def get_child_chunk_by_id(cls, child_chunk_id: str, tenant_id: str) -> Optional[ChildChunk]:
|
|
|
|
"""Get a child chunk by its ID."""
|
|
|
|
"""Get a child chunk by its ID."""
|
|
|
|
result = ChildChunk.query.filter(ChildChunk.id == child_chunk_id, ChildChunk.tenant_id == tenant_id).first()
|
|
|
|
result = (
|
|
|
|
|
|
|
|
db.session.query(ChildChunk)
|
|
|
|
|
|
|
|
.filter(ChildChunk.id == child_chunk_id, ChildChunk.tenant_id == tenant_id)
|
|
|
|
|
|
|
|
.first()
|
|
|
|
|
|
|
|
)
|
|
|
|
return result if isinstance(result, ChildChunk) else None
|
|
|
|
return result if isinstance(result, ChildChunk) else None
|
|
|
|
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
@classmethod
|
|
|
|
@ -2184,7 +2220,7 @@ class SegmentService:
|
|
|
|
limit: int = 20,
|
|
|
|
limit: int = 20,
|
|
|
|
):
|
|
|
|
):
|
|
|
|
"""Get segments for a document with optional filtering."""
|
|
|
|
"""Get segments for a document with optional filtering."""
|
|
|
|
query = DocumentSegment.query.filter(
|
|
|
|
query = db.session.query(DocumentSegment).filter(
|
|
|
|
DocumentSegment.document_id == document_id, DocumentSegment.tenant_id == tenant_id
|
|
|
|
DocumentSegment.document_id == document_id, DocumentSegment.tenant_id == tenant_id
|
|
|
|
)
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
@ -2236,9 +2272,11 @@ class SegmentService:
|
|
|
|
raise ValueError(ex.description)
|
|
|
|
raise ValueError(ex.description)
|
|
|
|
|
|
|
|
|
|
|
|
# check segment
|
|
|
|
# check segment
|
|
|
|
segment = DocumentSegment.query.filter(
|
|
|
|
segment = (
|
|
|
|
DocumentSegment.id == segment_id, DocumentSegment.tenant_id == user_id
|
|
|
|
db.session.query(DocumentSegment)
|
|
|
|
).first()
|
|
|
|
.filter(DocumentSegment.id == segment_id, DocumentSegment.tenant_id == user_id)
|
|
|
|
|
|
|
|
.first()
|
|
|
|
|
|
|
|
)
|
|
|
|
if not segment:
|
|
|
|
if not segment:
|
|
|
|
raise NotFound("Segment not found.")
|
|
|
|
raise NotFound("Segment not found.")
|
|
|
|
|
|
|
|
|
|
|
|
@ -2251,9 +2289,11 @@ class SegmentService:
|
|
|
|
@classmethod
|
|
|
|
@classmethod
|
|
|
|
def get_segment_by_id(cls, segment_id: str, tenant_id: str) -> Optional[DocumentSegment]:
|
|
|
|
def get_segment_by_id(cls, segment_id: str, tenant_id: str) -> Optional[DocumentSegment]:
|
|
|
|
"""Get a segment by its ID."""
|
|
|
|
"""Get a segment by its ID."""
|
|
|
|
result = DocumentSegment.query.filter(
|
|
|
|
result = (
|
|
|
|
DocumentSegment.id == segment_id, DocumentSegment.tenant_id == tenant_id
|
|
|
|
db.session.query(DocumentSegment)
|
|
|
|
).first()
|
|
|
|
.filter(DocumentSegment.id == segment_id, DocumentSegment.tenant_id == tenant_id)
|
|
|
|
|
|
|
|
.first()
|
|
|
|
|
|
|
|
)
|
|
|
|
return result if isinstance(result, DocumentSegment) else None
|
|
|
|
return result if isinstance(result, DocumentSegment) else None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
hjlarry
1 year ago