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

Feat/dify rag (#2528)

Browse Source 
Co-authored-by: jyong <jyong@dify.ai>
This commit is contained in:
Jyong
2024-02-22 23:31:57 +08:00
committed by GitHub
parent 97fe817186
commit 6c4e6bf1d6
119 changed files with 3181 additions and 5892 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
api/tasks/clean_document_task.py
View File

@@ -4,17 +4,18 @@ import time
import click
from celery import shared_task
from core.index.index import IndexBuilder
from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
from extensions.ext_database import db
from models.dataset import Dataset, DocumentSegment
@shared_task(queue='dataset')
def clean_document_task(document_id: str, dataset_id: str):
def clean_document_task(document_id: str, dataset_id: str, doc_form: str):
"""
Clean document when document deleted.
:param document_id: document id
:param dataset_id: dataset id
:param doc_form: doc_form
Usage: clean_document_task.delay(document_id, dataset_id)
"""
@@ -27,21 +28,12 @@ def clean_document_task(document_id: str, dataset_id: str):
if not dataset:
raise Exception('Document has no dataset')
vector_index = IndexBuilder.get_index(dataset, 'high_quality')
kw_index = IndexBuilder.get_index(dataset, 'economy')
segments = db.session.query(DocumentSegment).filter(DocumentSegment.document_id == document_id).all()
# check segment is exist
if segments:
index_node_ids = [segment.index_node_id for segment in segments]
# delete from vector index
if vector_index:
vector_index.delete_by_document_id(document_id)
# delete from keyword index
if index_node_ids:
kw_index.delete_by_ids(index_node_ids)
index_processor = IndexProcessorFactory(doc_form).init_index_processor()
index_processor.clean(dataset, index_node_ids)
for segment in segments:
db.session.delete(segment)