Merge branch 'main' into rerank

README-zh.md

@@ -242,7 +242,6 @@ if __name__ == "__main__":
 | **tokenizer** | `Tokenizer` | 用于将文本转换为 tokens（数字）以及使用遵循 TokenizerInterface 协议的 .encode() 和 .decode() 函数将 tokens 转换回文本的函数。 如果您不指定，它将使用默认的 Tiktoken tokenizer。 | `TiktokenTokenizer` |
 | **tiktoken_model_name** | `str` | 如果您使用的是默认的 Tiktoken tokenizer，那么这是要使用的特定 Tiktoken 模型的名称。如果您提供自己的 tokenizer，则忽略此设置。 | `gpt-4o-mini` |
 | **entity_extract_max_gleaning** | `int` | 实体提取过程中的循环次数，附加历史消息 | `1` |
-| **entity_summary_to_max_tokens** | `int` | 每个实体摘要的最大令牌大小 | `500` |
 | **node_embedding_algorithm** | `str` | 节点嵌入算法（当前未使用） | `node2vec` |
 | **node2vec_params** | `dict` | 节点嵌入的参数 | `{"dimensions": 1536,"num_walks": 10,"walk_length": 40,"window_size": 2,"iterations": 3,"random_seed": 3,}` |
 | **embedding_func** | `EmbeddingFunc` | 从文本生成嵌入向量的函数 | `openai_embed` |

README.md

@@ -249,7 +249,6 @@ A full list of LightRAG init parameters:
 | **tokenizer** | `Tokenizer` | The function used to convert text into tokens (numbers) and back using .encode() and .decode() functions following `TokenizerInterface` protocol. If you don't specify one, it will use the default Tiktoken tokenizer. | `TiktokenTokenizer` |
 | **tiktoken_model_name** | `str` | If you're using the default Tiktoken tokenizer, this is the name of the specific Tiktoken model to use. This setting is ignored if you provide your own tokenizer. | `gpt-4o-mini` |
 | **entity_extract_max_gleaning** | `int` | Number of loops in the entity extraction process, appending history messages | `1` |
-| **entity_summary_to_max_tokens** | `int` | Maximum token size for each entity summary | `500` |
 | **node_embedding_algorithm** | `str` | Algorithm for node embedding (currently not used) | `node2vec` |
 | **node2vec_params** | `dict` | Parameters for node embedding | `{"dimensions": 1536,"num_walks": 10,"walk_length": 40,"window_size": 2,"iterations": 3,"random_seed": 3,}` |
 | **embedding_func** | `EmbeddingFunc` | Function to generate embedding vectors from text | `openai_embed` |

env.example

@@ -75,8 +75,6 @@ OLLAMA_EMULATING_MODEL_TAG=latest
 SUMMARY_LANGUAGE=English
 ### Number of duplicated entities/edges to trigger LLM re-summary on merge ( at least 3 is recommented)
 # FORCE_LLM_SUMMARY_ON_MERGE=6
-### Max tokens for entity/relations description after merge
-# MAX_TOKEN_SUMMARY=500
 ### Maximum number of entity extraction attempts for ambiguous content
 # MAX_GLEANING=1
 

lightrag/api/README.md

@@ -181,9 +181,9 @@ The command-line `workspace` argument and the `WORKSPACE` environment variable i
 - **For local file-based databases, data isolation is achieved through workspace subdirectories:** `JsonKVStorage`, `JsonDocStatusStorage`, `NetworkXStorage`, `NanoVectorDBStorage`, `FaissVectorDBStorage`.
 - **For databases that store data in collections, it's done by adding a workspace prefix to the collection name:** `RedisKVStorage`, `RedisDocStatusStorage`, `MilvusVectorDBStorage`, `QdrantVectorDBStorage`, `MongoKVStorage`, `MongoDocStatusStorage`, `MongoVectorDBStorage`, `MongoGraphStorage`, `PGGraphStorage`.
 - **For relational databases, data isolation is achieved by adding a `workspace` field to the tables for logical data separation:** `PGKVStorage`, `PGVectorStorage`, `PGDocStatusStorage`.
-- **For the Neo4j graph database, logical data isolation is achieved through labels:** `Neo4JStorage`
+- **For graph databases, logical data isolation is achieved through labels:** `Neo4JStorage`, `MemgraphStorage`
 
-To maintain compatibility with legacy data, the default workspace for PostgreSQL is `default` and for Neo4j is `base` when no workspace is configured. For all external storages, the system provides dedicated workspace environment variables to override the common `WORKSPACE` environment variable configuration. These storage-specific workspace environment variables are: `REDIS_WORKSPACE`, `MILVUS_WORKSPACE`, `QDRANT_WORKSPACE`, `MONGODB_WORKSPACE`, `POSTGRES_WORKSPACE`, `NEO4J_WORKSPACE`.
+To maintain compatibility with legacy data, the default workspace for PostgreSQL is `default` and for Neo4j is `base` when no workspace is configured. For all external storages, the system provides dedicated workspace environment variables to override the common `WORKSPACE` environment variable configuration. These storage-specific workspace environment variables are: `REDIS_WORKSPACE`, `MILVUS_WORKSPACE`, `QDRANT_WORKSPACE`, `MONGODB_WORKSPACE`, `POSTGRES_WORKSPACE`, `NEO4J_WORKSPACE`, `MEMGRAPH_WORKSPACE`.
 
 ### Multiple workers for Gunicorn + Uvicorn
 
@@ -396,6 +396,7 @@ MongoKVStorage   MongoDB
 NetworkXStorage      NetworkX (default)
 Neo4JStorage         Neo4J
 PGGraphStorage       PostgreSQL with AGE plugin
+MemgraphStorage.     Memgraph
 ```
 
 > Testing has shown that Neo4J delivers superior performance in production environments compared to PostgreSQL with AGE plugin.

lightrag/api/utils_api.py

@@ -10,7 +10,6 @@ from ascii_colors import ASCIIColors
 from lightrag.api import __api_version__ as api_version
 from lightrag import __version__ as core_version
 from lightrag.constants import (
-    DEFAULT_MAX_TOKEN_SUMMARY,
     DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE,
 )
 from fastapi import HTTPException, Security, Request, status
@@ -280,9 +279,6 @@ def display_splash_screen(args: argparse.Namespace) -> None:
     ASCIIColors.white("    ├─ Top-K: ", end="")
     ASCIIColors.yellow(f"{args.top_k}")
     ASCIIColors.white("    ├─ Max Token Summary: ", end="")
-    ASCIIColors.yellow(
-        f"{get_env_value('MAX_TOKEN_SUMMARY', DEFAULT_MAX_TOKEN_SUMMARY, int)}"
-    )
     ASCIIColors.white("    └─ Force LLM Summary on Merge: ", end="")
     ASCIIColors.yellow(
         f"{get_env_value('FORCE_LLM_SUMMARY_ON_MERGE', DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE, int)}"

lightrag/constants.py

@@ -8,8 +8,7 @@ consistency and makes maintenance easier.
 
 # Default values for environment variables
 DEFAULT_MAX_GLEANING = 1
-DEFAULT_MAX_TOKEN_SUMMARY = 500
-DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE = 6
+DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE = 4
 DEFAULT_WOKERS = 2
 DEFAULT_TIMEOUT = 150
 

lightrag/kg/memgraph_impl.py

@@ -435,7 +435,7 @@ class MemgraphStorage(BaseGraphStorage):
 
     async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
         """
-        Upsert a node in the Neo4j database.
+        Upsert a node in the Memgraph database.
 
         Args:
             node_id: The unique identifier for the node (used as label)
@@ -448,7 +448,9 @@ class MemgraphStorage(BaseGraphStorage):
         properties = node_data
         entity_type = properties["entity_type"]
         if "entity_id" not in properties:
-            raise ValueError("Neo4j: node properties must contain an 'entity_id' field")
+            raise ValueError(
+                "Memgraph: node properties must contain an 'entity_id' field"
+            )
 
         try:
             async with self._driver.session(database=self._DATABASE) as session:
@@ -732,7 +734,7 @@ class MemgraphStorage(BaseGraphStorage):
         self,
         node_label: str,
         max_depth: int = 3,
-        max_nodes: int = MAX_GRAPH_NODES,
+        max_nodes: int = None,
     ) -> KnowledgeGraph:
         """
         Retrieve a connected subgraph of nodes where the label includes the specified `node_label`.
@@ -740,120 +742,118 @@ class MemgraphStorage(BaseGraphStorage):
         Args:
             node_label: Label of the starting node, * means all nodes
             max_depth: Maximum depth of the subgraph, Defaults to 3
-            max_nodes: Maxiumu nodes to return by BFS, Defaults to 1000
+            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
-
-        Raises:
-            Exception: If there is an error executing the query
         """
-        if self._driver is None:
-            raise RuntimeError(
-                "Memgraph driver is not initialized. Call 'await initialize()' first."
-            )
+        # Get max_nodes from global_config if not provided
+        if max_nodes is None:
+            max_nodes = self.global_config.get("max_graph_nodes", 1000)
+        else:
+            # Limit max_nodes to not exceed global_config max_graph_nodes
+            max_nodes = min(max_nodes, self.global_config.get("max_graph_nodes", 1000))
 
+        workspace_label = self._get_workspace_label()
         result = KnowledgeGraph()
         seen_nodes = set()
         seen_edges = set()
-        workspace_label = self._get_workspace_label()
+
         async with self._driver.session(
             database=self._DATABASE, default_access_mode="READ"
         ) as session:
             try:
                 if node_label == "*":
-                    # First check if database has any nodes
-                    count_query = "MATCH (n) RETURN count(n) as total"
+                    # First check total node count to determine if graph is truncated
+                    count_query = (
+                        f"MATCH (n:`{workspace_label}`) RETURN count(n) as total"
+                    )
                     count_result = None
-                    total_count = 0
                     try:
                         count_result = await session.run(count_query)
                         count_record = await count_result.single()
-                        if count_record:
-                            total_count = count_record["total"]
-                            if total_count == 0:
-                                logger.debug("No nodes found in database")
-                                return result
-                            if total_count > max_nodes:
-                                result.is_truncated = True
-                                logger.info(
-                                    f"Graph truncated: {total_count} nodes found, limited to {max_nodes}"
-                                )
+
+                        if count_record and count_record["total"] > max_nodes:
+                            result.is_truncated = True
+                            logger.info(
+                                f"Graph truncated: {count_record['total']} nodes found, limited to {max_nodes}"
+                            )
                     finally:
                         if count_result:
                             await count_result.consume()
 
-                    # Run the main query to get nodes with highest degree
+                    # Run main query to get nodes with highest degree
                     main_query = f"""
                     MATCH (n:`{workspace_label}`)
                     OPTIONAL MATCH (n)-[r]-()
                     WITH n, COALESCE(count(r), 0) AS degree
                     ORDER BY degree DESC
                     LIMIT $max_nodes
-                    WITH collect(n) AS kept_nodes
-                    MATCH (a)-[r]-(b)
+                    WITH collect({{node: n}}) AS filtered_nodes
+                    UNWIND filtered_nodes AS node_info
+                    WITH collect(node_info.node) AS kept_nodes, filtered_nodes
+                    OPTIONAL MATCH (a)-[r]-(b)
                     WHERE a IN kept_nodes AND b IN kept_nodes
-                    RETURN [node IN kept_nodes | {{node: node}}] AS node_info,
+                    RETURN filtered_nodes AS node_info,
                         collect(DISTINCT r) AS relationships
                     """
                     result_set = None
                     try:
                         result_set = await session.run(
-                            main_query, {"max_nodes": max_nodes}
+                            main_query,
+                            {"max_nodes": max_nodes},
                         )
                         record = await result_set.single()
-                        if not record:
-                            logger.debug("No record returned from main query")
-                            return result
                     finally:
                         if result_set:
                             await result_set.consume()
 
                 else:
-                    bfs_query = f"""
+                    # Run subgraph query for specific node_label
+                    subgraph_query = f"""
                     MATCH (start:`{workspace_label}`)
                     WHERE start.entity_id = $entity_id
-                    WITH start
-                    CALL {{
-                        WITH start
-                        MATCH path = (start)-[*0..{max_depth}]-(node)
-                        WITH nodes(path) AS path_nodes, relationships(path) AS path_rels
-                        UNWIND path_nodes AS n
-                        WITH collect(DISTINCT n) AS all_nodes, collect(DISTINCT path_rels) AS all_rel_lists
-                        WITH all_nodes, reduce(r = [], x IN all_rel_lists | r + x) AS all_rels
-                        RETURN all_nodes, all_rels
-                    }}
-                    WITH all_nodes AS nodes, all_rels AS relationships, size(all_nodes) AS total_nodes
+
+                    MATCH path = (start)-[*BFS 0..{max_depth}]-(end:`{workspace_label}`)
+                    WHERE ALL(n IN nodes(path) WHERE '{workspace_label}' IN labels(n))
+                    WITH collect(DISTINCT end) + start AS all_nodes_unlimited
                     WITH
                     CASE
-                        WHEN total_nodes <= {max_nodes} THEN nodes
-                        ELSE nodes[0..{max_nodes}]
+                        WHEN size(all_nodes_unlimited) <= $max_nodes THEN all_nodes_unlimited
+                        ELSE all_nodes_unlimited[0..$max_nodes]
                     END AS limited_nodes,
-                    relationships,
-                    total_nodes,
-                    total_nodes > {max_nodes} AS is_truncated
+                    size(all_nodes_unlimited) > $max_nodes AS is_truncated
+
+                    UNWIND limited_nodes AS n
+                    MATCH (n)-[r]-(m)
+                    WHERE m IN limited_nodes
+                    WITH collect(DISTINCT n) AS limited_nodes, collect(DISTINCT r) AS relationships, is_truncated
+
                     RETURN
                     [node IN limited_nodes | {{node: node}}] AS node_info,
                     relationships,
-                    total_nodes,
                     is_truncated
                     """
+
                     result_set = None
                     try:
                         result_set = await session.run(
-                            bfs_query,
+                            subgraph_query,
                             {
                                 "entity_id": node_label,
+                                "max_nodes": max_nodes,
                             },
                         )
                         record = await result_set.single()
+
+                        # If no record found, return empty KnowledgeGraph
                         if not record:
                             logger.debug(f"No nodes found for entity_id: {node_label}")
                             return result
 
-                        # Check if the query indicates truncation
-                        if "is_truncated" in record and record["is_truncated"]:
+                        # Check if the result was truncated
+                        if record.get("is_truncated"):
                             result.is_truncated = True
                             logger.info(
                                 f"Graph truncated: breadth-first search limited to {max_nodes} nodes"
@@ -863,13 +863,11 @@ class MemgraphStorage(BaseGraphStorage):
                         if result_set:
                             await result_set.consume()
 
-                # Process the record if it exists
-                if record and record["node_info"]:
+                if record:
                     for node_info in record["node_info"]:
                         node = node_info["node"]
                         node_id = node.id
                         if node_id not in seen_nodes:
-                            seen_nodes.add(node_id)
                             result.nodes.append(
                                 KnowledgeGraphNode(
                                     id=f"{node_id}",
@@ -877,11 +875,11 @@ class MemgraphStorage(BaseGraphStorage):
                                     properties=dict(node),
                                 )
                             )
+                            seen_nodes.add(node_id)
 
                     for rel in record["relationships"]:
                         edge_id = rel.id
                         if edge_id not in seen_edges:
-                            seen_edges.add(edge_id)
                             start = rel.start_node
                             end = rel.end_node
                             result.edges.append(
@@ -893,14 +891,13 @@ class MemgraphStorage(BaseGraphStorage):
                                     properties=dict(rel),
                                 )
                             )
+                            seen_edges.add(edge_id)
 
-                logger.info(
-                    f"Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}"
-                )
+                    logger.info(
+                        f"Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}"
+                    )
 
             except Exception as e:
-                logger.error(f"Error getting knowledge graph: {str(e)}")
-                # Return empty but properly initialized KnowledgeGraph on error
-                return KnowledgeGraph()
+                logger.warning(f"Memgraph error during subgraph query: {str(e)}")
 
         return result

lightrag/lightrag.py

@@ -23,7 +23,6 @@ from typing import (
 )
 from lightrag.constants import (
     DEFAULT_MAX_GLEANING,
-    DEFAULT_MAX_TOKEN_SUMMARY,
     DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE,
 )
 from lightrag.utils import get_env_value
@@ -134,10 +133,6 @@ class LightRAG:
     )
     """Maximum number of entity extraction attempts for ambiguous content."""
 
-    summary_to_max_tokens: int = field(
-        default=get_env_value("MAX_TOKEN_SUMMARY", DEFAULT_MAX_TOKEN_SUMMARY, int)
-    )
-
     force_llm_summary_on_merge: int = field(
         default=get_env_value(
             "FORCE_LLM_SUMMARY_ON_MERGE", DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE, int

lightrag/operate.py

@@ -123,7 +123,6 @@ async def _handle_entity_relation_summary(
 
     tokenizer: Tokenizer = global_config["tokenizer"]
     llm_max_tokens = global_config["llm_model_max_token_size"]
-    # summary_max_tokens = global_config["summary_to_max_tokens"]
 
     language = global_config["addon_params"].get(
         "language", PROMPTS["DEFAULT_LANGUAGE"]