feat: add neighborhood-based graph traversal for retrievers

cognee/api/v1/search/search.py

@@ -41,6 +41,7 @@ async def search(
     triplet_distance_penalty: Optional[float] = 3.5,
     verbose: bool = False,
     retriever_specific_config: Optional[dict] = None,
+    neighborhood_depth: Optional[int] = None,
 ) -> List[SearchResult]:
     """
     Search and query the knowledge graph for insights, information, and connections.
@@ -229,6 +230,7 @@ async def search(
             triplet_distance_penalty=triplet_distance_penalty,
             verbose=verbose,
             retriever_specific_config=retriever_specific_config,
+            neighborhood_depth=neighborhood_depth,
         )
 
         n = len(filtered_search_results) if filtered_search_results else 0

cognee/infrastructure/databases/graph/graph_db_interface.py

@@ -277,6 +277,30 @@ async def get_connections(
         """
         raise NotImplementedError
 
+    @abstractmethod
+    async def get_neighborhood(
+        self,
+        node_ids: List[str],
+        depth: int = 1,
+        edge_types: Optional[List[str]] = None,
+    ) -> Tuple[List[Node], List[EdgeData]]:
+        """
+        Get the k-hop neighborhood subgraph around a set of seed nodes.
+
+        Returns all nodes and edges within `depth` hops of any seed node,
+        in the same format as get_graph_data().
+        Optional edge_type filtering to constrain traversal paths.
+
+        Parameters:
+        -----------
+
+            - node_ids (List[str]): Seed node identifiers to start traversal from.
+            - depth (int): Number of hops to traverse from each seed node. (default 1)
+            - edge_types (Optional[List[str]]): If provided, only traverse edges of these
+              relationship types. (default None)
+        """
+        raise NotImplementedError
+
     @abstractmethod
     async def get_filtered_graph_data(
         self, attribute_filters: List[Dict[str, List[Union[str, int]]]]

cognee/infrastructure/databases/graph/kuzu/adapter.py

@@ -1263,6 +1263,106 @@ async def get_graph_data(
             logger.error(f"Failed to get graph data: {e}")
             raise
 
+    async def get_neighborhood(
+        self,
+        node_ids: List[str],
+        depth: int = 1,
+        edge_types: Optional[List[str]] = None,
+    ) -> Tuple[List[Tuple[str, Dict[str, Any]]], List[Tuple[str, str, str, Dict[str, Any]]]]:
+        """
+        Get the k-hop neighborhood subgraph around a set of seed nodes.
+
+        Returns all nodes and edges within `depth` hops of any seed node,
+        in the same format as get_graph_data().
+        """
+        import time
+
+        start_time = time.time()
+
+        try:
+            if not node_ids:
+                logger.warning("No node IDs provided for neighborhood retrieval.")
+                return [], []
+
+            # Use variable-length path to find all nodes within depth hops
+            path_query = f"""
+            MATCH (seed:Node)-[r*1..{depth}]-(neighbor:Node)
+            WHERE seed.id IN $node_ids{" AND ALL(rel IN r WHERE rel.relationship_name IN $edge_types)" if edge_types else ""}
+            RETURN DISTINCT neighbor.id
+            """
+            params = {"node_ids": node_ids}
+            if edge_types:
+                params["edge_types"] = edge_types
+
+            neighbor_rows = await self.query(path_query, params)
+            neighbor_ids = [row[0] for row in neighbor_rows if row[0]]
+
+            # Combine seed nodes and neighbor nodes
+            all_ids = list(set(node_ids) | set(neighbor_ids))
+
+            # Fetch all nodes
+            nodes_query = """
+            MATCH (n:Node)
+            WHERE n.id IN $ids
+            RETURN n.id, {
+                name: n.name,
+                type: n.type,
+                properties: n.properties
+            }
+            """
+            node_rows = await self.query(nodes_query, {"ids": all_ids})
+            formatted_nodes = []
+            for n in node_rows:
+                if n[0]:
+                    node_id = str(n[0])
+                    props = n[1]
+                    if props.get("properties"):
+                        try:
+                            additional_props = json.loads(props["properties"])
+                            props.update(additional_props)
+                            del props["properties"]
+                        except json.JSONDecodeError:
+                            logger.warning(f"Failed to parse properties JSON for node {node_id}")
+                    formatted_nodes.append((node_id, props))
+
+            if not formatted_nodes:
+                logger.warning("No nodes found in neighborhood.")
+                return [], []
+
+            # Fetch all edges between the collected nodes
+            edges_query = """
+            MATCH (n:Node)-[r]->(m:Node)
+            WHERE n.id IN $ids AND m.id IN $ids
+            RETURN n.id, m.id, r.relationship_name, r.properties
+            """
+            edge_rows = await self.query(edges_query, {"ids": all_ids})
+            formatted_edges = []
+            for e in edge_rows:
+                if e and len(e) >= 3:
+                    source_id = str(e[0])
+                    target_id = str(e[1])
+                    rel_type = str(e[2])
+                    props = {}
+                    if len(e) > 3 and e[3]:
+                        try:
+                            props = json.loads(e[3])
+                        except (json.JSONDecodeError, TypeError):
+                            logger.warning(
+                                f"Failed to parse edge properties for {source_id}->{target_id}"
+                            )
+                    formatted_edges.append((source_id, target_id, rel_type, props))
+
+            retrieval_time = time.time() - start_time
+            logger.info(
+                f"Neighborhood retrieval ({depth}-hop): {len(formatted_nodes)} nodes and "
+                f"{len(formatted_edges)} edges in {retrieval_time:.2f}s"
+            )
+            return formatted_nodes, formatted_edges
+
+        except Exception as e:
+            logger.error(f"Failed to get neighborhood: {e}")
+            raise
+
     async def get_nodeset_subgraph(
         self, node_type: Type[Any], node_name: List[str]
     ) -> Tuple[List[Tuple[str, dict]], List[Tuple[str, str, str, dict]]]:

cognee/infrastructure/databases/graph/neo4j_driver/adapter.py

@@ -991,6 +991,91 @@ async def get_graph_data(self):
             logger.error(f"Error during graph data retrieval: {str(e)}")
             raise
 
+    async def get_neighborhood(
+        self,
+        node_ids: List[str],
+        depth: int = 1,
+        edge_types: Optional[List[str]] = None,
+    ) -> Tuple[List[Tuple[str, Dict[str, Any]]], List[Tuple[str, str, str, Dict[str, Any]]]]:
+        """
+        Get the k-hop neighborhood subgraph around a set of seed nodes.
+
+        Returns all nodes and edges within `depth` hops of any seed node,
+        in the same format as get_graph_data().
+        """
+        import time
+
+        start_time = time.time()
+
+        try:
+            if not node_ids:
+                logger.warning("No node IDs provided for neighborhood retrieval.")
+                return [], []
+
+            # Collect all node IDs within depth hops, then fetch nodes and edges
+            if edge_types:
+                path_query = f"""
+                MATCH path = (seed)-[*1..{depth}]-(neighbor)
+                WHERE seed.id IN $node_ids
+                  AND ALL(r IN relationships(path) WHERE TYPE(r) IN $edge_types)
+                RETURN DISTINCT neighbor.id AS nid
+                """
+            else:
+                path_query = f"""
+                MATCH (seed)-[*1..{depth}]-(neighbor)
+                WHERE seed.id IN $node_ids
+                RETURN DISTINCT neighbor.id AS nid
+                """
+
+            params = {"node_ids": node_ids}
+            if edge_types:
+                params["edge_types"] = edge_types
+
+            result = await self.query(path_query, params)
+            neighbor_ids = [record["nid"] for record in result if record.get("nid")]
+
+            all_ids = list(set(node_ids) | set(neighbor_ids))
+
+            # Step 2: Fetch all nodes
+            nodes_query = """
+            MATCH (n)
+            WHERE n.id IN $ids
+            RETURN n.id AS id, properties(n) AS properties
+            """
+            nodes_result = await self.query(nodes_query, {"ids": all_ids})
+            nodes = []
+            for record in nodes_result:
+                nodes.append((record["properties"]["id"], record["properties"]))
+
+            # Step 3: Fetch all edges between collected nodes
+            edges_query = """
+            MATCH (n)-[r]->(m)
+            WHERE n.id IN $ids AND m.id IN $ids
+            RETURN properties(r) AS properties, TYPE(r) AS type
+            """
+            edges_result = await self.query(edges_query, {"ids": all_ids})
+            edges = []
+            for record in edges_result:
+                edges.append(
+                    (
+                        record["properties"]["source_node_id"],
+                        record["properties"]["target_node_id"],
+                        record["type"],
+                        record["properties"],
+                    )
+                )
+
+            retrieval_time = time.time() - start_time
+            logger.info(
+                f"Neighborhood retrieval ({depth}-hop): {len(nodes)} nodes and "
+                f"{len(edges)} edges in {retrieval_time:.2f}s"
+            )
+            return (nodes, edges)
+
+        except Exception as e:
+            logger.error(f"Error during neighborhood retrieval: {str(e)}")
+            raise
+
     async def get_id_filtered_graph_data(self, target_ids: list[str]):
         """
         Retrieve graph data filtered by specific node IDs, including their direct neighbors

cognee/infrastructure/databases/graph/neptune_driver/adapter.py

@@ -669,6 +669,75 @@ async def get_graph_data(self) -> Tuple[List[Node], List[EdgeData]]:
             logger.error(f"Failed to get graph data: {error_msg}")
             raise Exception(f"Failed to get graph data: {error_msg}") from e
 
+    async def get_neighborhood(
+        self,
+        node_ids: List[str],
+        depth: int = 1,
+        edge_types: Optional[List[str]] = None,
+    ) -> Tuple[List[Node], List[EdgeData]]:
+        """
+        Get the k-hop neighborhood subgraph around a set of seed nodes.
+
+        Returns all nodes and edges within `depth` hops of any seed node,
+        in the same format as get_graph_data().
+        """
+        try:
+            if not node_ids:
+                logger.warning("No node IDs provided for neighborhood retrieval.")
+                return [], []
+
+            # Step 1: Find all neighbor node IDs within depth hops
+            if edge_types:
+                allowed = "|".join(edge_types)
+                path_query = f"""
+                MATCH (seed:{self._GRAPH_NODE_LABEL})-[:{allowed}*1..{depth}]-(neighbor:{self._GRAPH_NODE_LABEL})
+                WHERE seed.`~id` IN $node_ids
+                RETURN DISTINCT id(neighbor) AS nid
+                """
+            else:
+                path_query = f"""
+                MATCH (seed:{self._GRAPH_NODE_LABEL})-[*1..{depth}]-(neighbor:{self._GRAPH_NODE_LABEL})
+                WHERE seed.`~id` IN $node_ids
+                RETURN DISTINCT id(neighbor) AS nid
+                """
+
+            result = await self.query(path_query, {"node_ids": node_ids})
+            neighbor_ids = [record["nid"] for record in result if record.get("nid")]
+
+            all_ids = list(set(node_ids) | set(neighbor_ids))
+
+            # Step 2: Fetch all nodes
+            nodes_query = f"""
+            MATCH (n:{self._GRAPH_NODE_LABEL})
+            WHERE id(n) IN $ids
+            RETURN id(n) AS node_id, properties(n) AS properties
+            """
+            nodes_result = await self.query(nodes_query, {"ids": all_ids})
+            nodes = [(r["node_id"], r["properties"]) for r in nodes_result]
+
+            # Step 3: Fetch all edges between collected nodes
+            edges_query = f"""
+            MATCH (source:{self._GRAPH_NODE_LABEL})-[r]->(target:{self._GRAPH_NODE_LABEL})
+            WHERE id(source) IN $ids AND id(target) IN $ids
+            RETURN id(source) AS source_id, id(target) AS target_id,
+                   type(r) AS relationship_name, properties(r) AS properties
+            """
+            edges_result = await self.query(edges_query, {"ids": all_ids})
+            edges = [
+                (r["source_id"], r["target_id"], r["relationship_name"], r["properties"])
+                for r in edges_result
+            ]
+
+            logger.debug(
+                f"Neighborhood retrieval ({depth}-hop): {len(nodes)} nodes and {len(edges)} edges"
+            )
+            return (nodes, edges)
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to get neighborhood: {error_msg}")
+            raise Exception(f"Failed to get neighborhood: {error_msg}") from e
+
     async def get_graph_metrics(self, include_optional: bool = False) -> Dict[str, Any]:
         """
         Fetch metrics and statistics of the graph, possibly including optional details.