Move graph edit function implemention to a utils_graph.py to educe the size of lightray.py

lightrag/utils.py

@@ -893,6 +893,365 @@ def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
         return new_loop
 
 
+async def aexport_data(
+    chunk_entity_relation_graph,
+    entities_vdb,
+    relationships_vdb,
+    output_path: str,
+    file_format: str = "csv",
+    include_vector_data: bool = False,
+) -> None:
+    """
+    Asynchronously exports all entities, relations, and relationships to various formats.
+    
+    Args:
+        chunk_entity_relation_graph: Graph storage instance for entities and relations
+        entities_vdb: Vector database storage for entities
+        relationships_vdb: Vector database storage for relationships
+        output_path: The path to the output file (including extension).
+        file_format: Output format - "csv", "excel", "md", "txt".
+            - csv: Comma-separated values file
+            - excel: Microsoft Excel file with multiple sheets
+            - md: Markdown tables
+            - txt: Plain text formatted output
+        include_vector_data: Whether to include data from the vector database.
+    """
+    # Collect data
+    entities_data = []
+    relations_data = []
+    relationships_data = []
+
+    # --- Entities ---
+    all_entities = await chunk_entity_relation_graph.get_all_labels()
+    for entity_name in all_entities:
+        # Get entity information from graph
+        node_data = await chunk_entity_relation_graph.get_node(entity_name)
+        source_id = node_data.get("source_id") if node_data else None
+        
+        entity_info = {
+            "graph_data": node_data,
+            "source_id": source_id,
+        }
+        
+        # Optional: Get vector database information
+        if include_vector_data:
+            entity_id = compute_mdhash_id(entity_name, prefix="ent-")
+            vector_data = await entities_vdb.get_by_id(entity_id)
+            entity_info["vector_data"] = vector_data
+            
+        entity_row = {
+            "entity_name": entity_name,
+            "source_id": source_id,
+            "graph_data": str(entity_info["graph_data"]),  # Convert to string to ensure compatibility
+        }
+        if include_vector_data and "vector_data" in entity_info:
+            entity_row["vector_data"] = str(entity_info["vector_data"])
+        entities_data.append(entity_row)
+
+    # --- Relations ---
+    for src_entity in all_entities:
+        for tgt_entity in all_entities:
+            if src_entity == tgt_entity:
+                continue
+
+            edge_exists = await chunk_entity_relation_graph.has_edge(
+                src_entity, tgt_entity
+            )
+            if edge_exists:
+                # Get edge information from graph
+                edge_data = await chunk_entity_relation_graph.get_edge(
+                    src_entity, tgt_entity
+                )
+                source_id = edge_data.get("source_id") if edge_data else None
+                
+                relation_info = {
+                    "graph_data": edge_data,
+                    "source_id": source_id,
+                }
+                
+                # Optional: Get vector database information
+                if include_vector_data:
+                    rel_id = compute_mdhash_id(src_entity + tgt_entity, prefix="rel-")
+                    vector_data = await relationships_vdb.get_by_id(rel_id)
+                    relation_info["vector_data"] = vector_data
+                
+                relation_row = {
+                    "src_entity": src_entity,
+                    "tgt_entity": tgt_entity,
+                    "source_id": relation_info["source_id"],
+                    "graph_data": str(relation_info["graph_data"]),  # Convert to string
+                }
+                if include_vector_data and "vector_data" in relation_info:
+                    relation_row["vector_data"] = str(relation_info["vector_data"])
+                relations_data.append(relation_row)
+
+    # --- Relationships (from VectorDB) ---
+    all_relationships = await relationships_vdb.client_storage
+    for rel in all_relationships["data"]:
+        relationships_data.append(
+            {
+                "relationship_id": rel["__id__"],
+                "data": str(rel),  # Convert to string for compatibility
+            }
+        )
+
+    # Export based on format
+    if file_format == "csv":
+        # CSV export
+        with open(output_path, "w", newline="", encoding="utf-8") as csvfile:
+            # Entities
+            if entities_data:
+                csvfile.write("# ENTITIES\n")
+                writer = csv.DictWriter(csvfile, fieldnames=entities_data[0].keys())
+                writer.writeheader()
+                writer.writerows(entities_data)
+                csvfile.write("\n\n")
+
+            # Relations
+            if relations_data:
+                csvfile.write("# RELATIONS\n")
+                writer = csv.DictWriter(
+                    csvfile, fieldnames=relations_data[0].keys()
+                )
+                writer.writeheader()
+                writer.writerows(relations_data)
+                csvfile.write("\n\n")
+
+            # Relationships
+            if relationships_data:
+                csvfile.write("# RELATIONSHIPS\n")
+                writer = csv.DictWriter(
+                    csvfile, fieldnames=relationships_data[0].keys()
+                )
+                writer.writeheader()
+                writer.writerows(relationships_data)
+
+    elif file_format == "excel":
+        # Excel export
+        import pandas as pd
+        
+        entities_df = (
+            pd.DataFrame(entities_data) if entities_data else pd.DataFrame()
+        )
+        relations_df = (
+            pd.DataFrame(relations_data) if relations_data else pd.DataFrame()
+        )
+        relationships_df = (
+            pd.DataFrame(relationships_data)
+            if relationships_data
+            else pd.DataFrame()
+        )
+
+        with pd.ExcelWriter(output_path, engine="xlsxwriter") as writer:
+            if not entities_df.empty:
+                entities_df.to_excel(writer, sheet_name="Entities", index=False)
+            if not relations_df.empty:
+                relations_df.to_excel(writer, sheet_name="Relations", index=False)
+            if not relationships_df.empty:
+                relationships_df.to_excel(
+                    writer, sheet_name="Relationships", index=False
+                )
+
+    elif file_format == "md":
+        # Markdown export
+        with open(output_path, "w", encoding="utf-8") as mdfile:
+            mdfile.write("# LightRAG Data Export\n\n")
+
+            # Entities
+            mdfile.write("## Entities\n\n")
+            if entities_data:
+                # Write header
+                mdfile.write("| " + " | ".join(entities_data[0].keys()) + " |\n")
+                mdfile.write(
+                    "| "
+                    + " | ".join(["---"] * len(entities_data[0].keys()))
+                    + " |\n"
+                )
+
+                # Write rows
+                for entity in entities_data:
+                    mdfile.write(
+                        "| " + " | ".join(str(v) for v in entity.values()) + " |\n"
+                    )
+                mdfile.write("\n\n")
+            else:
+                mdfile.write("*No entity data available*\n\n")
+
+            # Relations
+            mdfile.write("## Relations\n\n")
+            if relations_data:
+                # Write header
+                mdfile.write("| " + " | ".join(relations_data[0].keys()) + " |\n")
+                mdfile.write(
+                    "| "
+                    + " | ".join(["---"] * len(relations_data[0].keys()))
+                    + " |\n"
+                )
+
+                # Write rows
+                for relation in relations_data:
+                    mdfile.write(
+                        "| "
+                        + " | ".join(str(v) for v in relation.values())
+                        + " |\n"
+                    )
+                mdfile.write("\n\n")
+            else:
+                mdfile.write("*No relation data available*\n\n")
+
+            # Relationships
+            mdfile.write("## Relationships\n\n")
+            if relationships_data:
+                # Write header
+                mdfile.write(
+                    "| " + " | ".join(relationships_data[0].keys()) + " |\n"
+                )
+                mdfile.write(
+                    "| "
+                    + " | ".join(["---"] * len(relationships_data[0].keys()))
+                    + " |\n"
+                )
+
+                # Write rows
+                for relationship in relationships_data:
+                    mdfile.write(
+                        "| "
+                        + " | ".join(str(v) for v in relationship.values())
+                        + " |\n"
+                    )
+            else:
+                mdfile.write("*No relationship data available*\n\n")
+
+    elif file_format == "txt":
+        # Plain text export
+        with open(output_path, "w", encoding="utf-8") as txtfile:
+            txtfile.write("LIGHTRAG DATA EXPORT\n")
+            txtfile.write("=" * 80 + "\n\n")
+
+            # Entities
+            txtfile.write("ENTITIES\n")
+            txtfile.write("-" * 80 + "\n")
+            if entities_data:
+                # Create fixed width columns
+                col_widths = {
+                    k: max(len(k), max(len(str(e[k])) for e in entities_data))
+                    for k in entities_data[0]
+                }
+                header = "  ".join(k.ljust(col_widths[k]) for k in entities_data[0])
+                txtfile.write(header + "\n")
+                txtfile.write("-" * len(header) + "\n")
+
+                # Write rows
+                for entity in entities_data:
+                    row = "  ".join(
+                        str(v).ljust(col_widths[k]) for k, v in entity.items()
+                    )
+                    txtfile.write(row + "\n")
+                txtfile.write("\n\n")
+            else:
+                txtfile.write("No entity data available\n\n")
+
+            # Relations
+            txtfile.write("RELATIONS\n")
+            txtfile.write("-" * 80 + "\n")
+            if relations_data:
+                # Create fixed width columns
+                col_widths = {
+                    k: max(len(k), max(len(str(r[k])) for r in relations_data))
+                    for k in relations_data[0]
+                }
+                header = "  ".join(
+                    k.ljust(col_widths[k]) for k in relations_data[0]
+                )
+                txtfile.write(header + "\n")
+                txtfile.write("-" * len(header) + "\n")
+
+                # Write rows
+                for relation in relations_data:
+                    row = "  ".join(
+                        str(v).ljust(col_widths[k]) for k, v in relation.items()
+                    )
+                    txtfile.write(row + "\n")
+                txtfile.write("\n\n")
+            else:
+                txtfile.write("No relation data available\n\n")
+
+            # Relationships
+            txtfile.write("RELATIONSHIPS\n")
+            txtfile.write("-" * 80 + "\n")
+            if relationships_data:
+                # Create fixed width columns
+                col_widths = {
+                    k: max(len(k), max(len(str(r[k])) for r in relationships_data))
+                    for k in relationships_data[0]
+                }
+                header = "  ".join(
+                    k.ljust(col_widths[k]) for k in relationships_data[0]
+                )
+                txtfile.write(header + "\n")
+                txtfile.write("-" * len(header) + "\n")
+
+                # Write rows
+                for relationship in relationships_data:
+                    row = "  ".join(
+                        str(v).ljust(col_widths[k]) for k, v in relationship.items()
+                    )
+                    txtfile.write(row + "\n")
+            else:
+                txtfile.write("No relationship data available\n\n")
+
+    else:
+        raise ValueError(
+            f"Unsupported file format: {file_format}. "
+            f"Choose from: csv, excel, md, txt"
+        )
+    if file_format is not None:
+        print(f"Data exported to: {output_path} with format: {file_format}")
+    else:
+        print("Data displayed as table format")
+
+
+def export_data(
+    chunk_entity_relation_graph,
+    entities_vdb,
+    relationships_vdb,
+    output_path: str,
+    file_format: str = "csv",
+    include_vector_data: bool = False,
+) -> None:
+    """
+    Synchronously exports all entities, relations, and relationships to various formats.
+    
+    Args:
+        chunk_entity_relation_graph: Graph storage instance for entities and relations
+        entities_vdb: Vector database storage for entities
+        relationships_vdb: Vector database storage for relationships
+        output_path: The path to the output file (including extension).
+        file_format: Output format - "csv", "excel", "md", "txt".
+            - csv: Comma-separated values file
+            - excel: Microsoft Excel file with multiple sheets
+            - md: Markdown tables
+            - txt: Plain text formatted output
+        include_vector_data: Whether to include data from the vector database.
+    """
+    try:
+        loop = asyncio.get_event_loop()
+    except RuntimeError:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+    loop.run_until_complete(
+        aexport_data(
+            chunk_entity_relation_graph,
+            entities_vdb,
+            relationships_vdb,
+            output_path,
+            file_format,
+            include_vector_data
+        )
+    )
+
+
 def lazy_external_import(module_name: str, class_name: str) -> Callable[..., Any]:
     """Lazily import a class from an external module based on the package of the caller."""
     # Get the caller's module and package