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

feat: neo4j custom timeouts & fix: neo4j max_pool_size

Browse Source
This commit is contained in:
ultrageopro
2025-02-17 12:58:04 +03:00
parent 9a1603f594
commit d2f1e961ec
1 changed files with 15 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
lightrag/kg/neo4j_impl.py
View File

@@ -69,12 +69,24 @@ class Neo4JStorage(BaseGraphStorage):
config.get("neo4j", "connection_pool_size", fallback=800), config.get("neo4j", "connection_pool_size", fallback=800),
) )
) )
CONNECTION_TIMEOUT = float(
os.environ.get("NEO4J_CONNECTION_TIMEOUT"),
config.get("neo4j", "connection_timeout", fallback=60.0),
)
CONNECTION_ACQUISITION_TIMEOUT = float(
os.environ.get("NEO4J_CONNECTION_ACQUISITION_TIMEOUT"),
config.get("neo4j", "connection_acquisition_timeout", fallback=60.0),
)
DATABASE = os.environ.get( DATABASE = os.environ.get(
"NEO4J_DATABASE", re.sub(r"[^a-zA-Z0-9-]", "-", namespace) "NEO4J_DATABASE", re.sub(r"[^a-zA-Z0-9-]", "-", namespace)
) )
self._driver: AsyncDriver = AsyncGraphDatabase.driver( self._driver: AsyncDriver = AsyncGraphDatabase.driver(
URI, auth=(USERNAME, PASSWORD) URI,
auth=(USERNAME, PASSWORD),
max_connection_pool_size=MAX_CONNECTION_POOL_SIZE,
connection_timeout=CONNECTION_TIMEOUT,
connection_acquisition_timeout=CONNECTION_ACQUISITION_TIMEOUT,
) )
# Try to connect to the database # Try to connect to the database
@@ -82,6 +94,8 @@ class Neo4JStorage(BaseGraphStorage):
URI, URI,
auth=(USERNAME, PASSWORD), auth=(USERNAME, PASSWORD),
max_connection_pool_size=MAX_CONNECTION_POOL_SIZE, max_connection_pool_size=MAX_CONNECTION_POOL_SIZE,
connection_timeout=CONNECTION_TIMEOUT,
connection_acquisition_timeout=CONNECTION_ACQUISITION_TIMEOUT,
) as _sync_driver: ) as _sync_driver:
for database in (DATABASE, None): for database in (DATABASE, None):
self._DATABASE = database self._DATABASE = database