Convert a personal-graph to and from NetworkX

.github/workflows/personal-graph.yaml

@@ -62,5 +62,5 @@ jobs:
 
       - name: Test with pytest
         run: |
-         cd personal_graph 
+         cd tests 
          poetry run pytest -vvv

examples/pythonic_api.py

@@ -1,6 +1,7 @@
 import os
 import logging
 from personal_graph.graph import Graph
+from personal_graph.ml import to_networkx, from_networkx
 from personal_graph.models import Node, EdgeInput, KnowledgeGraph, Edge
 
 
@@ -104,6 +105,14 @@ def main(url, token):
         )
         logging.info(graph.visualize_graph(kg))
 
+        # Transforms to and from networkx do not alter the graph
+        g2 = from_networkx(
+            to_networkx(graph, post_visualize=True), post_visualize=True, override=False
+        )
+        if graph == g2:
+            logging.info("TRUE")
+        assert graph == g2
+
         graph.save()
 
 

personal_graph/database.py

@@ -539,6 +539,16 @@ def _find_node(cursor, connection):
     return _find_node
 
 
+def find_label(identifier: Any) -> CursorExecFunction:
+    def _find_label(cursor, connection):
+        node_label = cursor.execute(
+            "SELECT label from nodes where id=?", (identifier,)
+        ).fetchone()
+        return node_label
+
+    return _find_label
+
+
 def _parse_search_results(results: List[Tuple], idx: int = 0) -> List[Dict]:
     return [json.loads(item[idx]) for item in results]
 
@@ -650,6 +660,9 @@ def _merge(cursor, connection):
                 cursor, connection
             )
 
+            if similar_nodes is None:
+                continue
+
             if len(similar_nodes) < 1:
                 continue
 
@@ -800,6 +813,34 @@ def _fetch_nodes_from_db(cursor, connection):
     return _fetch_nodes_from_db
 
 
+def find_indegree_edges(target_id: Any) -> CursorExecFunction:
+    def _indegree_edges(cursor, connection):
+        indegree = cursor.execute(
+            "SELECT source, label, attributes from edges where target=? ", (target_id,)
+        )
+
+        if indegree:
+            indegree = indegree.fetchall()
+
+        return indegree
+
+    return _indegree_edges
+
+
+def find_outdegree_edges(source_id: Any) -> CursorExecFunction:
+    def _outdegree_edges(cursor, connection):
+        outdegree = cursor.execute(
+            "SELECT target, label, attributes from edges where source=? ", (source_id,)
+        )
+
+        if outdegree:
+            outdegree = outdegree.fetchall()
+
+        return outdegree
+
+    return _outdegree_edges
+
+
 def all_connected_nodes(node_or_edge: Union[Node | Edge]) -> CursorExecFunction:
     def _connected_nodes(cursor, connection):
         nodes = None

personal_graph/graph.py

@@ -28,6 +28,9 @@
     find_similar_nodes,
     nodes_list,
     vector_search_node,
+    find_label,
+    find_outdegree_edges,
+    find_indegree_edges,
 )
 from .natural import insert_into_graph, search_from_graph, visualize_knowledge_graph
 from .visualizers import graphviz_visualize
@@ -38,6 +41,11 @@ def __init__(self, db_url: Optional[str] = None, auth_token: Optional[str] = Non
         self.db_url = db_url
         self.auth_token = auth_token
 
+    def __eq__(self, other):
+        if not isinstance(other, Graph):
+            return "Not of Graph Type"
+        return self.db_url == other.db_url and self.auth_token == other.auth_token
+
     def __enter__(self, schema_file: str = "schema.sql") -> Graph:
         if not self.db_url:
             # Support for local SQLite database
@@ -131,6 +139,9 @@ def remove_nodes(self, ids: List[Any]) -> None:
     def search_node(self, node_id: Any) -> Any:
         return atomic(find_node(node_id), self.db_url, self.auth_token)
 
+    def search_node_label(self, node_id: Any) -> Any:
+        return atomic(find_label(node_id), self.db_url, self.auth_token)
+
     def traverse(
         self, source: Any, target: Optional[Any] = None, with_bodies: bool = False
     ) -> List:
@@ -170,6 +181,12 @@ def visualize(self, file: str, path: List[str]) -> Digraph:
     def fetch_ids_from_db(self) -> List[str]:
         return atomic(nodes_list(), self.db_url, self.auth_token)
 
+    def search_indegree_edges(self, target) -> List[Any]:
+        return atomic(find_indegree_edges(target), self.db_url, self.auth_token)
+
+    def search_outdegree_edges(self, source) -> List[Any]:
+        return atomic(find_outdegree_edges(source), self.db_url, self.auth_token)
+
     def is_unique_prompt(self, text: str, threshold: float) -> bool:
         similar_nodes = atomic(
             vector_search_node({"body": text}, threshold, 1),

personal_graph/ml.py

@@ -1,4 +1,183 @@
-"""
-Provide functions to convert the data in the graph db into formats compatible with
-graph analysis frameworks such as Networkx and DGL
-"""
+import json
+import os
+
+import networkx as nx  # type: ignore
+from typing import Dict, Any
+import matplotlib.pyplot as plt
+from graphviz import Digraph  # type: ignore
+
+import personal_graph.graph as pg
+from personal_graph.models import Node, KnowledgeGraph, Edge, EdgeInput
+
+
+def to_networkx(graph: pg.Graph, *, post_visualize: bool = False) -> nx.Graph:
+    """
+    Convert the graph database to a NetworkX DiGraph object.
+    """
+    G = nx.Graph()  # Empty Graph with no nodes and edges
+
+    node_ids = graph.fetch_ids_from_db()
+    # Add edges to networkX
+    for source_id in node_ids:
+        for target_id, edge_label, edge_data in graph.search_outdegree_edges(source_id):
+            edge_data = json.loads(edge_data)
+            edge_data["label"] = edge_label
+            G.add_edge(source_id, target_id, **edge_data)
+
+    for target_id in node_ids:
+        for source_id, edge_label, edge_data in graph.search_indegree_edges(target_id):
+            edge_data = json.loads(edge_data)
+            edge_data["label"] = edge_label
+            G.add_edge(source_id, target_id, **edge_data)
+
+    node_ids_with_edges = set([node for edge in G.edges() for node in edge])
+    for node_id in node_ids:
+        if node_id not in node_ids_with_edges:
+            node_data = graph.search_node(node_id)
+            node_label = graph.search_node_label(node_id)
+            node_data["label"] = node_label
+            G.add_node(node_id, **node_data)
+
+    if post_visualize:
+        # Visualizing the NetworkX Graph
+        plt.figure(figsize=(20, 20), dpi=100)  # Increase the figure size and resolution
+        pos = nx.spring_layout(
+            G, scale=6
+        )  # Use spring layout for better node positioning
+
+        nx.draw_networkx(
+            G,
+            pos,
+            with_labels=True,
+            nodelist=G.nodes(),
+            edgelist=G.edges(),
+            node_size=600,
+            node_color="skyblue",
+            edge_color="gray",
+            width=1.5,
+        )
+        nx.draw_networkx_edge_labels(
+            G, pos, edge_labels=nx.get_edge_attributes(G, "label")
+        )
+        plt.axis("off")  # Show the axes
+        plt.savefig("networkX_graph.png")
+
+    return G
+
+
+def from_networkx(
+    network_graph: nx, *, post_visualize: bool = False, override: bool = True
+) -> pg.Graph:
+    with pg.Graph(os.getenv("LIBSQL_URL"), os.getenv("LIBSQL_AUTH_TOKEN")) as graph:
+        if override:
+            node_ids = graph.fetch_ids_from_db()
+            graph.remove_nodes(node_ids)
+
+        node_ids_with_edges = set()
+        kg = KnowledgeGraph()
+
+        # Convert networkX edges to personal graph edges
+        for source_id, target_id, edge_data in network_graph.edges(data=True):
+            edge_attributes: Dict[str, Any] = edge_data
+            edge_label: str = edge_attributes["label"]
+            node_ids_with_edges.add(str(source_id))
+            node_ids_with_edges.add(str(target_id))
+
+            if not override:
+                # Check if the node with the given id exists, if not then firstly add the node.
+                source = graph.search_node(source_id)
+                if source is []:
+                    node_ids_with_edges.remove(str(source_id))
+                    graph.add_node(
+                        Node(
+                            id=str(source_id),
+                            label=edge_label if edge_label else "",
+                            attributes=edge_attributes,
+                        )
+                    )
+
+                target = graph.search_node(target_id)
+                if target is []:
+                    node_ids_with_edges.remove(str(target_id))
+                    graph.add_node(
+                        Node(
+                            id=str(target_id),
+                            label=edge_label if edge_label else "",
+                            attributes=edge_attributes,
+                        )
+                    )
+
+            # After adding the new nodes if exists , add an edge
+            edge = Edge(
+                source=str(source_id),
+                target=str(target_id),
+                label=edge_label if edge_label else "",
+                attributes=edge_attributes,
+            )
+            kg.edges.append(edge)
+
+        # Convert networkX nodes to personal graph nodes
+        for node_id, node_data in network_graph.nodes(data=True):
+            if str(node_id) not in node_ids_with_edges:
+                node_attributes: Dict[str, Any] = node_data
+                node_label: str = node_attributes.pop("label", "")
+                node = Node(
+                    id=str(node_id),
+                    label=node_label[0] if node_label else "",
+                    attributes=node_attributes,
+                )
+
+                if not override:
+                    # Check if the node exists
+                    if_node_exists = graph.search_node(node_id)
+
+                    if if_node_exists:
+                        graph.update_node(node)
+                    else:
+                        graph.add_node(node)
+                kg.nodes.append(node)
+
+        for edge in kg.edges:
+            source_node_attributes = graph.search_node(edge.source)
+            source_node_label = graph.search_node_label(edge.source)
+            target_node_attributes = graph.search_node(edge.target)
+            target_node_label = graph.search_node_label(edge.target)
+            final_edge_to_be_inserted = EdgeInput(
+                source=Node(
+                    id=edge.source,
+                    label=source_node_label
+                    if isinstance(source_node_label, str)
+                    else "Sample label",
+                    attributes=source_node_attributes
+                    if isinstance(source_node_attributes, Dict)
+                    else "Sample Attributes",
+                ),
+                target=Node(
+                    id=edge.target,
+                    label=target_node_label
+                    if isinstance(target_node_label, str)
+                    else "Sample label",
+                    attributes=target_node_attributes
+                    if isinstance(target_node_attributes, Dict)
+                    else "Sample Attributes",
+                ),
+                label=edge.label if isinstance(edge.label, str) else "Sample label",
+                attributes=edge.attributes
+                if isinstance(edge_attributes, Dict)
+                else "Sample Attributes",
+            )
+            graph.add_edge(final_edge_to_be_inserted)
+
+    if post_visualize:
+        # Visualize the personal graph using graphviz
+        dot = Digraph()
+
+        for node in kg.nodes:
+            dot.node(node.id, label=f"{node.label}: {node.id}")
+
+        for edge in kg.edges:
+            dot.edge(edge.source, edge.target, label=edge.label)
+
+        dot.render("personal_graph.gv", view=True)
+
+    return graph

personal_graph/models.py

@@ -27,7 +27,6 @@ class Edge(BaseModel):
     )
     label: str = Field(
         ...,
-        min_length=2,
         description="Most related and unique name associated with the edge.",
     )
     attributes: Union[str, Dict[str, str]] = Field(