fix dataset segment update api not effect issue (#4151)

This commit is contained in:
Jyong 2024-05-07 17:47:20 +08:00 committed by GitHub
parent 543a00e597
commit 591b993685
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -43,6 +43,7 @@ from services.vector_service import VectorService
from tasks.clean_notion_document_task import clean_notion_document_task
from tasks.deal_dataset_vector_index_task import deal_dataset_vector_index_task
from tasks.delete_segment_from_index_task import delete_segment_from_index_task
from tasks.disable_segment_from_index_task import disable_segment_from_index_task
from tasks.document_indexing_task import document_indexing_task
from tasks.document_indexing_update_task import document_indexing_update_task
from tasks.duplicate_document_indexing_task import duplicate_document_indexing_task
@ -1241,6 +1242,25 @@ class SegmentService:
cache_result = redis_client.get(indexing_cache_key)
if cache_result is not None:
raise ValueError("Segment is indexing, please try again later")
if 'enabled' in args and args['enabled'] is not None:
action = args['enabled']
if segment.enabled != action:
if not action:
segment.enabled = action
segment.disabled_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
segment.disabled_by = current_user.id
db.session.add(segment)
db.session.commit()
# Set cache to prevent indexing the same segment multiple times
redis_client.setex(indexing_cache_key, 600, 1)
disable_segment_from_index_task.delay(segment.id)
return segment
if not segment.enabled:
if 'enabled' in args and args['enabled'] is not None:
if not args['enabled']:
raise ValueError("Can't update disabled segment")
else:
raise ValueError("Can't update disabled segment")
try:
content = args['content']
if segment.content == content:
@ -1248,8 +1268,9 @@ class SegmentService:
segment.answer = args['answer']
if 'keywords' in args and args['keywords']:
segment.keywords = args['keywords']
if 'enabled' in args and args['enabled'] is not None:
segment.enabled = args['enabled']
segment.enabled = True
segment.disabled_at = None
segment.disabled_by = None
db.session.add(segment)
db.session.commit()
# update segment index task
@ -1294,12 +1315,16 @@ class SegmentService:
segment.completed_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
segment.updated_by = current_user.id
segment.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
segment.enabled = True
segment.disabled_at = None
segment.disabled_by = None
if document.doc_form == 'qa_model':
segment.answer = args['answer']
db.session.add(segment)
db.session.commit()
# update segment vector index
VectorService.update_segment_vector(args['keywords'], segment, dataset)
except Exception as e:
logging.exception("update segment index failed")
segment.enabled = False