lightrag-agensgraph 0.1.0__tar.gz

lightrag_agensgraph-0.1.0/PKG-INFO

@@ -0,0 +1,60 @@
+Metadata-Version: 2.4
+Name: lightrag-agensgraph
+Version: 0.1.0
+Summary: Lightrag graph store implementation for agensgraph
+Keywords: lightrag,agensgraph,graph store,integration
+Author: Muhammad Taha Naveed
+Author-email: skaisw@skaiworldwide.com
+Requires-Python: >=3.9,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Dist: lightrag-hku (>=1.3.9)
+Requires-Dist: psycopg[binary,pool] (>=3.1.0)
+Project-URL: Homepage, https://github.com/skaiworldwide-oss/agensgraph-ai/tree/main/lightrag
+Project-URL: Repository, https://github.com/skaiworldwide-oss/agensgraph-ai/tree/main/lightrag
+Description-Content-Type: text/markdown
+
+# LightRag Knowledge Graph Integration: AgensGraph
+
+This plugin adds support for storing and querying knowledge graphs in [AgensGraph](https://github.com/skaiworldwide-oss/agensgraph) with [LightRAG](https://lightrag.github.io/).
+
+## Build
+
+You can use the following commands to build the plugin from source.
+
+```bash
+pip install poetry
+poetry install
+poetry build
+```
+
+## Usage
+
+```
+pip install lightrag-hku lightrag_agensgraph-0.1.0-py3-none-any.whl
+```
+
+```python
+from lightrag import LightRag
+import lighrag_agensgraph
+
+os.environ["AGENSGRAPH_DB"] = ""
+os.environ["AGENSGRAPH_USER"] = ""
+os.environ["AGENSGRAPH_PASSWORD"] = ""
+os.environ["AGENSGRAPH_HOST"] = ""
+os.environ["AGENSGRAPH_PORT"] = ""
+os.environ["AGENSGRAPH_GRAPHNAME"] = ""
+
+rag = LightRAG(
+    graph_storage="AgensgraphStorage",
+    ...
+)
+```
+
+See [examples](./examples/) and [tests](./tests/) for more details on how to use the plugin with LightRAG.
+

lightrag_agensgraph-0.1.0/README.md

@@ -0,0 +1,38 @@
+# LightRag Knowledge Graph Integration: AgensGraph
+
+This plugin adds support for storing and querying knowledge graphs in [AgensGraph](https://github.com/skaiworldwide-oss/agensgraph) with [LightRAG](https://lightrag.github.io/).
+
+## Build
+
+You can use the following commands to build the plugin from source.
+
+```bash
+pip install poetry
+poetry install
+poetry build
+```
+
+## Usage
+
+```
+pip install lightrag-hku lightrag_agensgraph-0.1.0-py3-none-any.whl
+```
+
+```python
+from lightrag import LightRag
+import lighrag_agensgraph
+
+os.environ["AGENSGRAPH_DB"] = ""
+os.environ["AGENSGRAPH_USER"] = ""
+os.environ["AGENSGRAPH_PASSWORD"] = ""
+os.environ["AGENSGRAPH_HOST"] = ""
+os.environ["AGENSGRAPH_PORT"] = ""
+os.environ["AGENSGRAPH_GRAPHNAME"] = ""
+
+rag = LightRAG(
+    graph_storage="AgensgraphStorage",
+    ...
+)
+```
+
+See [examples](./examples/) and [tests](./tests/) for more details on how to use the plugin with LightRAG.

lightrag_agensgraph-0.1.0/lightrag_agensgraph/__init__.py

@@ -0,0 +1,6 @@
+import sys, types
+from lightrag_agensgraph.kg.agensgraph_impl import AgensgraphStorage
+import lightrag.kg
+
+lightrag.kg.STORAGE_IMPLEMENTATIONS["GRAPH_STORAGE"]["implementations"].append("AgensgraphStorage")
+lightrag.kg.STORAGES["AgensgraphStorage"] = "lightrag_agensgraph.kg.agensgraph_impl"

lightrag_agensgraph-0.1.0/lightrag_agensgraph/kg/agensgraph_impl.py

@@ -0,0 +1,1075 @@
+import asyncio
+import inspect
+import re, json
+import os
+import sys
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import Any, Dict, List, NamedTuple, Optional, Union, final
+import pipmaster as pm
+from lightrag.types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge
+from typing import Any, List, Dict, Optional, Tuple, NamedTuple, Pattern
+
+from tenacity import (
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from lightrag.utils import logger
+from lightrag.base import BaseGraphStorage
+try:
+    from lightrag.constants import GRAPH_FIELD_SEP
+except ImportError:
+    from lightrag.prompt import GRAPH_FIELD_SEP
+
+if sys.platform.startswith("win"):
+    import asyncio.windows_events
+
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+import psycopg
+from psycopg import sql
+from psycopg.types.json import Jsonb
+from psycopg.rows import namedtuple_row
+from psycopg_pool import AsyncConnectionPool, PoolTimeout
+
+class AgensgraphQueryException(Exception):
+    """Exception for the Agensgraph queries."""
+
+    def __init__(self, exception: Union[str, Dict]) -> None:
+        if isinstance(exception, dict):
+            self.message = exception["message"] if "message" in exception else "unknown"
+            self.details = exception["details"] if "details" in exception else "unknown"
+        else:
+            self.message = exception
+            self.details = "unknown"
+
+    def get_message(self) -> str:
+        return self.message
+
+    def get_details(self) -> Any:
+        return self.details
+
+
+@final
+@dataclass
+class AgensgraphStorage(BaseGraphStorage):
+    vertex_regex: Pattern = re.compile(r"(\w+)\[(\d+\.\d+)\](\{.*\})")
+    edge_regex: Pattern = re.compile(r"(\w+)\[(\d+\.\d+)\]\[(\d+\.\d+),\s*(\d+\.\d+)\](\{.*\})")
+
+    @staticmethod
+    def load_nx_graph(file_name):
+        print("no preloading of graph with Agensgraph in production")
+
+    def __init__(self, namespace, global_config, embedding_func):
+        super().__init__(
+            namespace=namespace,
+            global_config=global_config,
+            embedding_func=embedding_func,
+        )
+        self._driver = None
+        self._driver_lock = asyncio.Lock()
+        DB = os.environ["AGENSGRAPH_DB"]
+        USER = os.environ["AGENSGRAPH_USER"]
+        PASSWORD = os.environ["AGENSGRAPH_PASSWORD"]
+        HOST = os.environ.get("AGENSGRAPH_HOST", "localhost")
+        PORT = os.environ.get("AGENSGRAPH_PORT", "5432")
+        self.graph_name = namespace or os.environ.get("AGENSGRAPH_GRAPHNAME", "lightrag")
+
+        connection_string = f"dbname='{DB}' user='{USER}' password='{PASSWORD}' host='{HOST}' port={PORT}"
+
+        self._driver = AsyncConnectionPool(connection_string, open=False)
+
+        return None
+
+    async def initialize(self):
+        """
+        Initialize the Agensgraph storage by creating a connection pool.
+        """
+        if self._driver is None:
+            raise AgensgraphQueryException("Agensgraph driver is not initialized")
+
+        # create graph and set graph_path
+        async with self._driver_lock:
+            try:
+                await self._driver.open()
+            except psycopg.errors.InvalidSchemaName as e:
+                raise AgensgraphQueryException(
+                    f"Failed to open connection to Agensgraph: {str(e)}"
+                ) from e
+
+        async with self._get_pool_connection() as conn:
+            async with conn.cursor() as curs:
+                try:
+                    await curs.execute(sql.SQL("CREATE GRAPH IF NOT EXISTS {}").format(sql.Identifier(self.graph_name)))
+                    await curs.execute(sql.SQL("SET graph_path = {}").format(sql.Identifier(self.graph_name)))
+                    await curs.execute('CREATE VLABEL IF NOT EXISTS base')
+                    await curs.execute('CREATE ELABEL IF NOT EXISTS "DIRECTED"')
+                    await curs.execute('CREATE PROPERTY INDEX IF NOT EXISTS base_entity_idx ON base (entity_id)')
+                    await conn.commit()
+                except (
+                    psycopg.errors.InvalidSchemaName,
+                    psycopg.errors.UniqueViolation,
+                ):
+                    await conn.rollback()
+                    logger.warning(
+                        f"Graph {self.graph_name} already exists or could not be created."
+                    )
+                except psycopg.Error as e:
+                    await conn.rollback()
+                    raise AgensgraphQueryException(
+                        f"Error initializing graph {self.graph_name}: {str(e)}"
+                    ) from e
+            
+        logger.info(f"Agensgraph storage initialized for graph: {self.graph_name}")
+
+    async def finalize(self):
+        """Close the Agensgraph driver and release all resources"""
+        if self._driver:
+            await self._driver.close()
+            self._driver = None
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.finalize()
+
+    async def index_done_callback(self) -> None:
+        # Agensgraph handles persistence automatically
+        pass
+
+    async def has_node(self, node_id: str) -> bool:
+        """
+        Check if a node with the given label exists in the database
+
+        Args:
+            node_id: Label of the node to check
+
+        Returns:
+            bool: True if node exists, False otherwise
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
+        query = """
+                MATCH (n:base {entity_id: %(node_id)s})
+                RETURN true AS node_exists LIMIT 1
+                """
+        single_result = (await self._query(query, {"node_id": Jsonb(node_id)}))[0]
+        logger.debug(
+            "{%s}:query:{%s}:result:{%s}",
+            inspect.currentframe().f_code.co_name,
+            query,
+            single_result["node_exists"],
+        )
+
+        return single_result["node_exists"]
+
+    async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
+        """
+        Check if an edge exists between two nodes
+
+        Args:
+            source_node_id: Label of the source node
+            target_node_id: Label of the target node
+
+        Returns:
+            bool: True if edge exists, False otherwise
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
+        query = """
+                MATCH (a:base {entity_id: %(source_node_id)s})-[r]-(b:base {entity_id: %(target_node_id)s})
+                RETURN true AS "edgeExists" LIMIT 1
+                """
+        single_result = (await self._query(query, {
+            "source_node_id": Jsonb(source_node_id),
+            "target_node_id": Jsonb(target_node_id),
+        }))[0]
+        logger.debug(
+            "{%s}:query:{%s}:result:{%s}",
+            inspect.currentframe().f_code.co_name,
+            query,
+            single_result["edgeExists"],
+        )
+        return single_result["edgeExists"]
+
+    async def get_node(self, node_id: str) -> dict[str, str] | None:
+        """Get node by its label identifier, return only node properties
+
+        Args:
+            node_id: The node label to look up
+
+        Returns:
+            dict: Node properties if found
+            None: If node not found
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
+        query = """
+                MATCH (n:base {entity_id: %(node_id)s})
+                RETURN n
+                """
+        records = await self._query(query, {"node_id": Jsonb(node_id)})
+        if records:
+            # warn if there are multiple records returned
+            if len(records) > 1:
+                logger.warning(
+                    "Multiple nodes found for entity_id '%s'. Returning first result.",
+                    node_id,
+                )
+            node_dict = records[0]["n"]
+            logger.debug(
+                "{%s}: query: {%s}, result: {%s}",
+                inspect.currentframe().f_code.co_name,
+                query,
+                node_dict,
+            )
+            # Return the node properties as a dictionary
+            return node_dict
+
+        return None
+
+    async def get_nodes_batch(self, node_ids: list[str]) -> dict[str, dict]:
+        """
+        Retrieve multiple nodes in one query using UNWIND.
+
+        Args:
+            node_ids: List of node entity IDs to fetch.
+
+        Returns:
+            A dictionary mapping each node_id to its node data (or None if not found).
+        """
+        query = """
+                UNWIND %(node_ids)s AS id
+                MATCH (n:base {entity_id: id})
+                RETURN n.entity_id AS entity_id, n
+                """
+        records = await self._query(query, {"node_ids": Jsonb(node_ids)})
+        nodes = {}
+        if records:
+            for record in records:
+                entity_id = record["entity_id"]
+                node_dict = record["n"]
+                logger.debug(
+                    "{%s}: query: {%s}, result: {%s}",
+                    inspect.currentframe().f_code.co_name,
+                    query,
+                    node_dict,
+                )
+                # Return a dictionary with entity_id as key
+                nodes[entity_id] = node_dict
+            return nodes
+        return None
+
+    async def node_degree(self, node_id: str) -> int:
+        """Get the degree (number of relationships) of a node with the given label.
+        If multiple nodes have the same label, returns the degree of the first node.
+        If no node is found, returns 0.
+
+        Args:
+            node_id: The label of the node
+
+        Returns:
+            int: The number of relationships the node has, or 0 if no node found
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
+        query = """
+                MATCH (n:base {entity_id: %(node_id)s})
+                OPTIONAL MATCH (n)-[r]-()
+                RETURN COUNT(r) AS degree
+                """
+        record = (await self._query(query, {"node_id": Jsonb(node_id)}))[0]
+        if record:
+            edge_count = int(record["degree"])
+            logger.debug(
+                "{%s}:query:{%s}:result:{%s}",
+                inspect.currentframe().f_code.co_name,
+                query,
+                edge_count,
+            )
+            return edge_count
+        else:
+            logger.warning(f"No node found with label '{self.escape_str(node_id)}'")
+            return 0
+
+    async def node_degrees_batch(self, node_ids: list[str]) -> dict[str, int]:
+        """
+        Retrieve the degree for multiple nodes in a single query using UNWIND.
+
+        Args:
+            node_ids: List of node labels (entity_id values) to look up.
+
+        Returns:
+            A dictionary mapping each node_id to its degree (number of relationships).
+            If a node is not found, its degree will be set to 0.
+        """
+        query = """
+                UNWIND %(node_ids)s AS id
+                MATCH (n:base {entity_id: id})
+                OPTIONAL MATCH (n)-[r]-()
+                RETURN n.entity_id AS entity_id, count(r) AS degree
+                """
+        records = (await self._query(query, {"node_ids": Jsonb(node_ids)}))
+
+        if records:
+            degrees = {}
+            for record in records:
+                entity_id = record["entity_id"]
+                degree = int(record["degree"])
+                degrees[entity_id] = degree
+                logger.debug(
+                    "{%s}: query: {%s}, result: {%s}",
+                    inspect.currentframe().f_code.co_name,
+                    query,
+                    degrees,
+                )
+            # For any node_id that did not return a record, set degree to 0.
+            for nid in node_ids:
+                if nid not in degrees:
+                    logger.warning(f"No node found with label '{nid}'")
+                    degrees[nid] = 0
+            return degrees
+        else:
+            logger.warning("No nodes found for the provided labels.")
+            return {nid: 0 for nid in node_ids}
+
+    async def edge_degree(self, src_id: str, tgt_id: str) -> int:
+        """Get the total degree (sum of relationships) of two nodes.
+
+        Args:
+            src_id: Label of the source node
+            tgt_id: Label of the target node
+
+        Returns:
+            int: Sum of the degrees of both nodes
+        """
+        src_degree = await self.node_degree(src_id)
+        trg_degree = await self.node_degree(tgt_id)
+
+        # Convert None to 0 for addition
+        src_degree = 0 if src_degree is None else src_degree
+        trg_degree = 0 if trg_degree is None else trg_degree
+
+        degrees = int(src_degree) + int(trg_degree)
+        logger.debug(
+            "{%s}:query:src_Degree+trg_degree:result:{%s}",
+            inspect.currentframe().f_code.co_name,
+            degrees,
+        )
+        return degrees
+    
+    async def edge_degrees_batch(
+        self, edge_pairs: list[tuple[str, str]]
+    ) -> dict[tuple[str, str], int]:
+        """
+        Calculate the combined degree for each edge (sum of the source and target node degrees)
+        in batch using the already implemented node_degrees_batch.
+
+        Args:
+            edge_pairs: List of (src, tgt) tuples.
+
+        Returns:
+            A dictionary mapping each (src, tgt) tuple to the sum of their degrees.
+        """
+        # Collect unique node IDs from all edge pairs.
+        unique_node_ids = {src for src, _ in edge_pairs}
+        unique_node_ids.update({tgt for _, tgt in edge_pairs})
+
+        # Get degrees for all nodes in one go.
+        degrees = await self.node_degrees_batch(list(unique_node_ids))
+
+        # Sum up degrees for each edge pair.
+        edge_degrees = {}
+        for src, tgt in edge_pairs:
+            edge_degrees[(src, tgt)] = degrees.get(src, 0) + degrees.get(tgt, 0)
+        
+        logger.debug(
+            "{%s}:query:edge_degrees_batch:result:{%s}",
+            inspect.currentframe().f_code.co_name,
+            edge_degrees,
+        )
+        return edge_degrees
+
+    async def get_edge(
+        self, source_node_id: str, target_node_id: str
+    ) -> dict[str, str] | None:
+        """Get edge properties between two nodes.
+
+        Args:
+            source_node_id: Label of the source node
+            target_node_id: Label of the target node
+
+        Returns:
+            dict: Edge properties if found, default properties if not found or on error
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
+        query = """
+                MATCH (start:base {entity_id: %(source_node_id)s})-[r]-("end":base {entity_id: %(target_node_id)s})
+                RETURN properties(r) as edge_properties
+                """
+        records = await self._query(query, {
+            "source_node_id": Jsonb(source_node_id),
+            "target_node_id": Jsonb(target_node_id),
+        })
+
+        if records:
+            if len(records) > 1:
+                logger.warning(
+                    "Multiple edges found between '%s' and '%s'. Returning first result.",
+                    self.escape_str(source_node_id),
+                    self.escape_str(target_node_id),
+                )
+            edge_result = records[0]["edge_properties"]
+
+            required_keys = {
+                "weight": 0.0,
+                "source_id": None,
+                "description": None,
+                "keywords": None,
+            }
+            for key, default_value in required_keys.items():
+                if key not in edge_result:
+                    edge_result[key] = default_value
+                    logger.warning(
+                        f"Edge between {self.escape_str(source_node_id)} and {self.escape_str(target_node_id)} "
+                        f"missing {key}, using default: {default_value}"
+                    )
+            logger.debug(
+                "{%s}:query:{%s}:result:{%s}",
+                inspect.currentframe().f_code.co_name,
+                query,
+                edge_result,
+            )
+            return edge_result
+        else:
+            logger.warning(
+                "No edge found between '%s' and '%s'. Returning default properties.",
+                self.escape_str(source_node_id),
+                self.escape_str(target_node_id),
+            )
+            # Return None when no edge found
+            return None
+
+    async def get_edges_batch(
+        self, pairs: list[dict[str, str]]
+    ) -> dict[tuple[str, str], dict]:
+        """
+        Retrieve edge properties for multiple (src, tgt) pairs in one query.
+
+        Args:
+            pairs: List of dictionaries, e.g. [{"src": "node1", "tgt": "node2"}, ...]
+
+        Returns:
+            A dictionary mapping (src, tgt) tuples to their edge properties.
+        """
+        query = """
+                UNWIND %(pairs)s AS pair
+                MATCH (start:base {entity_id: pair.src})-[r:"DIRECTED"]-("end":base {entity_id: pair.tgt})
+                RETURN pair.src AS src_id, pair.tgt AS tgt_id, collect(properties(r)) AS edges
+                """
+        records = await self._query(query, {"pairs": Jsonb(pairs)})
+        edges_dict = {}
+        if records:
+            for record in records:
+                src = record["src_id"]
+                tgt = record["tgt_id"]
+                edges = record["edges"]
+                if edges and len(edges) > 0:
+                    edge_props = edges[0]  # choose the first if multiple exist
+                    # Ensure required keys exist with defaults
+                    for key, default in {
+                        "weight": 0.0,
+                        "source_id": None,
+                        "description": None,
+                        "keywords": None,
+                    }.items():
+                        if key not in edge_props:
+                            edge_props[key] = default
+                    edges_dict[(src, tgt)] = edge_props
+                else:
+                    edges_dict[(src, tgt)] = {
+                        "weight": 0.0,
+                        "source_id": None,
+                        "description": None,
+                        "keywords": None,
+                    }
+            logger.debug(
+                "{%s}:query:{%s}:result:{%s}",
+                inspect.currentframe().f_code.co_name,
+                query,
+                edges_dict,
+            )
+            return edges_dict
+        else:
+            logger.warning("No edges found for the provided pairs.")
+            return edges_dict
+    
+    async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None:
+        """Retrieves all edges (relationships) for a particular node identified by its label.
+
+        Args:
+            source_node_id: Label of the node to get edges for
+
+        Returns:
+            list[tuple[str, str]]: List of (source_label, target_label) tuples representing edges
+            None: If no edges found
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
+        query = """
+                MATCH (n:base {entity_id: %(source_node_id)s})
+                OPTIONAL MATCH (n)-[r]-(connected:base)
+                WHERE connected.entity_id IS NOT NULL
+                RETURN n, r, connected
+                """
+        results = await self._query(query, {"source_node_id": Jsonb(source_node_id)})
+        if results:
+            edges = []
+            for record in results:
+                source_node = record["n"] if record["n"] else None
+                connected_node = record["connected"] if record["connected"] else None
+
+                if not source_node or not connected_node:
+                    continue
+
+                source_label = (
+                    source_node.get("entity_id")
+                    if source_node.get("entity_id")
+                    else None
+                )
+                target_label = (
+                    connected_node.get("entity_id")
+                    if connected_node.get("entity_id")
+                    else None
+                )
+
+                if source_label and target_label:
+                    edges.append((source_label, target_label))
+        else:
+            logger.warning(f"No edges found for node with label '{source_node_id}'")
+            return None
+
+        logger.debug(
+            "{%s}:query:{%s}:result:{%s}",
+            inspect.currentframe().f_code.co_name,
+            query,
+            edges,
+        )
+        return edges
+
+    async def get_nodes_edges_batch(
+        self, node_ids: list[str]
+    ) -> dict[str, list[tuple[str, str]]]:
+        """
+        Batch retrieve edges for multiple nodes in one query using UNWIND.
+        For each node, returns both outgoing and incoming edges to properly represent
+        the undirected graph nature.
+
+        Args:
+            node_ids: List of node IDs (entity_id) for which to retrieve edges.
+
+        Returns:
+            A dictionary mapping each node ID to its list of edge tuples (source, target).
+            For each node, the list includes both:
+            - Outgoing edges: (queried_node, connected_node)
+            - Incoming edges: (connected_node, queried_node)
+        """
+        # Query to get both outgoing and incoming edges
+        query = """
+                UNWIND %(node_ids)s AS id
+                MATCH (n:base {entity_id: id})
+                OPTIONAL MATCH (n)-[r]-(connected:base)
+                RETURN id AS queried_id, n.entity_id AS node_entity_id,
+                        connected.entity_id AS connected_entity_id,
+                        startNode(r).entity_id AS start_entity_id
+                """
+        records = await self._query(query, {"node_ids": Jsonb(node_ids)})
+
+        # Initialize the dictionary with empty lists for each node ID
+        edges_dict = {node_id: [] for node_id in node_ids}
+
+        if records:
+            for record in records:
+                queried_id = record["queried_id"]
+                node_entity_id = record["node_entity_id"]
+                connected_entity_id = record["connected_entity_id"]
+                start_entity_id = record["start_entity_id"]
+
+                # Skip if either node is None
+                if not node_entity_id or not connected_entity_id:
+                    continue
+
+                # Determine the actual direction of the edge
+                # If the start node is the queried node, it's an outgoing edge
+                # Otherwise, it's an incoming edge
+                if start_entity_id == node_entity_id:
+                    # Outgoing edge: (queried_node -> connected_node)
+                    edges_dict[queried_id].append((node_entity_id, connected_entity_id))
+                else:
+                    # Incoming edge: (connected_node -> queried_node)
+                    edges_dict[queried_id].append((connected_entity_id, node_entity_id))
+        else:
+            logger.warning("No edges found for the provided node IDs.")
+            # If no edges found, return empty lists for each node ID
+            return edges_dict
+        
+        logger.debug(
+            "{%s}:query:{%s}:result:{%s}",
+            inspect.currentframe().f_code.co_name,
+            query,
+            edges_dict,
+        )
+        return edges_dict
+    
+    async def get_nodes_by_chunk_ids(self, chunk_ids: list[str]) -> list[dict]:
+        query = """
+                UNWIND %(chunk_ids)s AS chunk_id
+                MATCH (n:base)
+                WHERE n.source_id IS NOT NULL AND chunk_id <@ split(n.source_id, {GRAPH_FIELD_SEP})::jsonb
+                RETURN DISTINCT n
+                """
+
+        results = await self._query(sql.SQL(query).format(
+            GRAPH_FIELD_SEP=sql.Literal(GRAPH_FIELD_SEP)
+        ), {"chunk_ids": Jsonb(chunk_ids)})
+        nodes = []
+        for record in results:
+            node_dict = record["n"]
+            # Add node id (entity_id) to the dictionary for easier access
+            node_dict["id"] = node_dict.get("entity_id")
+            nodes.append(node_dict)
+        return nodes
+    
+    async def get_edges_by_chunk_ids(self, chunk_ids: list[str]) -> list[dict]:
+        query = """
+                UNWIND %(chunk_ids)s AS chunk_id
+                MATCH (a:base)-[r]-(b:base)
+                WHERE r.source_id IS NOT NULL AND chunk_id <@ split(r.source_id, {GRAPH_FIELD_SEP})::jsonb
+                RETURN DISTINCT a.entity_id AS source, b.entity_id AS target, properties(r) AS properties
+                """
+
+        results = await self._query(sql.SQL(query).format(
+            GRAPH_FIELD_SEP=sql.Literal(GRAPH_FIELD_SEP)
+        ), {"chunk_ids": Jsonb(chunk_ids)})
+        edges = []
+        for record in results:
+            edge_properties = record["properties"]
+            edge_properties["source"] = record["source"]
+            edge_properties["target"] = record["target"]
+            edges.append(edge_properties)
+        return edges
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        retry=retry_if_exception_type((AgensgraphQueryException,)),
+    )
+    async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
+        """
+        Upsert a node in the Agensgraph database.
+
+        Args:
+            node_id: The unique identifier for the node (used as label)
+            node_data: Dictionary of node properties
+        """
+        query = """
+                MERGE (n:base {entity_id: %(node_id)s})
+                SET n += %(node_data)s
+                """
+        try:
+            await self._query(query, {
+                "node_id": Jsonb(node_id),
+                "node_data": Jsonb(node_data),
+            })
+            logger.debug(
+                "Upserted node with node_id '{%s}' and properties: {%s}",
+                node_id,
+                node_data,
+            )
+        except Exception as e:
+            logger.error("Error during upsert: {%s}", e)
+            raise
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        retry=retry_if_exception_type((AgensgraphQueryException,)),
+    )
+    async def upsert_edge(
+        self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
+    ) -> None:
+        """
+        Upsert an edge and its properties between two nodes identified by their labels.
+        Ensures both source and target nodes exist and are unique before creating the edge.
+        Uses entity_id property to uniquely identify nodes.
+
+        Args:
+            source_node_id (str): Label of the source node (used as identifier)
+            target_node_id (str): Label of the target node (used as identifier)
+            edge_data (dict): Dictionary of properties to set on the edge
+        """
+        query = """
+                MATCH (source:base {entity_id: %(source_node_id)s})
+                WITH source
+                MATCH (target:base {entity_id: %(target_node_id)s})
+                MERGE (source)-[r:"DIRECTED"]-(target)
+                SET r += %(edge_data)s
+                RETURN r, source, target
+                """
+        try:
+            await self._query(query, {
+                "source_node_id": Jsonb(source_node_id),
+                "target_node_id": Jsonb(target_node_id),
+                "edge_data": Jsonb(edge_data),
+            })
+            logger.debug(
+                "Upserted edge from '{%s}' to '{%s}' with properties: {%s}",
+                source_node_id,
+                target_node_id,
+                edge_data,
+            )
+        except Exception as e:
+            logger.error("Error during edge upsert: {%s}", e)
+            raise
+
+    async def get_knowledge_graph(
+        self, node_label: str, max_depth: int = 3, max_nodes: int = 1000
+    ) -> KnowledgeGraph:
+        """
+        Retrieve a connected subgraph of nodes where the label includes the specified `node_label`.
+
+        Args:
+            node_label: Label of the starting node, * means all nodes
+            max_depth: Maximum depth of the subgraph, Defaults to 3
+            max_nodes: Maximum nodes to return by BFS, Defaults to 1000
+
+        Returns:
+            KnowledgeGraph object containing nodes and edges, with an is_truncated flag
+            indicating whether the graph was truncated due to max_nodes limit
+        """
+        from collections import deque
+
+        result = KnowledgeGraph()
+        visited_nodes = set()
+        visited_edges = set()
+        visited_edge_pairs = set()
+        queue = deque()
+
+        # Step 1: Get starting nodes
+        if node_label == "*":
+            query = """
+                    MATCH (n:base)
+                    RETURN DISTINCT id(n) AS node_id, n
+                    LIMIT %(max_nodes)s
+                    """
+            node_results = await self._query(query, {"max_nodes": max_nodes})
+        else:
+            query = """
+                    MATCH (n:base {entity_id: %(node_label)s})
+                    RETURN id(n) AS node_id, n
+                    """
+            node_results = await self._query(query, {"node_label": Jsonb(node_label)})
+
+        for record in node_results:
+            node_data = record["n"]
+            if not node_data.get("entity_id"):
+                continue
+            start_node = KnowledgeGraphNode(
+                id=str(node_data["entity_id"]),
+                labels=[node_data["entity_id"]],
+                properties=node_data,
+            )
+            queue.append((start_node, None, 0))
+
+        # Step 2: BFS traversal
+        while queue and len(visited_nodes) < max_nodes:
+            current_node, current_edge, current_depth = queue.popleft()
+
+            if current_node.id in visited_nodes or current_depth > max_depth:
+                continue
+
+            result.nodes.append(current_node)
+            visited_nodes.add(current_node.id)
+
+            if current_edge and current_edge.id not in visited_edges:
+                result.edges.append(current_edge)
+                visited_edges.add(current_edge.id)
+
+            if len(visited_nodes) >= max_nodes:
+                result.is_truncated = True
+                break
+
+            # Step 3: Query neighbors
+            query = """
+            MATCH (a:base {entity_id: %(current_node_id)s})-[r]-(b:base)
+            RETURN type(r) as rel_type, properties(r) as r, b, id(r) AS edge_id, id(b) AS target_id
+            """
+            records = await self._query(query, {"current_node_id": Jsonb(current_node.id)})
+
+            for record in records:
+                rel_type = record["rel_type"]
+                rel = record["r"]
+                b_node = record["b"]
+                edge_id = str(record["edge_id"])
+                target_id = b_node.get("entity_id")
+
+                if not target_id:
+                    continue
+
+                target_node = KnowledgeGraphNode(
+                    id=str(target_id),
+                    labels=[target_id],
+                    properties=b_node,
+                )
+
+                target_edge = KnowledgeGraphEdge(
+                    id=edge_id,
+                    type=rel_type,
+                    source=current_node.id,
+                    target=target_id,
+                    properties=rel,
+                )
+
+                sorted_pair = tuple(sorted([current_node.id, target_id]))
+                if sorted_pair not in visited_edge_pairs:
+                    if (
+                        target_id in visited_nodes or
+                        (target_id not in visited_nodes and current_depth < max_depth)
+                    ):
+                        result.edges.append(target_edge)
+                        visited_edges.add(edge_id)
+                        visited_edge_pairs.add(sorted_pair)
+
+                if target_id not in visited_nodes and current_depth < max_depth:
+                    queue.append((target_node, None, current_depth + 1))
+
+        return result
+
+    async def get_all_labels(self) -> list[str]:
+        """Get all node labels in the database
+
+        Returns:
+            ["label1", "label2", ...]  # Alphabetically sorted label list
+        """
+        query = """
+                MATCH (n:base)
+                WHERE n.entity_id IS NOT NULL
+                WITH DISTINCT n.entity_id AS label
+                ORDER BY label
+                RETURN collect(label) AS node_labels
+                """
+        results = await self._query(query)
+
+        if not results:
+            logger.warning("No labels found in the graph.")
+            return []
+
+        labels = results[0]["node_labels"]
+        logger.debug(
+            "{%s}:query:{%s}:result:{%s}",
+            inspect.currentframe().f_code.co_name,
+            query,
+            labels,
+        )
+        return labels
+
+    async def delete_node(self, node_id: str) -> None:
+        """Delete a node with the specified label
+
+        Args:
+            node_id: The label of the node to delete
+        """
+        query = """
+                MATCH (n:base {entity_id: %(node_id)s})
+                DETACH DELETE n
+                """
+        try:
+            await self._query(query, {"node_id": Jsonb(node_id)})
+            logger.debug(f"Deleted node with label '{self.escape_str(node_id)}'")
+        except Exception as e:
+            logger.error(f"Error during node deletion: {str(e)}")
+            raise
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        retry=retry_if_exception_type((AgensgraphQueryException,)),
+    )
+    async def remove_nodes(self, nodes: list[str]):
+        """Delete multiple nodes
+
+        Args:
+            nodes: List of node labels to be deleted
+        """
+        for node in nodes:
+            await self.delete_node(node)
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        retry=retry_if_exception_type((AgensgraphQueryException,)),
+    )
+    async def remove_edges(self, edges: list[tuple[str, str]]):
+        """Delete multiple edges
+
+        Args:
+            edges: List of edges to be deleted, each edge is a (source, target) tuple
+        """
+        for source, target in edges:
+            query = """
+                    MATCH (source:base {entity_id: %(source)s})-[r]-(target:base {entity_id: %(target)s})
+                    DELETE r
+                    """
+            try:
+                await self._query(query, {
+                    "source": Jsonb(source),
+                    "target": Jsonb(target),
+                })
+                logger.debug(
+                    f"Deleted edge from '{self.escape_str(source)}' to '{self.escape_str(target)}'"
+                )
+            except Exception as e:
+                logger.error(f"Error during edge deletion: {str(e)}")
+                raise
+
+    async def drop(self) -> dict[str, str]:
+        """Drop the storage by removing all nodes and relationships in the graph.
+
+        Returns:
+            dict[str, str]: Status of the operation with keys 'status' and 'message'
+        """
+        try:
+            query = """
+                    MATCH (n)
+                    DETACH DELETE n
+                    """
+            await self._query(query)
+            logger.info(f"Successfully dropped all data from graph {self.graph_name}")
+            return {"status": "success", "message": "graph data dropped"}
+        except Exception as e:
+            logger.error(f"Error dropping graph {self.graph_name}: {e}")
+            return {"status": "error", "message": str(e)}
+
+    @staticmethod
+    def _record_to_dict(record: NamedTuple) -> Dict[str, Any]:
+        """
+        Convert a record returned from an agensgraph query to a dictionary
+
+        Args:
+            record (): a record from an agensgraph query result
+
+        Returns:
+            Dict[str, Any]: a dictionary representation of the record where
+                the dictionary key is the field name and the value is the
+                value converted to a python type
+        """
+        # result holder
+        d = {}
+
+        # prebuild a mapping of vertex_id to vertex mappings to be used
+        # later to build edges
+        vertices = {}
+        for k in record._fields:
+            v = getattr(record, k)
+
+            # records comes back label[id]{properties} which must be parsed
+            if isinstance(v, str):
+                vertex = AgensgraphStorage.vertex_regex.match(v)
+                if vertex:
+                    label, vertex_id, properties = vertex.groups()
+                    properties = json.loads(properties)
+                    vertices[str(vertex_id)] = properties
+
+        # iterate returned fields and parse appropriately
+        for k in record._fields:
+            v = getattr(record, k)
+
+            if isinstance(v, str):
+                vertex = AgensgraphStorage.vertex_regex.match(v)
+                edge = AgensgraphStorage.edge_regex.match(v)
+
+                if vertex:
+                    d[k] = json.loads(vertex.group(3))
+                elif edge:
+                    elabel, edge_id, start_id, end_id, properties = edge.groups()
+                    d[k] = (
+                        vertices.get(start_id, {}),
+                        elabel,
+                        vertices.get(end_id, {}),
+                    )
+                else:
+                    d[k] = v
+
+            else:
+                d[k] = v
+
+        return d
+
+    @staticmethod
+    def escape_str(val: str) -> str:
+        return val.replace("'", "''").replace("\\", "\\\\").replace('"', '\\"')
+
+    async def _query(self, query: str, params: Dict = {}) -> List[Dict[str, Any]]:
+        """
+        Query the graph by taking a cypher query, converting it to an
+        age compatible query, executing it and converting the result
+
+        Args:
+            query (str): a cypher query to be executed
+            params (dict): parameters for the query
+
+        Returns:
+            List[Dict[str, Any]]: a list of dictionaries containing the result set
+        """
+        await self._driver.open()
+
+        # execute the query, rolling back on an error
+        async with self._get_pool_connection() as conn:
+            async with conn.cursor(row_factory=namedtuple_row) as curs:
+                try:
+                    await curs.execute(sql.SQL("SET graph_path = {}").format(sql.Identifier(self.graph_name)))
+                    await curs.execute(query, params)
+                    await conn.commit()
+                except psycopg.Error as e:
+                    await conn.rollback()
+                    raise AgensgraphQueryException(
+                        {
+                            "message": f"Error executing graph query: {query}",
+                            "detail": str(e),
+                        }
+                    ) from e
+                try:
+                    data = await curs.fetchall()
+                except psycopg.ProgrammingError:
+                    data = []  # Handle queries that don’t return data
+                if data is None:
+                    result = []
+                # decode records
+                else:
+                    result = [AgensgraphStorage._record_to_dict(d) for d in data]
+
+                return result
+
+    @asynccontextmanager
+    async def _get_pool_connection(self, timeout: Optional[float] = None):
+        """Workaround for a psycopg_pool bug"""
+
+        try:
+            connection = await self._driver.getconn(timeout=timeout)
+        except PoolTimeout:
+            await self._driver._add_connection(None)  # workaround...
+            connection = await self._driver.getconn(timeout=timeout)
+
+        try:
+            async with connection:
+                yield connection
+        finally:
+            await self._driver.putconn(connection)

lightrag_agensgraph-0.1.0/pyproject.toml

@@ -0,0 +1,21 @@
+[build-system]
+build-backend = "poetry.core.masonry.api"
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+
+[tool.poetry]
+name = "lightrag-agensgraph"
+version = "0.1.0"
+description = "Lightrag graph store implementation for agensgraph"
+authors = ["Muhammad Taha Naveed <skaisw@skaiworldwide.com>"]
+readme = "README.md"
+homepage = "https://github.com/skaiworldwide-oss/agensgraph-ai/tree/main/lightrag"
+repository = "https://github.com/skaiworldwide-oss/agensgraph-ai/tree/main/lightrag"
+keywords = ["lightrag", "agensgraph", "graph store", "integration"]
+
+[tool.poetry.dependencies]
+python = ">=3.9,<4.0"
+psycopg = { version = ">=3.1.0", extras = ["binary", "pool"] }
+lightrag-hku = ">=1.3.9"
+
+[[tool.poetry.packages]]
+include = "lightrag_agensgraph"