gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improved task handling and error management in LightRAG

Browse Source 
- Created tasks with references for cancellation
- Prioritized entity relation task execution
- Implemented task cancellation on failure
- Added error logging to pipeline_status
This commit is contained in:
yangdx
2025-03-04 14:59:50 +08:00
parent 25cbb9bf92
commit 8cf9f04dbc
1 changed files with 27 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
lightrag/lightrag.py
View File

@@ -845,7 +845,8 @@ class LightRAG:
) )
} }
# Process document (text chunks and full docs) in parallel # 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( self.doc_status.upsert(
{ {
doc_id: { doc_id: {
@@ -857,15 +858,20 @@ class LightRAG:
"created_at": status_doc.created_at, "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))
self.full_docs.upsert( entity_relation_task = asyncio.create_task(self._process_entity_relation_graph(chunks))
{doc_id: {"content": status_doc.content}} 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: try:
# Wait for entity_relation_task first as it's critical
await entity_relation_task
# If successful, wait for other tasks
await asyncio.gather(*tasks) await asyncio.gather(*tasks)
await self.doc_status.upsert( await self.doc_status.upsert(
{ {
@@ -881,9 +887,18 @@ class LightRAG:
} }
) )
except Exception as e: except Exception as e:
logger.error( # Log error and update pipeline status
f"Failed to process document {doc_id}: {str(e)}" 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, full_docs_task, text_chunks_task]:
if not task.done():
task.cancel()
# Update document status to failed
await self.doc_status.upsert( await self.doc_status.upsert(
{ {
doc_id: { doc_id: {