From 3ed7f11b4c0c60ecf86f01cecbf1abc8a6545346 Mon Sep 17 00:00:00 2001 From: yangdx Date: Wed, 2 Apr 2025 16:36:02 +0800 Subject: [PATCH] Add index on entity_id for Neo4j --- lightrag/kg/neo4j_impl.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/lightrag/kg/neo4j_impl.py b/lightrag/kg/neo4j_impl.py index c8b75ebc..66642256 100644 --- a/lightrag/kg/neo4j_impl.py +++ b/lightrag/kg/neo4j_impl.py @@ -151,6 +151,37 @@ class Neo4JStorage(BaseGraphStorage): raise e if connected: + # Create index for base nodes on entity_id if it doesn't exist + try: + async with self._driver.session(database=database) as session: + # Check if index exists first + check_query = """ + CALL db.indexes() YIELD name, labelsOrTypes, properties + WHERE labelsOrTypes = ['base'] AND properties = ['entity_id'] + RETURN count(*) > 0 AS exists + """ + try: + check_result = await session.run(check_query) + record = await check_result.single() + await check_result.consume() + + index_exists = record and record.get("exists", False) + + if not index_exists: + # Create index only if it doesn't exist + result = await session.run( + "CREATE INDEX FOR (n:base) ON (n.entity_id)" + ) + await result.consume() + logger.info(f"Created index for base nodes on entity_id in {database}") + except Exception: + # Fallback if db.indexes() is not supported in this Neo4j version + result = await session.run( + "CREATE INDEX IF NOT EXISTS FOR (n:base) ON (n.entity_id)" + ) + await result.consume() + except Exception as e: + logger.warning(f"Failed to create index: {str(e)}") break async def finalize(self):