Repository: topoteretes/cognee
File: cognee/api/health.py (Python)

Code before:
class HealthResponse(BaseModel):
    status: HealthStatus
    timestamp: str
    version: str
    uptime: int
    components: Dict[str, ComponentHealth]


class HealthChecker:
    def __init__(self):
        self.start_time = time.time()

    async def check_relational_db(self) -> ComponentHealth:
        """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.config import get_relational_config

            config = get_relational_config()
            engine = get_relational_engine()

            # Test connection by creating a session
            session = await engine.get_session()
            if session:
                await session.close()

            response_time = int((time.time() - start_time) * 1000)
            return ComponentHealth(
                status=HealthStatus.HEALTHY,
                provider=config.db_provider,
                response_time_ms=response_time,
                details="Connection successful",
            )
        except Exception as e:
            response_time = int((time.time() - start_time) * 1000)
            return ComponentHealth(
                status=HealthStatus.UNHEALTHY,
                provider="unknown",
                response_time_ms=response_time,
                details=f"Connection failed: {str(e)}",
            )

    async def check_vector_db(self) -> ComponentHealth:
        """Check vector database health."""
        start_time = time.time()
        try:
            from cognee.infrastructure.databases.vector.get_vector_engine import get_vector_engine
            from cognee.infrastructure.databases.vector.config import get_vectordb_config

Review comment (suggestion):
we can't await here
```suggestion
            session = engine.get_session()
```

Code after:
class HealthResponse(BaseModel):
    status: HealthStatus
    timestamp: str
    version: str
    uptime: int
    components: Dict[str, ComponentHealth]


class HealthChecker:
    def __init__(self):
        self.start_time = time.time()

    async def check_relational_db(self) -> ComponentHealth:
        """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.config import get_relational_config

            config = get_relational_config()
            engine = get_relational_engine()

            # Test connection by creating a session
            session = engine.get_session()
            if session:
                await session.close()

            response_time = int((time.time() - start_time) * 1000)
            return ComponentHealth(
                status=HealthStatus.HEALTHY,
                provider=config.db_provider,
                response_time_ms=response_time,
                details="Connection successful",
            )
        except Exception as e:
            response_time = int((time.time() - start_time) * 1000)
            logger.error(f"Relational DB health check failed: {str(e)}", exc_info=True)
            return ComponentHealth(
                status=HealthStatus.UNHEALTHY,
                provider="unknown",
                response_time_ms=response_time,
                details=f"Connection failed: {str(e)}",
            )

    async def check_vector_db(self) -> ComponentHealth:
        """Check vector database health."""
        start_time = time.time()
        try:
            from cognee.infrastructure.databases.vector.get_vector_engine import get_vector_engine