Repository: topoteretes/cognee
File: cognee/modules/graph/cognee_graph/CogneeGraph.py (Python)

Code before:
        node = self.get_node(node_id)
        if node:
            return node.skeleton_edges
        else:
            raise EntityNotFoundError(message=f"Node with id {node_id} does not exist.")

    def get_edges(self) -> List[Edge]:
        return self.edges

    async def project_graph_from_db(
        self,
        adapter: Union[GraphDBInterface],
        node_properties_to_project: List[str],
        edge_properties_to_project: List[str],
        directed=True,
        node_dimension=1,
        edge_dimension=1,
        memory_fragment_filter=[],
        node_type: Optional[Type] = None,
        node_name: Optional[List[str]] = None,
    ) -> None:
        logger.info("Starting graph projection from database")
        logger.debug(
            f"Projection parameters: directed={directed}, node_dimension={node_dimension}, edge_dimension={edge_dimension}"
        )
        logger.debug(f"Node properties to project: {node_properties_to_project}")
        logger.debug(f"Edge properties to project: {edge_properties_to_project}")

        if node_dimension < 1 or edge_dimension < 1:
            logger.error(
                f"Invalid dimensions: node_dimension={node_dimension}, edge_dimension={edge_dimension}"
            )
            raise InvalidValueError(message="Dimensions must be positive integers")

        try:
            import time

            start_time = time.time()

            # Determine projection strategy and log it
            if node_type is not None and node_name is not None:
                logger.info(
                    f"Using nodeset subgraph projection strategy for node_type={node_type.__name__}, node_names={node_name}"
                )
                nodes_data, edges_data = await adapter.get_nodeset_subgraph(
                    node_type=node_type, node_name=node_name
                )
                if not nodes_data or not edges_data:
                    logger.warning("Nodeset subgraph projection returned empty results")
                    raise EntityNotFoundError(
                        message="Nodeset does not exist, or empty nodetes projected from the database."

Review comment (bug):
I dont think this I would log. Its better to keep aggragated info, time, and node or edge in the case something fails. Otherwise this will dump the whole graph together with properties and everything.

Code after:
    def get_edges(self) -> List[Edge]:
        return self.edges

    async def project_graph_from_db(
        self,
        adapter: Union[GraphDBInterface],
        node_properties_to_project: List[str],
        edge_properties_to_project: List[str],
        directed=True,
        node_dimension=1,
        edge_dimension=1,
        memory_fragment_filter=[],
        node_type: Optional[Type] = None,
        node_name: Optional[List[str]] = None,
    ) -> None:
        if node_dimension < 1 or edge_dimension < 1:
            raise InvalidValueError(message="Dimensions must be positive integers")

        try:
            import time

            start_time = time.time()

            # Determine projection strategy
            if node_type is not None and node_name is not None:
                nodes_data, edges_data = await adapter.get_nodeset_subgraph(
                    node_type=node_type, node_name=node_name
                )
                if not nodes_data or not edges_data:
                    raise EntityNotFoundError(
                        message="Nodeset does not exist, or empty nodetes projected from the database."
                    )
            elif len(memory_fragment_filter) == 0:
                nodes_data, edges_data = await adapter.get_graph_data()
                if not nodes_data or not edges_data:
                    raise EntityNotFoundError(message="Empty graph projected from the database.")
            else:
                nodes_data, edges_data = await adapter.get_filtered_graph_data(
                    attribute_filters=memory_fragment_filter
                )
                if not nodes_data or not edges_data:
                    raise EntityNotFoundError(
                        message="Empty filtered graph projected from the database."
                    )

            # Process nodes
            for node_id, properties in nodes_data:
                node_attributes = {key: properties.get(key) for key in node_properties_to_project}
                self.add_node(Node(str(node_id), node_attributes, dimension=node_dimension))

            # Process edges