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

refactor code

Browse source
This commit is contained in:
Pavan Chilukuri 2025-08-12 10:45:43 -05:00
parent b2d83c47cd
commit a19ba4d88e
1 changed files with 19 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
cognee/api/health.py
View file

@ -51,7 +51,7 @@ class HealthChecker:
engine = get_relational_engine() engine = get_relational_engine()
# Test connection by creating a session # Test connection by creating a session
session = await engine.get_session() session = engine.get_session()
if session: if session:
await session.close() await session.close()
@ -64,6 +64,7 @@ class HealthChecker:
) )
except Exception as e: except Exception as e:
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
logger.error(f"Relational DB health check failed: {str(e)}", exc_info=True)
return ComponentHealth( return ComponentHealth(
status=HealthStatus.UNHEALTHY, status=HealthStatus.UNHEALTHY,
provider="unknown", provider="unknown",
@ -97,6 +98,7 @@ class HealthChecker:
) )
except Exception as e: except Exception as e:
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
logger.error(f"Vector DB health check failed: {str(e)}", exc_info=True)
return ComponentHealth( return ComponentHealth(
status=HealthStatus.UNHEALTHY, status=HealthStatus.UNHEALTHY,
provider="unknown", provider="unknown",
@ -118,9 +120,9 @@ class HealthChecker:
if hasattr(engine, "execute"): if hasattr(engine, "execute"):
# For SQL-like graph DBs (Neo4j, Memgraph) # For SQL-like graph DBs (Neo4j, Memgraph)
await engine.execute("MATCH () RETURN count(*) LIMIT 1") await engine.execute("MATCH () RETURN count(*) LIMIT 1")
elif hasattr(engine, "get_nodes"): elif hasattr(engine, "query"):
# For other graph engines - try to get nodes # For other graph engines
list(engine.get_nodes(limit=1)) engine.query("MATCH () RETURN count(*) LIMIT 1", {})
# If engine exists but no test method, consider it healthy # If engine exists but no test method, consider it healthy
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
@ -132,6 +134,7 @@ class HealthChecker:
) )
except Exception as e: except Exception as e:
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
logger.error(f"Graph DB health check failed: {str(e)}", exc_info=True)
return ComponentHealth( return ComponentHealth(
status=HealthStatus.UNHEALTHY, status=HealthStatus.UNHEALTHY,
provider="unknown", provider="unknown",
@ -194,7 +197,7 @@ class HealthChecker:
# Test actual API connection with minimal request # Test actual API connection with minimal request
client = get_llm_client() client = get_llm_client()
await client.acomplete("test", max_tokens=1) await client.show_prompt("test", "test")
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
return ComponentHealth( return ComponentHealth(
@ -205,6 +208,7 @@ class HealthChecker:
) )
except Exception as e: except Exception as e:
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
logger.error(f"LLM provider health check failed: {str(e)}", exc_info=True)
return ComponentHealth( return ComponentHealth(
status=HealthStatus.DEGRADED, status=HealthStatus.DEGRADED,
provider="unknown", provider="unknown",
@ -245,6 +249,15 @@ class HealthChecker:
components = {} components = {}
# Critical services # Critical services
critical_components = [
"relational_db",
"vector_db",
"graph_db",
"file_storage",
"llm_provider",
"embedding_service",
]
critical_checks = [ critical_checks = [
("relational_db", self.check_relational_db()), ("relational_db", self.check_relational_db()),
("vector_db", self.check_vector_db()), ("vector_db", self.check_vector_db()),
@ -294,15 +307,7 @@ class HealthChecker:
critical_unhealthy = any( critical_unhealthy = any(
comp.status == HealthStatus.UNHEALTHY comp.status == HealthStatus.UNHEALTHY
for name, comp in components.items() for name, comp in components.items()
if name if name in critical_components
in [
"relational_db",
"vector_db",
"graph_db",
"file_storage",
"llm_provider",
"embedding_service",
]
) )
has_degraded = any(comp.status == HealthStatus.DEGRADED for comp in components.values()) has_degraded = any(comp.status == HealthStatus.DEGRADED for comp in components.values())