From e1045f01c666e92c77b3316545f2ac4cb0a2fd3c Mon Sep 17 00:00:00 2001 From: takatost Date: Mon, 12 Feb 2024 13:50:43 +0800 Subject: [PATCH] pref: optimize add hit count query performance when dataset hit (#2436) --- .../callback_handler/index_tool_callback_handler.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/api/core/callback_handler/index_tool_callback_handler.py b/api/core/callback_handler/index_tool_callback_handler.py index 7c8a3ce478..903953486a 100644 --- a/api/core/callback_handler/index_tool_callback_handler.py +++ b/api/core/callback_handler/index_tool_callback_handler.py @@ -42,12 +42,16 @@ class DatasetIndexToolCallbackHandler: def on_tool_end(self, documents: list[Document]) -> None: """Handle tool end.""" for document in documents: - doc_id = document.metadata['doc_id'] + query = db.session.query(DocumentSegment).filter( + DocumentSegment.index_node_id == document.metadata['doc_id'] + ) + + # if 'dataset_id' in document.metadata: + if 'dataset_id' in document.metadata: + query = query.filter(DocumentSegment.dataset_id == document.metadata['dataset_id']) # add hit count to document segment - db.session.query(DocumentSegment).filter( - DocumentSegment.index_node_id == doc_id - ).update( + query.update( {DocumentSegment.hit_count: DocumentSegment.hit_count + 1}, synchronize_session=False )