topoteretes--cognee/distributed/workers/graph_saving_worker.py

import os
import modal
import asyncio
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential

from distributed.app import app
from distributed.signal import QueueSignal
from distributed.modal_image import image
from distributed.queues import add_nodes_and_edges_queue
from distributed.graph_write_batch import group_graph_writes, apply_grouped_graph_writes

from cognee.shared.logging_utils import get_logger
from cognee.infrastructure.databases.graph import get_graph_engine
from cognee.infrastructure.databases.graph.config import get_graph_config

logger = get_logger("graph_saving_worker")


class GraphDatabaseDeadlockError(Exception):
    message = "A deadlock occurred while trying to add data points to the vector database."


def is_deadlock_error(error):
    graph_config = get_graph_config()

    if graph_config.graph_database_provider == "neo4j":
        # Neo4j
        from neo4j.exceptions import TransientError

        if isinstance(error, TransientError) and (
            error.code == "Neo.TransientError.Transaction.DeadlockDetected"
        ):
            return True

    # Ladybug
    if "deadlock" in str(error).lower() or "cannot acquire lock" in str(error).lower():
        return True

    return False


secret_name = os.environ.get("MODAL_SECRET_NAME", "distributed_cognee")


@app.function(
    retries=3,
    image=image,
    timeout=86400,
    max_containers=1,
    secrets=[modal.Secret.from_name(secret_name)],
)
async def graph_saving_worker():
    print("Started processing of nodes and edges; starting graph engine queue.")
    graph_engine = await get_graph_engine()
    # Defines how many data packets do we glue together from the queue before ingesting them into the graph database
    BATCH_SIZE = 25
    stop_seen = False

    while True:
        if stop_seen:
            print("Finished processing all data points; stopping graph engine queue consumer.")
            return True

        if await add_nodes_and_edges_queue.len.aio() != 0:
            try:
                print("Remaining elements in queue:")
                print(await add_nodes_and_edges_queue.len.aio())

                all_items = []
                for _ in range(min(BATCH_SIZE, await add_nodes_and_edges_queue.len.aio())):
                    nodes_and_edges = await add_nodes_and_edges_queue.get.aio(block=False)

                    if not nodes_and_edges:
                        continue

                    if nodes_and_edges == QueueSignal.STOP:
                        await add_nodes_and_edges_queue.put.aio(QueueSignal.STOP)
                        stop_seen = True
                        break

                    # Payload: (nodes, edges, source_ref_key, pipeline_run_id).
                    if isinstance(nodes_and_edges, (list, tuple)) and len(nodes_and_edges) == 4:
                        all_items.append(nodes_and_edges)
                    else:
                        print(f"Malformed queue item skipped: {type(nodes_and_edges)!r}")

                # Group by provenance key so each data item's source ref is folded
                # onto its own artifacts; all nodes are still written before any
                # edge (across groups) so cross-item edges find their endpoints.
                groups = group_graph_writes(all_items)

                if groups:
                    total_nodes = sum(len(nodes) for nodes, _edges in groups.values())
                    total_edges = sum(len(edges) for _nodes, edges in groups.values())
                    print(
                        f"Adding {total_nodes} nodes and {total_edges} edges "
                        f"across {len(groups)} provenance group(s)."
                    )

                    @retry(
                        retry=retry_if_exception_type(GraphDatabaseDeadlockError),
                        stop=stop_after_attempt(3),
                        wait=wait_exponential(multiplier=2, min=1, max=6),
                    )
                    async def save_graph_nodes(new_nodes, source_ref_key, pipeline_run_id):
                        try:
                            await graph_engine.add_nodes(
                                new_nodes,
                                source_ref_key=source_ref_key,
                                pipeline_run_id=pipeline_run_id,
                                distributed=False,
                            )
                        except Exception as error:
                            if is_deadlock_error(error):
                                raise GraphDatabaseDeadlockError()

                    @retry(
                        retry=retry_if_exception_type(GraphDatabaseDeadlockError),
                        stop=stop_after_attempt(3),
                        wait=wait_exponential(multiplier=2, min=1, max=6),
                    )
                    async def save_graph_edges(new_edges, source_ref_key, pipeline_run_id):
                        try:
                            await graph_engine.add_edges(
                                new_edges,
                                source_ref_key=source_ref_key,
                                pipeline_run_id=pipeline_run_id,
                                distributed=False,
                            )
                        except Exception as error:
                            if is_deadlock_error(error):
                                raise GraphDatabaseDeadlockError()

                    await apply_grouped_graph_writes(groups, save_graph_nodes, save_graph_edges)

                    print("Finished adding nodes and edges.")

            except modal.exception.DeserializationError as error:
                logger.error(f"Deserialization error: {str(error)}")
                continue

        else:
            print("No jobs, go to sleep.")
            await asyncio.sleep(5)