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

format using ruff

Browse source
This commit is contained in:
Pavan Chilukuri 2025-08-02 13:52:55 -05:00
parent be33be7438
commit 585ad96c47
2 changed files with 66 additions and 73 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
cognee/api/client.py
View file

@ -185,16 +185,13 @@ async def readiness_check():
if health_status.status == HealthStatus.UNHEALTHY: if health_status.status == HealthStatus.UNHEALTHY:
return JSONResponse( return JSONResponse(
status_code=503, status_code=503,
content={"status": "not ready", "reason": "critical services unhealthy"} content={"status": "not ready", "reason": "critical services unhealthy"},
) )
return JSONResponse( return JSONResponse(status_code=200, content={"status": "ready"})
status_code=200,
content={"status": "ready"}
)
except Exception as e: except Exception as e:
return JSONResponse( return JSONResponse(
status_code=503, status_code=503,
content={"status": "not ready", "reason": f"health check failed: {str(e)}"} content={"status": "not ready", "reason": f"health check failed: {str(e)}"},
) )
@ -211,17 +208,11 @@ async def detailed_health_check():
elif health_status.status == HealthStatus.DEGRADED: elif health_status.status == HealthStatus.DEGRADED:
status_code = 200 # Degraded is still operational status_code = 200 # Degraded is still operational
return JSONResponse( return JSONResponse(status_code=status_code, content=health_status.model_dump())
status_code=status_code,
content=health_status.model_dump()
)
except Exception as e: except Exception as e:
return JSONResponse( return JSONResponse(
status_code=503, status_code=503,
content={ content={"status": "unhealthy", "error": f"Health check system failure: {str(e)}"},
"status": "unhealthy",
"error": f"Health check system failure: {str(e)}"
}
) )

56
cognee/api/health.py
View file

@ -42,7 +42,9 @@ class HealthChecker:
"""Check relational database health.""" """Check relational database health."""
start_time = time.time() start_time = time.time()
try: try:
from cognee.infrastructure.databases.relational.get_relational_engine import get_relational_engine from cognee.infrastructure.databases.relational.get_relational_engine import (
get_relational_engine,
)
from cognee.infrastructure.databases.relational.config import get_relational_config from cognee.infrastructure.databases.relational.config import get_relational_config
config = get_relational_config() config = get_relational_config()
@ -58,7 +60,7 @@ class HealthChecker:
status=HealthStatus.HEALTHY, status=HealthStatus.HEALTHY,
provider=config.db_provider, provider=config.db_provider,
response_time_ms=response_time, response_time_ms=response_time,
details="Connection successful" details="Connection successful",
) )
except Exception as e: except Exception as e:
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
@ -66,7 +68,7 @@ class HealthChecker:
status=HealthStatus.UNHEALTHY, status=HealthStatus.UNHEALTHY,
provider="unknown", provider="unknown",
response_time_ms=response_time, response_time_ms=response_time,
details=f"Connection failed: {str(e)}" details=f"Connection failed: {str(e)}",
) )
async def check_vector_db(self) -> ComponentHealth: async def check_vector_db(self) -> ComponentHealth:
@ -80,9 +82,9 @@ class HealthChecker:
engine = get_vector_engine() engine = get_vector_engine()
# Test basic operation - just check if engine is accessible # Test basic operation - just check if engine is accessible
if hasattr(engine, 'health_check'): if hasattr(engine, "health_check"):
await engine.health_check() await engine.health_check()
elif hasattr(engine, 'list_tables'): elif hasattr(engine, "list_tables"):
# For LanceDB and similar # For LanceDB and similar
engine.list_tables() engine.list_tables()
@ -91,7 +93,7 @@ class HealthChecker:
status=HealthStatus.HEALTHY, status=HealthStatus.HEALTHY,
provider=config.vector_db_provider, provider=config.vector_db_provider,
response_time_ms=response_time, response_time_ms=response_time,
details="Index accessible" details="Index accessible",
) )
except Exception as e: except Exception as e:
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
@ -99,7 +101,7 @@ class HealthChecker:
status=HealthStatus.UNHEALTHY, status=HealthStatus.UNHEALTHY,
provider="unknown", provider="unknown",
response_time_ms=response_time, response_time_ms=response_time,
details=f"Connection failed: {str(e)}" details=f"Connection failed: {str(e)}",
) )
async def check_graph_db(self) -> ComponentHealth: async def check_graph_db(self) -> ComponentHealth:
@ -113,9 +115,9 @@ class HealthChecker:
engine = await get_graph_engine() engine = await get_graph_engine()
# Test basic operation - just check if engine is accessible # Test basic operation - just check if engine is accessible
if hasattr(engine, 'health_check'): if hasattr(engine, "health_check"):
await engine.health_check() await engine.health_check()
elif hasattr(engine, 'get_nodes'): elif hasattr(engine, "get_nodes"):
# Basic connectivity test # Basic connectivity test
pass pass
@ -124,7 +126,7 @@ class HealthChecker:
status=HealthStatus.HEALTHY, status=HealthStatus.HEALTHY,
provider=config.graph_database_provider, provider=config.graph_database_provider,
response_time_ms=response_time, response_time_ms=response_time,
details="Schema validated" details="Schema validated",
) )
except Exception as e: except Exception as e:
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
@ -132,7 +134,7 @@ class HealthChecker:
status=HealthStatus.UNHEALTHY, status=HealthStatus.UNHEALTHY,
provider="unknown", provider="unknown",
response_time_ms=response_time, response_time_ms=response_time,
details=f"Connection failed: {str(e)}" details=f"Connection failed: {str(e)}",
) )
async def check_file_storage(self) -> ComponentHealth: async def check_file_storage(self) -> ComponentHealth:
@ -154,7 +156,7 @@ class HealthChecker:
os.makedirs(base_config.data_root_directory, exist_ok=True) os.makedirs(base_config.data_root_directory, exist_ok=True)
# Simple write/read test # Simple write/read test
test_file = os.path.join(base_config.data_root_directory, "health_check_test") test_file = os.path.join(base_config.data_root_directory, "health_check_test")
with open(test_file, 'w') as f: with open(test_file, "w") as f:
f.write("test") f.write("test")
os.remove(test_file) os.remove(test_file)
else: else:
@ -168,7 +170,7 @@ class HealthChecker:
status=HealthStatus.HEALTHY, status=HealthStatus.HEALTHY,
provider=provider, provider=provider,
response_time_ms=response_time, response_time_ms=response_time,
details="Storage accessible" details="Storage accessible",
) )
except Exception as e: except Exception as e:
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
@ -176,7 +178,7 @@ class HealthChecker:
status=HealthStatus.UNHEALTHY, status=HealthStatus.UNHEALTHY,
provider="unknown", provider="unknown",
response_time_ms=response_time, response_time_ms=response_time,
details=f"Storage test failed: {str(e)}" details=f"Storage test failed: {str(e)}",
) )
async def check_llm_provider(self) -> ComponentHealth: async def check_llm_provider(self) -> ComponentHealth:
@ -201,7 +203,7 @@ class HealthChecker:
status=status, status=status,
provider=config.llm_provider, provider=config.llm_provider,
response_time_ms=response_time, response_time_ms=response_time,
details=details details=details,
) )
except Exception as e: except Exception as e:
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
@ -209,24 +211,26 @@ class HealthChecker:
status=HealthStatus.DEGRADED, status=HealthStatus.DEGRADED,
provider="unknown", provider="unknown",
response_time_ms=response_time, response_time_ms=response_time,
details=f"Config check failed: {str(e)}" details=f"Config check failed: {str(e)}",
) )
async def check_embedding_service(self) -> ComponentHealth: async def check_embedding_service(self) -> ComponentHealth:
"""Check embedding service health (non-critical).""" """Check embedding service health (non-critical)."""
start_time = time.time() start_time = time.time()
try: try:
from cognee.infrastructure.databases.vector.embeddings.get_embedding_engine import get_embedding_engine from cognee.infrastructure.databases.vector.embeddings.get_embedding_engine import (
get_embedding_engine,
)
# Just check if we can get the engine without calling it # Just check if we can get the engine without calling it
engine = get_embedding_engine() get_embedding_engine()
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
return ComponentHealth( return ComponentHealth(
status=HealthStatus.HEALTHY, status=HealthStatus.HEALTHY,
provider="configured", provider="configured",
response_time_ms=response_time, response_time_ms=response_time,
details="Embedding engine accessible" details="Embedding engine accessible",
) )
except Exception as e: except Exception as e:
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
@ -234,7 +238,7 @@ class HealthChecker:
status=HealthStatus.DEGRADED, status=HealthStatus.DEGRADED,
provider="unknown", provider="unknown",
response_time_ms=response_time, response_time_ms=response_time,
details=f"Embedding engine failed: {str(e)}" details=f"Embedding engine failed: {str(e)}",
) )
async def get_health_status(self, detailed: bool = False) -> HealthResponse: async def get_health_status(self, detailed: bool = False) -> HealthResponse:
@ -257,8 +261,7 @@ class HealthChecker:
# Run critical checks # Run critical checks
critical_results = await asyncio.gather( critical_results = await asyncio.gather(
*[check for _, check in critical_checks], *[check for _, check in critical_checks], return_exceptions=True
return_exceptions=True
) )
for (name, _), result in zip(critical_checks, critical_results): for (name, _), result in zip(critical_checks, critical_results):
@ -267,7 +270,7 @@ class HealthChecker:
status=HealthStatus.UNHEALTHY, status=HealthStatus.UNHEALTHY,
provider="unknown", provider="unknown",
response_time_ms=0, response_time_ms=0,
details=f"Health check failed: {str(result)}" details=f"Health check failed: {str(result)}",
) )
else: else:
components[name] = result components[name] = result
@ -275,8 +278,7 @@ class HealthChecker:
# Run non-critical checks if detailed # Run non-critical checks if detailed
if detailed: if detailed:
non_critical_results = await asyncio.gather( non_critical_results = await asyncio.gather(
*[check for _, check in non_critical_checks], *[check for _, check in non_critical_checks], return_exceptions=True
return_exceptions=True
) )
for (name, _), result in zip(non_critical_checks, non_critical_results): for (name, _), result in zip(non_critical_checks, non_critical_results):
@ -285,7 +287,7 @@ class HealthChecker:
status=HealthStatus.DEGRADED, status=HealthStatus.DEGRADED,
provider="unknown", provider="unknown",
response_time_ms=0, response_time_ms=0,
details=f"Health check failed: {str(result)}" details=f"Health check failed: {str(result)}",
) )
else: else:
components[name] = result components[name] = result
@ -311,7 +313,7 @@ class HealthChecker:
timestamp=datetime.now(timezone.utc).isoformat(), timestamp=datetime.now(timezone.utc).isoformat(),
version=get_cognee_version(), version=get_cognee_version(),
uptime=int(time.time() - self.start_time), uptime=int(time.time() - self.start_time),
components=components components=components,
) )