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

Code before:
        """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

            config = get_vectordb_config()
            engine = get_vector_engine()

            # Test basic operation - just check if engine is accessible
            if hasattr(engine, "health_check"):
                await engine.health_check()
            elif hasattr(engine, "list_tables"):
                # For LanceDB and similar
                engine.list_tables()

            response_time = int((time.time() - start_time) * 1000)
            return ComponentHealth(
                status=HealthStatus.HEALTHY,
                provider=config.vector_db_provider,
                response_time_ms=response_time,
                details="Index accessible",
            )
        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_graph_db(self) -> ComponentHealth:
        """Check graph database health."""
        start_time = time.time()
        try:
            from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine
            from cognee.infrastructure.databases.graph.config import get_graph_config

            config = get_graph_config()
            engine = await get_graph_engine()

            # Test basic operation with actual graph query
            if hasattr(engine, "execute"):
                # For SQL-like graph DBs (Neo4j, Memgraph)
                await engine.execute("MATCH () RETURN count(*) LIMIT 1")
            elif hasattr(engine, "get_nodes"):
                # For other graph engines - try to get nodes
                list(engine.get_nodes(limit=1))
            # If engine exists but no test method, consider it healthy


Review comment (bug):
Same as in relational_db health check
```suggestion
            response_time = int((time.time() - start_time) * 1000)
            logger.error(f"Vector DB health check failed: {str(e)}", exc_info=True)
            return ComponentHealth(
```

Code after:
    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

            config = get_vectordb_config()
            engine = get_vector_engine()

            # Test basic operation - just check if engine is accessible
            if hasattr(engine, "health_check"):
                await engine.health_check()
            elif hasattr(engine, "list_tables"):
                # For LanceDB and similar
                engine.list_tables()

            response_time = int((time.time() - start_time) * 1000)
            return ComponentHealth(
                status=HealthStatus.HEALTHY,
                provider=config.vector_db_provider,
                response_time_ms=response_time,
                details="Index accessible",
            )
        except Exception as e:
            response_time = int((time.time() - start_time) * 1000)
            logger.error(f"Vector 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_graph_db(self) -> ComponentHealth:
        """Check graph database health."""
        start_time = time.time()
        try:
            from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine
            from cognee.infrastructure.databases.graph.config import get_graph_config

            config = get_graph_config()
            engine = await get_graph_engine()

            # Test basic operation with actual graph query
            if hasattr(engine, "execute"):
                # For SQL-like graph DBs (Neo4j, Memgraph)
                await engine.execute("MATCH () RETURN count(*) LIMIT 1")
            elif hasattr(engine, "query"):
                # For other graph engines
                engine.query("MATCH () RETURN count(*) LIMIT 1", {})