Add Apache AGE graph storage

examples/lightrag_ollama_age_demo.py

@@ -0,0 +1,80 @@
+import asyncio
+import inspect
+import logging
+import os
+
+from lightrag import LightRAG, QueryParam
+from lightrag.llm import ollama_embedding, ollama_model_complete
+from lightrag.utils import EmbeddingFunc
+
+WORKING_DIR = "./dickens_age"
+
+logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO)
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+# AGE
+os.environ["AGE_POSTGRES_DB"] = "postgresDB"
+os.environ["AGE_POSTGRES_USER"] = "postgresUser"
+os.environ["AGE_POSTGRES_PASSWORD"] = "postgresPW"
+os.environ["AGE_POSTGRES_HOST"] = "localhost"
+os.environ["AGE_POSTGRES_PORT"] = "5455"
+os.environ["AGE_GRAPH_NAME"] = "dickens"
+
+rag = LightRAG(
+    working_dir=WORKING_DIR,
+    llm_model_func=ollama_model_complete,
+    llm_model_name="llama3.1:8b",
+    llm_model_max_async=4,
+    llm_model_max_token_size=32768,
+    llm_model_kwargs={"host": "http://localhost:11434", "options": {"num_ctx": 32768}},
+    embedding_func=EmbeddingFunc(
+        embedding_dim=768,
+        max_token_size=8192,
+        func=lambda texts: ollama_embedding(
+            texts, embed_model="nomic-embed-text", host="http://localhost:11434"
+        ),
+    ),
+    graph_storage="AGEStorage",
+)
+
+with open("./book.txt", "r", encoding="utf-8") as f:
+    rag.insert(f.read())
+
+# Perform naive search
+print(
+    rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
+)
+
+# Perform local search
+print(
+    rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
+)
+
+# Perform global search
+print(
+    rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
+)
+
+# Perform hybrid search
+print(
+    rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
+)
+
+# stream response
+resp = rag.query(
+    "What are the top themes in this story?",
+    param=QueryParam(mode="hybrid", stream=True),
+)
+
+
+async def print_stream(stream):
+    async for chunk in stream:
+        print(chunk, end="", flush=True)
+
+
+if inspect.isasyncgen(resp):
+    asyncio.run(print_stream(resp))
+else:
+    print(resp)

lightrag/kg/age_impl.py

@@ -0,0 +1,613 @@
+import asyncio
+import inspect
+import json
+import os
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import Any, Dict, List, NamedTuple, Optional, Tuple, Union
+
+import psycopg
+from psycopg.rows import namedtuple_row
+from psycopg_pool import AsyncConnectionPool, PoolTimeout
+from tenacity import (
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from lightrag.utils import logger
+
+from ..base import BaseGraphStorage
+
+
+class AGEQueryException(Exception):
+    """Exception for the AGE 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
+
+
+@dataclass
+class AGEStorage(BaseGraphStorage):
+    @staticmethod
+    def load_nx_graph(file_name):
+        print("no preloading of graph with AGE 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["AGE_POSTGRES_DB"].replace("\\", "\\\\").replace("'", "\\'")
+        USER = os.environ["AGE_POSTGRES_USER"].replace("\\", "\\\\").replace("'", "\\'")
+        PASSWORD = (
+            os.environ["AGE_POSTGRES_PASSWORD"]
+            .replace("\\", "\\\\")
+            .replace("'", "\\'")
+        )
+        HOST = os.environ["AGE_POSTGRES_HOST"].replace("\\", "\\\\").replace("'", "\\'")
+        PORT = int(os.environ["AGE_POSTGRES_PORT"])
+        self.graph_name = os.environ["AGE_GRAPH_NAME"]
+
+        connection_string = f"dbname='{DB}' user='{USER}' password='{PASSWORD}' host='{HOST}' port={PORT}"
+
+        self._driver = AsyncConnectionPool(connection_string, open=False)
+
+        return None
+
+    def __post_init__(self):
+        self._node_embed_algorithms = {
+            "node2vec": self._node2vec_embed,
+        }
+
+    async def close(self):
+        if self._driver:
+            await self._driver.close()
+            self._driver = None
+
+    async def __aexit__(self, exc_type, exc, tb):
+        if self._driver:
+            await self._driver.close()
+
+    async def index_done_callback(self):
+        print("KG successfully indexed.")
+
+    @staticmethod
+    def _record_to_dict(record: NamedTuple) -> Dict[str, Any]:
+        """
+        Convert a record returned from an age query to a dictionary
+
+        Args:
+            record (): a record from an age 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)
+            # agtype comes back '{key: value}::type' which must be parsed
+            if isinstance(v, str) and "::" in v:
+                dtype = v.split("::")[-1]
+                v = v.split("::")[0]
+                if dtype == "vertex":
+                    vertex = json.loads(v)
+                    vertices[vertex["id"]] = vertex.get("properties")
+
+        # iterate returned fields and parse appropriately
+        for k in record._fields:
+            v = getattr(record, k)
+            if isinstance(v, str) and "::" in v:
+                dtype = v.split("::")[-1]
+                v = v.split("::")[0]
+            else:
+                dtype = ""
+
+            if dtype == "vertex":
+                vertex = json.loads(v)
+                field = json.loads(v).get("properties")
+                if not field:
+                    field = {}
+                field["label"] = AGEStorage._decode_graph_label(vertex["label"])
+                d[k] = field
+            # convert edge from id-label->id by replacing id with node information
+            # we only do this if the vertex was also returned in the query
+            # this is an attempt to be consistent with neo4j implementation
+            elif dtype == "edge":
+                edge = json.loads(v)
+                d[k] = (
+                    vertices.get(edge["start_id"], {}),
+                    edge[
+                        "label"
+                    ],  # we don't use decode_graph_label(), since edge label is always "DIRECTED"
+                    vertices.get(edge["end_id"], {}),
+                )
+            else:
+                d[k] = json.loads(v) if isinstance(v, str) else v
+
+        return d
+
+    @staticmethod
+    def _format_properties(
+        properties: Dict[str, Any], _id: Union[str, None] = None
+    ) -> str:
+        """
+        Convert a dictionary of properties to a string representation that
+        can be used in a cypher query insert/merge statement.
+
+        Args:
+            properties (Dict[str,str]): a dictionary containing node/edge properties
+            id (Union[str, None]): the id of the node or None if none exists
+
+        Returns:
+            str: the properties dictionary as a properly formatted string
+        """
+        props = []
+        # wrap property key in backticks to escape
+        for k, v in properties.items():
+            prop = f"`{k}`: {json.dumps(v)}"
+            props.append(prop)
+        if _id is not None and "id" not in properties:
+            props.append(
+                f"id: {json.dumps(_id)}" if isinstance(_id, str) else f"id: {_id}"
+            )
+        return "{" + ", ".join(props) + "}"
+
+    @staticmethod
+    def _encode_graph_label(label: str) -> str:
+        """
+        Since AGE suports only alphanumerical labels, we will encode generic label as HEX string
+
+        Args:
+            label (str): the original label
+
+        Returns:
+            str: the encoded label
+        """
+        return "x" + label.encode().hex()
+
+    @staticmethod
+    def _decode_graph_label(encoded_label: str) -> str:
+        """
+        Since AGE suports only alphanumerical labels, we will encode generic label as HEX string
+
+        Args:
+            encoded_label (str): the encoded label
+
+        Returns:
+            str: the decoded label
+        """
+        return bytes.fromhex(encoded_label.removeprefix("x")).decode()
+
+    @staticmethod
+    def _get_col_name(field: str, idx: int) -> str:
+        """
+        Convert a cypher return field to a pgsql select field
+        If possible keep the cypher column name, but create a generic name if necessary
+
+        Args:
+            field (str): a return field from a cypher query to be formatted for pgsql
+            idx (int): the position of the field in the return statement
+
+        Returns:
+            str: the field to be used in the pgsql select statement
+        """
+        # remove white space
+        field = field.strip()
+        # if an alias is provided for the field, use it
+        if " as " in field:
+            return field.split(" as ")[-1].strip()
+        # if the return value is an unnamed primitive, give it a generic name
+        if field.isnumeric() or field in ("true", "false", "null"):
+            return f"column_{idx}"
+        # otherwise return the value stripping out some common special chars
+        return field.replace("(", "_").replace(")", "")
+
+    @staticmethod
+    def _wrap_query(query: str, graph_name: str, **params: str) -> str:
+        """
+        Convert a cypher query to an Apache Age compatible
+        sql query by wrapping the cypher query in ag_catalog.cypher,
+        casting results to agtype and building a select statement
+
+        Args:
+            query (str): a valid cypher query
+            graph_name (str): the name of the graph to query
+            params (dict): parameters for the query
+
+        Returns:
+            str: an equivalent pgsql query
+        """
+
+        # pgsql template
+        template = """SELECT {projection} FROM ag_catalog.cypher('{graph_name}', $$
+            {query}
+        $$) AS ({fields});"""
+
+        # if there are any returned fields they must be added to the pgsql query
+        if "return" in query.lower():
+            # parse return statement to identify returned fields
+            fields = (
+                query.lower()
+                .split("return")[-1]
+                .split("distinct")[-1]
+                .split("order by")[0]
+                .split("skip")[0]
+                .split("limit")[0]
+                .split(",")
+            )
+
+            # raise exception if RETURN * is found as we can't resolve the fields
+            if "*" in [x.strip() for x in fields]:
+                raise ValueError(
+                    "AGE graph does not support 'RETURN *'"
+                    + " statements in Cypher queries"
+                )
+
+            # get pgsql formatted field names
+            fields = [
+                AGEStorage._get_col_name(field, idx) for idx, field in enumerate(fields)
+            ]
+
+            # build resulting pgsql relation
+            fields_str = ", ".join(
+                [field.split(".")[-1] + " agtype" for field in fields]
+            )
+
+        # if no return statement we still need to return a single field of type agtype
+        else:
+            fields_str = "a agtype"
+
+        select_str = "*"
+
+        query = query.format(**params)
+
+        return template.format(
+            graph_name=graph_name,
+            query=query,
+            fields=fields_str,
+            projection=select_str,
+        )
+
+    async def _query(self, query: str, **params: str) -> 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
+        """
+        # convert cypher query to pgsql/age query
+        wrapped_query = self._wrap_query(query, self.graph_name, **params)
+
+        await self._driver.open()
+
+        # create graph if it doesn't exist
+        async with self._get_pool_connection() as conn:
+            async with conn.cursor() as curs:
+                try:
+                    await curs.execute('SET search_path = ag_catalog, "$user", public')
+                    await curs.execute(f"SELECT create_graph('{self.graph_name}')")
+                    await conn.commit()
+                except (
+                    psycopg.errors.InvalidSchemaName,
+                    psycopg.errors.UniqueViolation,
+                ):
+                    await conn.rollback()
+
+        # 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('SET search_path = ag_catalog, "$user", public')
+                    await curs.execute(wrapped_query)
+                    await conn.commit()
+                except psycopg.Error as e:
+                    await conn.rollback()
+                    raise AGEQueryException(
+                        {
+                            "message": f"Error executing graph query: {query.format(**params)}",
+                            "detail": str(e),
+                        }
+                    ) from e
+
+                data = await curs.fetchall()
+                if data is None:
+                    result = []
+                # decode records
+                else:
+                    result = [AGEStorage._record_to_dict(d) for d in data]
+
+                return result
+
+    async def has_node(self, node_id: str) -> bool:
+        entity_name_label = node_id.strip('"')
+
+        query = "MATCH (n:`{label}`) RETURN count(n) > 0 AS node_exists"
+        single_result = (
+            await self._query(
+                query, label=AGEStorage._encode_graph_label(entity_name_label)
+            )
+        )[0]
+        logger.debug(
+            "{%s}:query:{%s}:result:{%s}",
+            inspect.currentframe().f_code.co_name,
+            query,
+            single_result[0],
+        )
+
+        return single_result["node_exists"].lower() == "true"
+
+    async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
+        entity_name_label_source = source_node_id.strip('"')
+        entity_name_label_target = target_node_id.strip('"')
+
+        query = (
+            "MATCH (a:`{src_label}`)-[r]-(b:`{tgt_label}`) "
+            "RETURN COUNT(r) > 0 AS edgeExists"
+        )
+        single_result = (
+            await self._query(
+                query,
+                src_label=AGEStorage._encode_graph_label(entity_name_label_source),
+                tgt_label=AGEStorage._encode_graph_label(entity_name_label_target),
+            )
+        )[0]
+        logger.debug(
+            "{%s}:query:{query}:result:{%s}",
+            inspect.currentframe().f_code.co_name,
+            single_result[0],
+        )
+        return single_result["edgeExists"].lower() == "true"
+
+    async def get_node(self, node_id: str) -> Union[dict, None]:
+        entity_name_label = node_id.strip('"')
+        query = "MATCH (n:`{label}`) RETURN n"
+        record = await self._query(
+            query, label=AGEStorage._encode_graph_label(entity_name_label)
+        )
+        if record:
+            node = record[0]
+            node_dict = node["n"]
+            logger.debug(
+                "{%s}: query: {%s}, result: {%s}",
+                inspect.currentframe().f_code.co_name,
+                query,
+                node_dict,
+            )
+            return node_dict
+        return None
+
+    async def node_degree(self, node_id: str) -> int:
+        entity_name_label = node_id.strip('"')
+
+        query = """
+            MATCH (n:`{label}`)-[]->(x)
+            RETURN count(x) AS total_edge_count
+        """
+        record = (
+            await self._query(
+                query, label=AGEStorage._encode_graph_label(entity_name_label)
+            )
+        )[0]
+        if record:
+            edge_count = int(record["total_edge_count"])
+            logger.debug(
+                "{%s}:query:{%s}:result:{%s}",
+                inspect.currentframe().f_code.co_name,
+                query,
+                edge_count,
+            )
+            return edge_count
+
+    async def edge_degree(self, src_id: str, tgt_id: str) -> int:
+        entity_name_label_source = src_id.strip('"')
+        entity_name_label_target = tgt_id.strip('"')
+        src_degree = await self.node_degree(entity_name_label_source)
+        trg_degree = await self.node_degree(entity_name_label_target)
+
+        # 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 get_edge(
+        self, source_node_id: str, target_node_id: str
+    ) -> Union[dict, None]:
+        """
+        Find all edges between nodes of two given labels
+
+        Args:
+            source_node_label (str): Label of the source nodes
+            target_node_label (str): Label of the target nodes
+
+        Returns:
+            list: List of all relationships/edges found
+        """
+        entity_name_label_source = source_node_id.strip('"')
+        entity_name_label_target = target_node_id.strip('"')
+
+        query = """
+        MATCH (a:`{src_label}`)-[r]->(b:`{tgt_label}`)
+        RETURN properties(r) as edge_properties
+        LIMIT 1
+        """
+
+        record = await self._query(
+            query,
+            src_label=AGEStorage._encode_graph_label(entity_name_label_source),
+            tgt_label=AGEStorage._encode_graph_label(entity_name_label_target),
+        )
+        if record and record[0] and record[0]["edge_properties"]:
+            result = record[0]["edge_properties"]
+            logger.debug(
+                "{%s}:query:{%s}:result:{%s}",
+                inspect.currentframe().f_code.co_name,
+                query,
+                result,
+            )
+            return result
+
+    async def get_node_edges(self, source_node_id: str) -> List[Tuple[str, str]]:
+        """
+        Retrieves all edges (relationships) for a particular node identified by its label.
+        :return: List of dictionaries containing edge information
+        """
+        node_label = source_node_id.strip('"')
+
+        query = """MATCH (n:`{label}`)
+                OPTIONAL MATCH (n)-[r]-(connected)
+                RETURN n, r, connected"""
+        results = await self._query(
+            query, label=AGEStorage._encode_graph_label(node_label)
+        )
+        edges = []
+        for record in results:
+            source_node = record["n"] if record["n"] else None
+            connected_node = record["connected"] if record["connected"] else None
+
+            source_label = (
+                source_node["label"] if source_node and source_node["label"] else None
+            )
+            target_label = (
+                connected_node["label"]
+                if connected_node and connected_node["label"]
+                else None
+            )
+
+            if source_label and target_label:
+                edges.append((source_label, target_label))
+
+        return edges
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        retry=retry_if_exception_type((AGEQueryException,)),
+    )
+    async def upsert_node(self, node_id: str, node_data: Dict[str, Any]):
+        """
+        Upsert a node in the AGE database.
+
+        Args:
+            node_id: The unique identifier for the node (used as label)
+            node_data: Dictionary of node properties
+        """
+        label = node_id.strip('"')
+        properties = node_data
+
+        query = """
+        MERGE (n:`{label}`)
+        SET n += {properties}
+        """
+        try:
+            await self._query(
+                query,
+                label=AGEStorage._encode_graph_label(label),
+                properties=AGEStorage._format_properties(properties),
+            )
+            logger.debug(
+                "Upserted node with label '{%s}' and properties: {%s}",
+                label,
+                properties,
+            )
+        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((AGEQueryException,)),
+    )
+    async def upsert_edge(
+        self, source_node_id: str, target_node_id: str, edge_data: Dict[str, Any]
+    ):
+        """
+        Upsert an edge and its properties between two nodes identified by their labels.
+
+        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
+        """
+        source_node_label = source_node_id.strip('"')
+        target_node_label = target_node_id.strip('"')
+        edge_properties = edge_data
+
+        query = """
+        MATCH (source:`{src_label}`)
+        WITH source
+        MATCH (target:`{tgt_label}`)
+        MERGE (source)-[r:DIRECTED]->(target)
+        SET r += {properties}
+        RETURN r
+        """
+        try:
+            await self._query(
+                query,
+                src_label=AGEStorage._encode_graph_label(source_node_label),
+                tgt_label=AGEStorage._encode_graph_label(target_node_label),
+                properties=AGEStorage._format_properties(edge_properties),
+            )
+            logger.debug(
+                "Upserted edge from '{%s}' to '{%s}' with properties: {%s}",
+                source_node_label,
+                target_node_label,
+                edge_properties,
+            )
+        except Exception as e:
+            logger.error("Error during edge upsert: {%s}", e)
+            raise
+
+    async def _node2vec_embed(self):
+        print("Implemented but never called.")
+
+    @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/lightrag.py

@@ -79,6 +79,7 @@ MongoKVStorage = lazy_external_import(".kg.mongo_impl", "MongoKVStorage")
 ChromaVectorDBStorage = lazy_external_import(".kg.chroma_impl", "ChromaVectorDBStorage")
 TiDBKVStorage = lazy_external_import(".kg.tidb_impl", "TiDBKVStorage")
 TiDBVectorDBStorage = lazy_external_import(".kg.tidb_impl", "TiDBVectorDBStorage")
+AGEStorage = lazy_external_import(".kg.age_impl", "AGEStorage")
 
 
 def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
@@ -273,6 +274,7 @@ class LightRAG:
             "NetworkXStorage": NetworkXStorage,
             "Neo4JStorage": Neo4JStorage,
             "OracleGraphStorage": OracleGraphStorage,
+            "AGEStorage": AGEStorage,
             # "ArangoDBStorage": ArangoDBStorage
         }
 

requirements.txt

@@ -11,6 +11,7 @@ networkx
 ollama
 openai
 oracledb
+psycopg[binary,pool]
 pymilvus
 pymongo
 pymysql