From fa99d72269dea53324472d72aff3c3ace1c0e805 Mon Sep 17 00:00:00 2001 From: Yannick Stephan Date: Thu, 20 Feb 2025 00:19:47 +0100 Subject: [PATCH] added lock --- lightrag/lightrag.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index 156bed49..22686ab2 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -7,7 +7,7 @@ from dataclasses import asdict, dataclass, field from datetime import datetime from functools import partial from typing import Any, AsyncIterator, Callable, Iterator, cast - +from asyncio import Lock from .base import ( BaseGraphStorage, BaseKVStorage, @@ -357,6 +357,9 @@ class LightRAG: convert_response_to_json ) + # Lock for entity extraction + _entity_lock = Lock() + # Custom Chunking Function chunking_func: Callable[ [ @@ -823,7 +826,6 @@ class LightRAG: self.tiktoken_model_name, ) } - # Process document (text chunks and full docs) in parallel tasks = [ self.doc_status.upsert( @@ -896,8 +898,9 @@ class LightRAG: if new_kg is None: logger.info("No new entities or relationships extracted.") else: - logger.info("New entities or relationships extracted.") - self.chunk_entity_relation_graph = new_kg + async with self._entity_lock: + logger.info("New entities or relationships extracted.") + self.chunk_entity_relation_graph = new_kg except Exception as e: logger.error("Failed to extract entities and relationships")