From 2a853e604e6296f3838166443fb2cc1e37865c33 Mon Sep 17 00:00:00 2001 From: Boris Arzentar Date: Thu, 24 Apr 2025 17:30:25 +0200 Subject: [PATCH] fix: deadlock retry adapt for class --- .../infrastructure/databases/graph/neo4j_driver/adapter.py | 2 +- .../databases/graph/neo4j_driver/deadlock_retry.py | 4 ++-- cognee/infrastructure/llm/config.py | 6 +++--- .../infrastructure/databases/graph/neo4j_deadlock_test.py | 6 +++--- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py b/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py index 4623c1951..321b68284 100644 --- a/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py +++ b/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py @@ -51,7 +51,7 @@ class Neo4jAdapter(GraphDBInterface): async with self.driver.session() as session: yield session - @deadlock_retry + @deadlock_retry() async def query( self, query: str, diff --git a/cognee/infrastructure/databases/graph/neo4j_driver/deadlock_retry.py b/cognee/infrastructure/databases/graph/neo4j_driver/deadlock_retry.py index c1591ed92..82d2a1c29 100644 --- a/cognee/infrastructure/databases/graph/neo4j_driver/deadlock_retry.py +++ b/cognee/infrastructure/databases/graph/neo4j_driver/deadlock_retry.py @@ -27,7 +27,7 @@ def deadlock_retry(max_retries=5): def decorator(func): @wraps(func) - async def wrapper(*args, **kwargs): + async def wrapper(self, *args, **kwargs): from neo4j.exceptions import Neo4jError, DatabaseUnavailable attempt = 0 @@ -43,7 +43,7 @@ def deadlock_retry(max_retries=5): while attempt <= max_retries: try: attempt += 1 - return await func(*args, **kwargs) + return await func(self, *args, **kwargs) except Neo4jError as error: if attempt > max_retries: raise # Re-raise the original error diff --git a/cognee/infrastructure/llm/config.py b/cognee/infrastructure/llm/config.py index 155179570..2942c4793 100644 --- a/cognee/infrastructure/llm/config.py +++ b/cognee/infrastructure/llm/config.py @@ -23,9 +23,9 @@ class LLMConfig(BaseSettings): embedding_rate_limit_requests: int = 60 embedding_rate_limit_interval: int = 60 # in seconds (default is 60 requests per minute) - fallback_api_key: str = None - fallback_endpoint: str = None - fallback_model: str = None + fallback_api_key: str = "" + fallback_endpoint: str = "" + fallback_model: str = "" model_config = SettingsConfigDict(env_file=".env", extra="allow") diff --git a/cognee/tests/unit/infrastructure/databases/graph/neo4j_deadlock_test.py b/cognee/tests/unit/infrastructure/databases/graph/neo4j_deadlock_test.py index f73b7b60b..401189080 100644 --- a/cognee/tests/unit/infrastructure/databases/graph/neo4j_deadlock_test.py +++ b/cognee/tests/unit/infrastructure/databases/graph/neo4j_deadlock_test.py @@ -15,7 +15,7 @@ async def test_deadlock_retry_errored(): wrapped_function = deadlock_retry(max_retries=1)(mock_function) with pytest.raises(Neo4jError): - await wrapped_function() + await wrapped_function(self=None) async def test_deadlock_retry(): @@ -25,7 +25,7 @@ async def test_deadlock_retry(): wrapped_function = deadlock_retry(max_retries=2)(mock_function) - result = await wrapped_function() + result = await wrapped_function(self=None) assert result == True, "Function should have succeded on second time" @@ -38,7 +38,7 @@ async def test_deadlock_retry_exhaustive(): wrapped_function = deadlock_retry(max_retries=2)(mock_function) - result = await wrapped_function() + result = await wrapped_function(self=None) assert result == True, "Function should have succeded on second time"