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:
return JSONResponse(
status_code=503,
content={"status": "not ready", "reason": "critical services unhealthy"}
)
return JSONResponse(
status_code=200,
content={"status": "ready"}
content={"status": "not ready", "reason": "critical services unhealthy"},
)
return JSONResponse(status_code=200, content={"status": "ready"})
except Exception as e:
return JSONResponse(
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:
status_code = 200 # Degraded is still operational
return JSONResponse(
status_code=status_code,
content=health_status.model_dump()
)
return JSONResponse(status_code=status_code, content=health_status.model_dump())
except Exception as e:
return JSONResponse(
status_code=503,
content={
"status": "unhealthy",
"error": f"Health check system failure: {str(e)}"
}
content={"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."""
start_time = time.time()
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
config = get_relational_config()
@ -58,7 +60,7 @@ class HealthChecker:
status=HealthStatus.HEALTHY,
provider=config.db_provider,
response_time_ms=response_time,
details="Connection successful"
details="Connection successful",
)
except Exception as e:
response_time = int((time.time() - start_time) * 1000)
@ -66,7 +68,7 @@ class HealthChecker:
status=HealthStatus.UNHEALTHY,
provider="unknown",
response_time_ms=response_time,
details=f"Connection failed: {str(e)}"
details=f"Connection failed: {str(e)}",
)
async def check_vector_db(self) -> ComponentHealth:
@ -80,9 +82,9 @@ class HealthChecker:
engine = get_vector_engine()
# Test basic operation - just check if engine is accessible
if hasattr(engine, 'health_check'):
if hasattr(engine, "health_check"):
await engine.health_check()
elif hasattr(engine, 'list_tables'):
elif hasattr(engine, "list_tables"):
# For LanceDB and similar
engine.list_tables()
@ -91,7 +93,7 @@ class HealthChecker:
status=HealthStatus.HEALTHY,
provider=config.vector_db_provider,
response_time_ms=response_time,
details="Index accessible"
details="Index accessible",
)
except Exception as e:
response_time = int((time.time() - start_time) * 1000)
@ -99,7 +101,7 @@ class HealthChecker:
status=HealthStatus.UNHEALTHY,
provider="unknown",
response_time_ms=response_time,
details=f"Connection failed: {str(e)}"
details=f"Connection failed: {str(e)}",
)
async def check_graph_db(self) -> ComponentHealth:
@ -113,9 +115,9 @@ class HealthChecker:
engine = await get_graph_engine()
# Test basic operation - just check if engine is accessible
if hasattr(engine, 'health_check'):
if hasattr(engine, "health_check"):
await engine.health_check()
elif hasattr(engine, 'get_nodes'):
elif hasattr(engine, "get_nodes"):
# Basic connectivity test
pass
@ -124,7 +126,7 @@ class HealthChecker:
status=HealthStatus.HEALTHY,
provider=config.graph_database_provider,
response_time_ms=response_time,
details="Schema validated"
details="Schema validated",
)
except Exception as e:
response_time = int((time.time() - start_time) * 1000)
@ -132,7 +134,7 @@ class HealthChecker:
status=HealthStatus.UNHEALTHY,
provider="unknown",
response_time_ms=response_time,
details=f"Connection failed: {str(e)}"
details=f"Connection failed: {str(e)}",
)
async def check_file_storage(self) -> ComponentHealth:
@ -154,7 +156,7 @@ class HealthChecker:
os.makedirs(base_config.data_root_directory, exist_ok=True)
# Simple write/read 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")
os.remove(test_file)
else:
@ -168,7 +170,7 @@ class HealthChecker:
status=HealthStatus.HEALTHY,
provider=provider,
response_time_ms=response_time,
details="Storage accessible"
details="Storage accessible",
)
except Exception as e:
response_time = int((time.time() - start_time) * 1000)
@ -176,7 +178,7 @@ class HealthChecker:
status=HealthStatus.UNHEALTHY,
provider="unknown",
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:
@ -201,7 +203,7 @@ class HealthChecker:
status=status,
provider=config.llm_provider,
response_time_ms=response_time,
details=details
details=details,
)
except Exception as e:
response_time = int((time.time() - start_time) * 1000)
@ -209,24 +211,26 @@ class HealthChecker:
status=HealthStatus.DEGRADED,
provider="unknown",
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:
"""Check embedding service health (non-critical)."""
start_time = time.time()
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
engine = get_embedding_engine()
get_embedding_engine()
response_time = int((time.time() - start_time) * 1000)
return ComponentHealth(
status=HealthStatus.HEALTHY,
provider="configured",
response_time_ms=response_time,
details="Embedding engine accessible"
details="Embedding engine accessible",
)
except Exception as e:
response_time = int((time.time() - start_time) * 1000)
@ -234,7 +238,7 @@ class HealthChecker:
status=HealthStatus.DEGRADED,
provider="unknown",
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:
@ -257,8 +261,7 @@ class HealthChecker:
# Run critical checks
critical_results = await asyncio.gather(
*[check for _, check in critical_checks],
return_exceptions=True
*[check for _, check in critical_checks], return_exceptions=True
)
for (name, _), result in zip(critical_checks, critical_results):
@ -267,7 +270,7 @@ class HealthChecker:
status=HealthStatus.UNHEALTHY,
provider="unknown",
response_time_ms=0,
details=f"Health check failed: {str(result)}"
details=f"Health check failed: {str(result)}",
)
else:
components[name] = result
@ -275,8 +278,7 @@ class HealthChecker:
# Run non-critical checks if detailed
if detailed:
non_critical_results = await asyncio.gather(
*[check for _, check in non_critical_checks],
return_exceptions=True
*[check for _, check in non_critical_checks], return_exceptions=True
)
for (name, _), result in zip(non_critical_checks, non_critical_results):
@ -285,7 +287,7 @@ class HealthChecker:
status=HealthStatus.DEGRADED,
provider="unknown",
response_time_ms=0,
details=f"Health check failed: {str(result)}"
details=f"Health check failed: {str(result)}",
)
else:
components[name] = result
@ -311,7 +313,7 @@ class HealthChecker:
timestamp=datetime.now(timezone.utc).isoformat(),
version=get_cognee_version(),
uptime=int(time.time() - self.start_time),
components=components
components=components,
)