diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index cad1826d..0554ab76 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -845,7 +845,8 @@ class LightRAG: ) } # Process document (text chunks and full docs) in parallel - tasks = [ + # Create tasks with references for potential cancellation + doc_status_task = asyncio.create_task( self.doc_status.upsert( { doc_id: { @@ -857,13 +858,28 @@ class LightRAG: "created_at": status_doc.created_at, } } - ), - self.chunks_vdb.upsert(chunks), - self._process_entity_relation_graph(chunks), + ) + ) + chunks_vdb_task = asyncio.create_task( + self.chunks_vdb.upsert(chunks) + ) + entity_relation_task = asyncio.create_task( + self._process_entity_relation_graph(chunks) + ) + full_docs_task = asyncio.create_task( self.full_docs.upsert( {doc_id: {"content": status_doc.content}} - ), - self.text_chunks.upsert(chunks), + ) + ) + text_chunks_task = asyncio.create_task( + self.text_chunks.upsert(chunks) + ) + tasks = [ + doc_status_task, + chunks_vdb_task, + entity_relation_task, + full_docs_task, + text_chunks_task, ] try: await asyncio.gather(*tasks) @@ -881,9 +897,25 @@ class LightRAG: } ) except Exception as e: - logger.error( + # Log error and update pipeline status + error_msg = ( f"Failed to process document {doc_id}: {str(e)}" ) + logger.error(error_msg) + pipeline_status["latest_message"] = error_msg + pipeline_status["history_messages"].append(error_msg) + + # Cancel other tasks as they are no longer meaningful + for task in [ + chunks_vdb_task, + entity_relation_task, + full_docs_task, + text_chunks_task, + ]: + if not task.done(): + task.cancel() + + # Update document status to failed await self.doc_status.upsert( { doc_id: { @@ -926,7 +958,7 @@ class LightRAG: pipeline_status["latest_message"] = log_message pipeline_status["history_messages"].append(log_message) - # 获取新的待处理文档 + # Check for pending documents again processing_docs, failed_docs, pending_docs = await asyncio.gather( self.doc_status.get_docs_by_status(DocStatus.PROCESSING), self.doc_status.get_docs_by_status(DocStatus.FAILED),