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

update dataset embedding model, update document status to be indexing (#7145)

Browse Source
This commit is contained in:
Jyong 2024-08-09 16:47:15 +08:00 committed by GitHub
parent f667ef98cb
commit c6b0dc6a29
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 70 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
api/tasks/deal_dataset_vector_index_task.py
View File

@ -42,13 +42,20 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
).all()
if dataset_documents:
documents = []
dataset_documents_ids = [doc.id for doc in dataset_documents]
db.session.query(DatasetDocument).filter(DatasetDocument.id.in_(dataset_documents_ids)) \
.update({"indexing_status": "indexing"}, synchronize_session=False)
db.session.commit()
for dataset_document in dataset_documents:
# delete from vector index
try:
# add from vector index
segments = db.session.query(DocumentSegment).filter(
DocumentSegment.document_id == dataset_document.id,
DocumentSegment.enabled == True
) .order_by(DocumentSegment.position.asc()).all()
if segments:
documents = []
for segment in segments:
document = Document(
page_content=segment.content,
@ -61,12 +68,16 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
)
documents.append(document)
# save vector index
index_processor.load(dataset, documents, with_keywords=False)
db.session.query(DatasetDocument).filter(DatasetDocument.id == dataset_document.id) \
.update({"indexing_status": "completed"}, synchronize_session=False)
db.session.commit()
except Exception as e:
db.session.query(DatasetDocument).filter(DatasetDocument.id == dataset_document.id) \
.update({"indexing_status": "error", "error": str(e)}, synchronize_session=False)
db.session.commit()
elif action == 'update':
# clean index
index_processor.clean(dataset, None, with_keywords=False)
dataset_documents = db.session.query(DatasetDocument).filter(
DatasetDocument.dataset_id == dataset_id,
DatasetDocument.indexing_status == 'completed',
@ -75,13 +86,24 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
).all()
# add new index
if dataset_documents:
documents = []
# update document status
dataset_documents_ids = [doc.id for doc in dataset_documents]
db.session.query(DatasetDocument).filter(DatasetDocument.id.in_(dataset_documents_ids)) \
.update({"indexing_status": "indexing"}, synchronize_session=False)
db.session.commit()
# clean index
index_processor.clean(dataset, None, with_keywords=False)
for dataset_document in dataset_documents:
# delete from vector index
# update from vector index
try:
segments = db.session.query(DocumentSegment).filter(
DocumentSegment.document_id == dataset_document.id,
DocumentSegment.enabled == True
).order_by(DocumentSegment.position.asc()).all()
if segments:
documents = []
for segment in segments:
document = Document(
page_content=segment.content,
@ -94,9 +116,16 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
)
documents.append(document)
# save vector index
index_processor.load(dataset, documents, with_keywords=False)
db.session.query(DatasetDocument).filter(DatasetDocument.id == dataset_document.id) \
.update({"indexing_status": "completed"}, synchronize_session=False)
db.session.commit()
except Exception as e:
db.session.query(DatasetDocument).filter(DatasetDocument.id == dataset_document.id) \
.update({"indexing_status": "error", "error": str(e)}, synchronize_session=False)
db.session.commit()
end_at = time.perf_counter()
logging.info(