gmakstutis/LightRAG
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

added lock

Browse source
This commit is contained in:
Yannick Stephan 2025-02-20 00:19:47 +01:00
parent 80a61d7e7a
commit fa99d72269
1 changed files with 7 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
lightrag/lightrag.py
View file

@ -7,7 +7,7 @@ from dataclasses import asdict, dataclass, field
from datetime import datetime from datetime import datetime
from functools import partial from functools import partial
from typing import Any, AsyncIterator, Callable, Iterator, cast from typing import Any, AsyncIterator, Callable, Iterator, cast
from asyncio import Lock
from .base import ( from .base import (
BaseGraphStorage, BaseGraphStorage,
BaseKVStorage, BaseKVStorage,
@ -357,6 +357,9 @@ class LightRAG:
convert_response_to_json convert_response_to_json
) )
# Lock for entity extraction
_entity_lock = Lock()
# Custom Chunking Function # Custom Chunking Function
chunking_func: Callable[ chunking_func: Callable[
[ [
@ -823,7 +826,6 @@ class LightRAG:
self.tiktoken_model_name, self.tiktoken_model_name,
) )
} }
# Process document (text chunks and full docs) in parallel # Process document (text chunks and full docs) in parallel
tasks = [ tasks = [
self.doc_status.upsert( self.doc_status.upsert(
@ -896,8 +898,9 @@ class LightRAG:
if new_kg is None: if new_kg is None:
logger.info("No new entities or relationships extracted.") logger.info("No new entities or relationships extracted.")
else: else:
logger.info("New entities or relationships extracted.") async with self._entity_lock:
self.chunk_entity_relation_graph = new_kg logger.info("New entities or relationships extracted.")
self.chunk_entity_relation_graph = new_kg
except Exception as e: except Exception as e:
logger.error("Failed to extract entities and relationships") logger.error("Failed to extract entities and relationships")