diff --git a/README-zh.md b/README-zh.md index 49c79f78..5f9614e3 100644 --- a/README-zh.md +++ b/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` | diff --git a/README.md b/README.md index 823ba7ba..0fa6c3d1 100644 --- a/README.md +++ b/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` | diff --git a/env.example b/env.example index d1b04b03..41b5d629 100644 --- a/env.example +++ b/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 diff --git a/lightrag/api/README.md b/lightrag/api/README.md index 4ba9b2cf..590c4ef5 100644 --- a/lightrag/api/README.md +++ b/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. diff --git a/lightrag/api/utils_api.py b/lightrag/api/utils_api.py index a724069d..b7099bb3 100644 --- a/lightrag/api/utils_api.py +++ b/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)}" diff --git a/lightrag/constants.py b/lightrag/constants.py index 6aa9673e..ea6cf0bb 100644 --- a/lightrag/constants.py +++ b/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 diff --git a/lightrag/kg/memgraph_impl.py b/lightrag/kg/memgraph_impl.py index 8c6d6574..3d2c131e 100644 --- a/lightrag/kg/memgraph_impl.py +++ b/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 diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index 81bac69c..df3929f6 100644 --- a/lightrag/lightrag.py +++ b/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 diff --git a/lightrag/operate.py b/lightrag/operate.py index 94a8a4fe..8de3a1da 100644 --- a/lightrag/operate.py +++ b/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"]