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

make it smalled

Browse Source
This commit is contained in:
Yannick Stephan
2025-02-19 23:45:51 +01:00
parent 98d005dc1c
commit 4875283d9f
1 changed files with 49 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

138
lightrag/lightrag.py
View File

@@ -768,14 +768,17 @@ class LightRAG:
4. Update the document status 4. Update the document status
""" """
# 1. Get all pending, failed, and abnormally terminated processing documents. # 1. Get all pending, failed, and abnormally terminated processing documents.
to_process_docs: dict[str, DocProcessingStatus] = {} # Run the asynchronous status retrievals in parallel using asyncio.gather
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),
self.doc_status.get_docs_by_status(DocStatus.PENDING),
)
processing_docs = await self.doc_status.get_docs_by_status(DocStatus.PROCESSING) to_process_docs: dict[str, DocProcessingStatus] = {}
to_process_docs.update(processing_docs) to_process_docs.update(processing_docs)
failed_docs = await self.doc_status.get_docs_by_status(DocStatus.FAILED)
to_process_docs.update(failed_docs) to_process_docs.update(failed_docs)
pendings_docs = await self.doc_status.get_docs_by_status(DocStatus.PENDING) to_process_docs.update(pending_docs)
to_process_docs.update(pendings_docs)
if not to_process_docs: if not to_process_docs:
logger.info("All documents have been processed or are duplicates") logger.info("All documents have been processed or are duplicates")
@@ -789,10 +792,11 @@ class LightRAG:
] ]
logger.info(f"Number of batches to process: {len(docs_batches)}.") logger.info(f"Number of batches to process: {len(docs_batches)}.")
tasks: list[tuple[str, DocProcessingStatus, dict[str, Any], Any]] = []
# 3. iterate over batches # 3. iterate over batches
for batch_idx, docs_batch in enumerate(docs_batches): for batch_idx, docs_batch in enumerate(docs_batches):
logger.info(
f"Start processing batch {batch_idx + 1} of {len(docs_batches)}."
)
# 4. iterate over batch # 4. iterate over batch
for doc_id_processing_status in docs_batch: for doc_id_processing_status in docs_batch:
doc_id, status_doc = doc_id_processing_status doc_id, status_doc = doc_id_processing_status
@@ -826,91 +830,47 @@ class LightRAG:
) )
} }
# Prepare async tasks with full context # Process document (text chunks and full docs) in parallel
tasks.extend( tasks = [
[ self.chunks_vdb.upsert(chunks),
( self._process_entity_relation_graph(chunks),
doc_status_id, self.full_docs.upsert({doc_id: {"content": status_doc.content}}),
status_doc, self.text_chunks.upsert(chunks),
chunks, self.doc_status.upsert(
self.chunks_vdb.upsert(chunks), {
), doc_status_id: {
( "status": DocStatus.PROCESSED,
doc_status_id, "chunks_count": len(chunks),
status_doc, "content": status_doc.content,
chunks, "content_summary": status_doc.content_summary,
self._process_entity_relation_graph(chunks), "content_length": status_doc.content_length,
), "created_at": status_doc.created_at,
( "updated_at": datetime.now().isoformat(),
doc_status_id, }
status_doc,
chunks,
self.full_docs.upsert(
{doc_id: {"content": status_doc.content}}
),
),
(
doc_status_id,
status_doc,
chunks,
self.text_chunks.upsert(chunks),
),
]
)
# Execute tasks as they complete
for future in asyncio.as_completed([task for _, _, _, task in tasks]):
try:
# Wait for the completed task
await future
# Retrieve the full context of the completed task
completed_doc_status_id, status_doc, chunks, _ = next(
(doc_id, s_doc, ch, task)
for doc_id, s_doc, ch, task in tasks
if task == future
)
# Update status to processed
await self.doc_status.upsert(
{
completed_doc_status_id: {
"status": DocStatus.PROCESSED,
"chunks_count": len(chunks),
"content": status_doc.content,
"content_summary": status_doc.content_summary,
"content_length": status_doc.content_length,
"created_at": status_doc.created_at,
"updated_at": datetime.now().isoformat(),
} }
} ),
) ]
logger.info(f"Completed doc_id: {completed_doc_status_id}") try:
except Exception as e: await asyncio.gather(*tasks)
# Retrieve the context of the failed task await self._insert_done()
failed_doc_status_id, status_doc, chunks, _ = next(
(doc_id, s_doc, ch, task)
for doc_id, s_doc, ch, task in tasks
if task == future
)
logger.error(
f"Failed to process document {failed_doc_status_id}: {str(e)}"
)
await self.doc_status.upsert( except Exception as e:
{ logger.error(f"Failed to process document {doc_id}: {str(e)}")
failed_doc_status_id: { await self.doc_status.upsert(
"status": DocStatus.FAILED, {
"error": str(e), doc_status_id: {
"content": status_doc.content, "status": DocStatus.FAILED,
"content_summary": status_doc.content_summary, "error": str(e),
"content_length": status_doc.content_length, "content": status_doc.content,
"created_at": status_doc.created_at, "content_summary": status_doc.content_summary,
"updated_at": datetime.now().isoformat(), "content_length": status_doc.content_length,
"created_at": status_doc.created_at,
"updated_at": datetime.now().isoformat(),
}
} }
} )
) continue
await self._insert_done() logger.info(f"Completed batch {batch_idx + 1} of {len(docs_batches)}.")
async def _process_entity_relation_graph(self, chunk: dict[str, Any]) -> None: async def _process_entity_relation_graph(self, chunk: dict[str, Any]) -> None:
try: try: