added lock
This commit is contained in:
@@ -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")
|
||||||
|
Reference in New Issue
Block a user