gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: deadlock retry adapt for class

Browse source
This commit is contained in:
Boris Arzentar 2025-04-24 17:30:25 +02:00
parent 36507502ef
commit 2a853e604e
4 changed files with 9 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cognee/infrastructure/databases/graph/neo4j_driver/adapter.py
View file

@ -51,7 +51,7 @@ class Neo4jAdapter(GraphDBInterface):
async with self.driver.session() as session: async with self.driver.session() as session:
yield session yield session
@deadlock_retry @deadlock_retry()
async def query( async def query(
self, self,
query: str, query: str,

4
cognee/infrastructure/databases/graph/neo4j_driver/deadlock_retry.py
View file

@ -27,7 +27,7 @@ def deadlock_retry(max_retries=5):
def decorator(func): def decorator(func):
@wraps(func) @wraps(func)
async def wrapper(*args, **kwargs): async def wrapper(self, *args, **kwargs):
from neo4j.exceptions import Neo4jError, DatabaseUnavailable from neo4j.exceptions import Neo4jError, DatabaseUnavailable
attempt = 0 attempt = 0
@ -43,7 +43,7 @@ def deadlock_retry(max_retries=5):
while attempt <= max_retries: while attempt <= max_retries:
try: try:
attempt += 1 attempt += 1
return await func(*args, **kwargs) return await func(self, *args, **kwargs)
except Neo4jError as error: except Neo4jError as error:
if attempt > max_retries: if attempt > max_retries:
raise # Re-raise the original error raise # Re-raise the original error

6
cognee/infrastructure/llm/config.py
View file

@ -23,9 +23,9 @@ class LLMConfig(BaseSettings):
embedding_rate_limit_requests: int = 60 embedding_rate_limit_requests: int = 60
embedding_rate_limit_interval: int = 60 # in seconds (default is 60 requests per minute) embedding_rate_limit_interval: int = 60 # in seconds (default is 60 requests per minute)
fallback_api_key: str = None fallback_api_key: str = ""
fallback_endpoint: str = None fallback_endpoint: str = ""
fallback_model: str = None fallback_model: str = ""
model_config = SettingsConfigDict(env_file=".env", extra="allow") model_config = SettingsConfigDict(env_file=".env", extra="allow")

6
cognee/tests/unit/infrastructure/databases/graph/neo4j_deadlock_test.py
View file

@ -15,7 +15,7 @@ async def test_deadlock_retry_errored():
wrapped_function = deadlock_retry(max_retries=1)(mock_function) wrapped_function = deadlock_retry(max_retries=1)(mock_function)
with pytest.raises(Neo4jError): with pytest.raises(Neo4jError):
await wrapped_function() await wrapped_function(self=None)
async def test_deadlock_retry(): async def test_deadlock_retry():
@ -25,7 +25,7 @@ async def test_deadlock_retry():
wrapped_function = deadlock_retry(max_retries=2)(mock_function) 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" 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) 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" assert result == True, "Function should have succeded on second time"