Add session expired retry to neo4j graph (#26182)

Co-authored-by: Erick Friis <erick@langchain.dev>
This commit is contained in:
Tomaz Bratanic 2024-09-09 03:40:43 +09:00 committed by GitHub
parent b3c7ed4913
commit 181e4fc0e0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -410,7 +410,9 @@ class Neo4jGraph(GraphStore):
"""Returns the structured schema of the Graph"""
return self.structured_schema
def query(self, query: str, params: dict = {}) -> List[Dict[str, Any]]:
def query(
self, query: str, params: dict = {}, retry_on_session_expired: bool = True
) -> List[Dict[str, Any]]:
"""Query Neo4j database.
Args:
@ -421,7 +423,7 @@ class Neo4jGraph(GraphStore):
List[Dict[str, Any]]: The list of dictionaries containing the query results.
"""
from neo4j import Query
from neo4j.exceptions import CypherSyntaxError
from neo4j.exceptions import CypherSyntaxError, SessionExpired
with self._driver.session(database=self._database) as session:
try:
@ -432,6 +434,15 @@ class Neo4jGraph(GraphStore):
return json_data
except CypherSyntaxError as e:
raise ValueError(f"Generated Cypher Statement is not valid\n{e}")
except (
SessionExpired
) as e: # Session expired is a transient error that can be retried
if retry_on_session_expired:
return self.query(
query, params=params, retry_on_session_expired=False
)
else:
raise e
def refresh_schema(self) -> None:
"""