From baa505eeb0bda03c5c13aefd44e81a91e80d3f04 Mon Sep 17 00:00:00 2001 From: Abyl Ikhsanov Date: Sat, 1 Mar 2025 13:26:02 +0100 Subject: [PATCH 01/16] adding full_doc_id to insert --- lightrag/lightrag.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index 15338874..b81ca08c 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -835,11 +835,11 @@ class LightRAG: await asyncio.gather(*tasks) logger.info("All Insert done") - def insert_custom_kg(self, custom_kg: dict[str, Any]) -> None: + def insert_custom_kg(self, custom_kg: dict[str, Any], full_doc_id: str = None) -> None: loop = always_get_an_event_loop() - loop.run_until_complete(self.ainsert_custom_kg(custom_kg)) + loop.run_until_complete(self.ainsert_custom_kg(custom_kg, full_doc_id)) - async def ainsert_custom_kg(self, custom_kg: dict[str, Any]) -> None: + async def ainsert_custom_kg(self, custom_kg: dict[str, Any], full_doc_id: str | None) -> None: update_storage = False try: # Insert chunks into vector storage @@ -865,7 +865,7 @@ class LightRAG: "source_id": source_id, "tokens": tokens, "chunk_order_index": chunk_order_index, - "full_doc_id": source_id, + "full_doc_id": full_doc_id if full_doc_id is not None else source_id, "status": DocStatus.PROCESSED, } all_chunks_data[chunk_id] = chunk_entry From bf0ddc645030b8aa22cd09828fc2d7341855f720 Mon Sep 17 00:00:00 2001 From: Abyl Ikhsanov Date: Sat, 1 Mar 2025 13:34:49 +0100 Subject: [PATCH 02/16] making optional for ainsert --- lightrag/lightrag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index b81ca08c..73747a91 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -839,7 +839,7 @@ class LightRAG: loop = always_get_an_event_loop() loop.run_until_complete(self.ainsert_custom_kg(custom_kg, full_doc_id)) - async def ainsert_custom_kg(self, custom_kg: dict[str, Any], full_doc_id: str | None) -> None: + async def ainsert_custom_kg(self, custom_kg: dict[str, Any], full_doc_id: str = None) -> None: update_storage = False try: # Insert chunks into vector storage From 1ca6837219ea38c512a3cf13504c930c3cddf162 Mon Sep 17 00:00:00 2001 From: yangdx Date: Sun, 2 Mar 2025 12:52:25 +0800 Subject: [PATCH 03/16] Add max nodes limit for graph retrieval of networkX MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit • Set MAX_GRAPH_NODES env var (default 1000) • Change edge type to "RELATED" --- .env.example | 1 + lightrag/api/routers/graph_routes.py | 19 +++++++++++++++++-- lightrag/kg/networkx_impl.py | 14 ++++++++------ 3 files changed, 26 insertions(+), 8 deletions(-) diff --git a/.env.example b/.env.example index de9b6452..70cb575c 100644 --- a/.env.example +++ b/.env.example @@ -3,6 +3,7 @@ # PORT=9621 # WORKERS=1 # NAMESPACE_PREFIX=lightrag # separating data from difference Lightrag instances +# MAX_GRAPH_NODES=1000 # Max nodes return from grap retrieval # CORS_ORIGINS=http://localhost:3000,http://localhost:8080 ### Optional SSL Configuration diff --git a/lightrag/api/routers/graph_routes.py b/lightrag/api/routers/graph_routes.py index 95a72758..aa1803c2 100644 --- a/lightrag/api/routers/graph_routes.py +++ b/lightrag/api/routers/graph_routes.py @@ -16,12 +16,27 @@ def create_graph_routes(rag, api_key: Optional[str] = None): @router.get("/graph/label/list", dependencies=[Depends(optional_api_key)]) async def get_graph_labels(): - """Get all graph labels""" + """ + Get all graph labels + + Returns: + List[str]: List of graph labels + """ return await rag.get_graph_labels() @router.get("/graphs", dependencies=[Depends(optional_api_key)]) async def get_knowledge_graph(label: str, max_depth: int = 3): - """Get knowledge graph for a specific label""" + """ + Get knowledge graph for a specific label. + Maximum number of nodes is limited to env MAX_GRAPH_NODES(default: 1000) + + Args: + label (str): Label to get knowledge graph for + max_depth (int, optional): Maximum depth of graph. Defaults to 3. + + Returns: + Dict[str, List[str]]: Knowledge graph for label + """ return await rag.get_knowledge_graph(node_label=label, max_depth=max_depth) return router diff --git a/lightrag/kg/networkx_impl.py b/lightrag/kg/networkx_impl.py index f11e9c0e..b1cc45fe 100644 --- a/lightrag/kg/networkx_impl.py +++ b/lightrag/kg/networkx_impl.py @@ -24,6 +24,8 @@ from .shared_storage import ( is_multiprocess, ) +MAX_GRAPH_NODES = int(os.getenv("MAX_GRAPH_NODES", 1000)) + @final @dataclass @@ -234,6 +236,7 @@ class NetworkXStorage(BaseGraphStorage): ) -> KnowledgeGraph: """ Get complete connected subgraph for specified node (including the starting node itself) + Maximum number of nodes is limited to env MAX_GRAPH_NODES(default: 1000) Args: node_label: Label of the starting node @@ -269,18 +272,17 @@ class NetworkXStorage(BaseGraphStorage): subgraph = nx.ego_graph(graph, nodes_to_explore[0], radius=max_depth) # Check if number of nodes exceeds max_graph_nodes - max_graph_nodes = 500 - if len(subgraph.nodes()) > max_graph_nodes: + if len(subgraph.nodes()) > MAX_GRAPH_NODES: origin_nodes = len(subgraph.nodes()) node_degrees = dict(subgraph.degree()) top_nodes = sorted(node_degrees.items(), key=lambda x: x[1], reverse=True)[ - :max_graph_nodes + :MAX_GRAPH_NODES ] top_node_ids = [node[0] for node in top_nodes] - # Create new subgraph with only top nodes + # Create new subgraph and keep nodes only with most degree subgraph = subgraph.subgraph(top_node_ids) logger.info( - f"Reduced graph from {origin_nodes} nodes to {max_graph_nodes} nodes (depth={max_depth})" + f"Reduced graph from {origin_nodes} nodes to {MAX_GRAPH_NODES} nodes (depth={max_depth})" ) # Add nodes to result @@ -320,7 +322,7 @@ class NetworkXStorage(BaseGraphStorage): result.edges.append( KnowledgeGraphEdge( id=edge_id, - type="DIRECTED", + type="RELATED", source=str(source), target=str(target), properties=edge_data, From 0f1eb42c8dd7e2440f6c4f1c18afbfc37ad2b9c0 Mon Sep 17 00:00:00 2001 From: yangdx Date: Sun, 2 Mar 2025 15:39:14 +0800 Subject: [PATCH 04/16] Add node limit and prioritization for knowledge graph retrieval MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit • Add MAX_GRAPH_NODES limit from env var • Prioritize nodes by label match & connection --- lightrag/kg/neo4j_impl.py | 56 ++++++++++++++++++++++++++++-------- lightrag/kg/networkx_impl.py | 47 +++++++++++++++++++++++++++--- 2 files changed, 87 insertions(+), 16 deletions(-) diff --git a/lightrag/kg/neo4j_impl.py b/lightrag/kg/neo4j_impl.py index f5c2237a..2fb2c494 100644 --- a/lightrag/kg/neo4j_impl.py +++ b/lightrag/kg/neo4j_impl.py @@ -23,7 +23,7 @@ import pipmaster as pm if not pm.is_installed("neo4j"): pm.install("neo4j") -from neo4j import ( +from neo4j import ( # type: ignore AsyncGraphDatabase, exceptions as neo4jExceptions, AsyncDriver, @@ -34,6 +34,9 @@ from neo4j import ( config = configparser.ConfigParser() config.read("config.ini", "utf-8") +# 从环境变量获取最大图节点数，默认为1000 +MAX_GRAPH_NODES = int(os.getenv("MAX_GRAPH_NODES", 1000)) + @final @dataclass @@ -471,12 +474,17 @@ class Neo4JStorage(BaseGraphStorage): ) -> KnowledgeGraph: """ Get complete connected subgraph for specified node (including the starting node itself) + Maximum number of nodes is constrained by the environment variable `MAX_GRAPH_NODES` (default: 1000). + When reducing the number of nodes, the prioritization criteria are as follows: + 1. Label matching nodes take precedence + 2. Followed by nodes directly connected to the matching nodes + 3. Finally, the degree of the nodes - Key fixes: - 1. Include the starting node itself - 2. Handle multi-label nodes - 3. Clarify relationship directions - 4. Add depth control + Args: + node_label (str): Label of the starting node + max_depth (int, optional): Maximum depth of the graph. Defaults to 5. + Returns: + KnowledgeGraph: Complete connected subgraph for specified node """ label = node_label.strip('"') result = KnowledgeGraph() @@ -485,14 +493,22 @@ class Neo4JStorage(BaseGraphStorage): async with self._driver.session(database=self._DATABASE) as session: try: - main_query = "" if label == "*": main_query = """ MATCH (n) - WITH collect(DISTINCT n) AS nodes - MATCH ()-[r]-() - RETURN nodes, collect(DISTINCT r) AS relationships; + OPTIONAL MATCH (n)-[r]-() + WITH n, count(r) AS degree + ORDER BY degree DESC + LIMIT $max_nodes + WITH collect(n) AS nodes + MATCH (a)-[r]->(b) + WHERE a IN nodes AND b IN nodes + RETURN nodes, collect(DISTINCT r) AS relationships """ + result_set = await session.run( + main_query, {"max_nodes": MAX_GRAPH_NODES} + ) + else: # Critical debug step: first verify if starting node exists validate_query = f"MATCH (n:`{label}`) RETURN n LIMIT 1" @@ -512,9 +528,25 @@ class Neo4JStorage(BaseGraphStorage): bfs: true }}) YIELD nodes, relationships - RETURN nodes, relationships + WITH start, nodes, relationships + UNWIND nodes AS node + OPTIONAL MATCH (node)-[r]-() + WITH node, count(r) AS degree, start, nodes, relationships, + CASE + WHEN id(node) = id(start) THEN 2 + WHEN EXISTS((start)-->(node)) OR EXISTS((node)-->(start)) THEN 1 + ELSE 0 + END AS priority + ORDER BY priority DESC, degree DESC + LIMIT $max_nodes + WITH collect(node) AS filtered_nodes, nodes, relationships + RETURN filtered_nodes AS nodes, + [rel IN relationships WHERE startNode(rel) IN filtered_nodes AND endNode(rel) IN filtered_nodes] AS relationships """ - result_set = await session.run(main_query) + result_set = await session.run( + main_query, {"max_nodes": MAX_GRAPH_NODES} + ) + record = await result_set.single() if record: diff --git a/lightrag/kg/networkx_impl.py b/lightrag/kg/networkx_impl.py index b1cc45fe..462fb832 100644 --- a/lightrag/kg/networkx_impl.py +++ b/lightrag/kg/networkx_impl.py @@ -236,7 +236,11 @@ class NetworkXStorage(BaseGraphStorage): ) -> KnowledgeGraph: """ Get complete connected subgraph for specified node (including the starting node itself) - Maximum number of nodes is limited to env MAX_GRAPH_NODES(default: 1000) + Maximum number of nodes is constrained by the environment variable `MAX_GRAPH_NODES` (default: 1000). + When reducing the number of nodes, the prioritization criteria are as follows: + 1. Label matching nodes take precedence + 2. Followed by nodes directly connected to the matching nodes + 3. Finally, the degree of the nodes Args: node_label: Label of the starting node @@ -268,14 +272,49 @@ class NetworkXStorage(BaseGraphStorage): logger.warning(f"No nodes found with label {node_label}") return result - # Get subgraph using ego_graph - subgraph = nx.ego_graph(graph, nodes_to_explore[0], radius=max_depth) + # Get subgraph using ego_graph from all matching nodes + combined_subgraph = nx.Graph() + for start_node in nodes_to_explore: + node_subgraph = nx.ego_graph(graph, start_node, radius=max_depth) + combined_subgraph = nx.compose(combined_subgraph, node_subgraph) + subgraph = combined_subgraph # Check if number of nodes exceeds max_graph_nodes if len(subgraph.nodes()) > MAX_GRAPH_NODES: origin_nodes = len(subgraph.nodes()) + + # 获取节点度数 node_degrees = dict(subgraph.degree()) - top_nodes = sorted(node_degrees.items(), key=lambda x: x[1], reverse=True)[ + + # 标记起点节点和直接连接的节点 + start_nodes = set() + direct_connected_nodes = set() + + if node_label != "*" and nodes_to_explore: + # 所有在 nodes_to_explore 中的节点都是起点节点 + start_nodes = set(nodes_to_explore) + + # 获取与所有起点直接连接的节点 + for start_node in start_nodes: + direct_connected_nodes.update(subgraph.neighbors(start_node)) + + # 从直接连接节点中移除起点节点（避免重复） + direct_connected_nodes -= start_nodes + + # 按优先级和度数排序 + def priority_key(node_item): + node, degree = node_item + # 优先级排序：起点(2) > 直接连接(1) > 其他节点(0) + if node in start_nodes: + priority = 2 + elif node in direct_connected_nodes: + priority = 1 + else: + priority = 0 + return (priority, degree) # 先按优先级，再按度数 + + # 排序并选择前MAX_GRAPH_NODES个节点 + top_nodes = sorted(node_degrees.items(), key=priority_key, reverse=True)[ :MAX_GRAPH_NODES ] top_node_ids = [node[0] for node in top_nodes] From 68bf02abb6224ba1212af0ada9da54c23b0b3185 Mon Sep 17 00:00:00 2001 From: yangdx Date: Sun, 2 Mar 2025 16:20:37 +0800 Subject: [PATCH 05/16] refactor: improve graph querying with label substring matching and security fixes --- lightrag/kg/neo4j_impl.py | 26 ++++++++++++++++---------- lightrag/kg/networkx_impl.py | 2 +- 2 files changed, 17 insertions(+), 11 deletions(-) diff --git a/lightrag/kg/neo4j_impl.py b/lightrag/kg/neo4j_impl.py index 2fb2c494..8052b1f7 100644 --- a/lightrag/kg/neo4j_impl.py +++ b/lightrag/kg/neo4j_impl.py @@ -34,7 +34,7 @@ from neo4j import ( # type: ignore config = configparser.ConfigParser() config.read("config.ini", "utf-8") -# 从环境变量获取最大图节点数，默认为1000 +# Get maximum number of graph nodes from environment variable, default is 1000 MAX_GRAPH_NODES = int(os.getenv("MAX_GRAPH_NODES", 1000)) @@ -473,20 +473,22 @@ class Neo4JStorage(BaseGraphStorage): self, node_label: str, max_depth: int = 5 ) -> KnowledgeGraph: """ - Get complete connected subgraph for specified node (including the starting node itself) + Retrieve a connected subgraph of nodes where the label includes the specified `node_label`. Maximum number of nodes is constrained by the environment variable `MAX_GRAPH_NODES` (default: 1000). When reducing the number of nodes, the prioritization criteria are as follows: - 1. Label matching nodes take precedence + 1. Label matching nodes take precedence (nodes containing the specified label string) 2. Followed by nodes directly connected to the matching nodes 3. Finally, the degree of the nodes Args: - node_label (str): Label of the starting node + node_label (str): String to match in node labels (will match any node containing this string in its label) max_depth (int, optional): Maximum depth of the graph. Defaults to 5. Returns: KnowledgeGraph: Complete connected subgraph for specified node """ label = node_label.strip('"') + # Escape single quotes to prevent injection attacks + escaped_label = label.replace("'", "\\'") result = KnowledgeGraph() seen_nodes = set() seen_edges = set() @@ -510,16 +512,20 @@ class Neo4JStorage(BaseGraphStorage): ) else: - # Critical debug step: first verify if starting node exists - validate_query = f"MATCH (n:`{label}`) RETURN n LIMIT 1" + validate_query = f""" + MATCH (n) + WHERE any(label IN labels(n) WHERE label CONTAINS '{escaped_label}') + RETURN n LIMIT 1 + """ validate_result = await session.run(validate_query) if not await validate_result.single(): - logger.warning(f"Starting node {label} does not exist!") + logger.warning(f"No nodes containing '{label}' in their labels found!") return result - # Optimized query (including direction handling and self-loops) + # Main query uses partial matching main_query = f""" - MATCH (start:`{label}`) + MATCH (start) + WHERE any(label IN labels(start) WHERE label CONTAINS '{escaped_label}') WITH start CALL apoc.path.subgraphAll(start, {{ relationshipFilter: '>', @@ -598,7 +604,7 @@ class Neo4JStorage(BaseGraphStorage): result = {"nodes": [], "edges": []} visited_nodes = set() visited_edges = set() - + async def traverse(current_label: str, current_depth: int): if current_depth > max_depth: return diff --git a/lightrag/kg/networkx_impl.py b/lightrag/kg/networkx_impl.py index 462fb832..92d36fa6 100644 --- a/lightrag/kg/networkx_impl.py +++ b/lightrag/kg/networkx_impl.py @@ -235,7 +235,7 @@ class NetworkXStorage(BaseGraphStorage): self, node_label: str, max_depth: int = 5 ) -> KnowledgeGraph: """ - Get complete connected subgraph for specified node (including the starting node itself) + Retrieve a connected subgraph of nodes where the label includes the specified `node_label`. Maximum number of nodes is constrained by the environment variable `MAX_GRAPH_NODES` (default: 1000). When reducing the number of nodes, the prioritization criteria are as follows: 1. Label matching nodes take precedence From 465737efed6e0d4b81854d87a142762a9d631b98 Mon Sep 17 00:00:00 2001 From: yangdx Date: Sun, 2 Mar 2025 17:32:25 +0800 Subject: [PATCH 06/16] Fix linting --- lightrag/api/routers/graph_routes.py | 7 ++++++- lightrag/kg/neo4j_impl.py | 8 +++++--- lightrag/kg/networkx_impl.py | 2 +- 3 files changed, 12 insertions(+), 5 deletions(-) diff --git a/lightrag/api/routers/graph_routes.py b/lightrag/api/routers/graph_routes.py index aa1803c2..e6f894a2 100644 --- a/lightrag/api/routers/graph_routes.py +++ b/lightrag/api/routers/graph_routes.py @@ -27,7 +27,12 @@ def create_graph_routes(rag, api_key: Optional[str] = None): @router.get("/graphs", dependencies=[Depends(optional_api_key)]) async def get_knowledge_graph(label: str, max_depth: int = 3): """ - Get knowledge graph for a specific label. + Retrieve a connected subgraph of nodes where the label includes the specified label. + Maximum number of nodes is constrained by the environment variable `MAX_GRAPH_NODES` (default: 1000). + When reducing the number of nodes, the prioritization criteria are as follows: + 1. Label matching nodes take precedence + 2. Followed by nodes directly connected to the matching nodes + 3. Finally, the degree of the nodes Maximum number of nodes is limited to env MAX_GRAPH_NODES(default: 1000) Args: diff --git a/lightrag/kg/neo4j_impl.py b/lightrag/kg/neo4j_impl.py index 8052b1f7..dccee330 100644 --- a/lightrag/kg/neo4j_impl.py +++ b/lightrag/kg/neo4j_impl.py @@ -475,7 +475,7 @@ class Neo4JStorage(BaseGraphStorage): """ Retrieve a connected subgraph of nodes where the label includes the specified `node_label`. Maximum number of nodes is constrained by the environment variable `MAX_GRAPH_NODES` (default: 1000). - When reducing the number of nodes, the prioritization criteria are as follows: + When reducing the number of nodes, the prioritization criteria are as follows: 1. Label matching nodes take precedence (nodes containing the specified label string) 2. Followed by nodes directly connected to the matching nodes 3. Finally, the degree of the nodes @@ -519,7 +519,9 @@ class Neo4JStorage(BaseGraphStorage): """ validate_result = await session.run(validate_query) if not await validate_result.single(): - logger.warning(f"No nodes containing '{label}' in their labels found!") + logger.warning( + f"No nodes containing '{label}' in their labels found!" + ) return result # Main query uses partial matching @@ -604,7 +606,7 @@ class Neo4JStorage(BaseGraphStorage): result = {"nodes": [], "edges": []} visited_nodes = set() visited_edges = set() - + async def traverse(current_label: str, current_depth: int): if current_depth > max_depth: return diff --git a/lightrag/kg/networkx_impl.py b/lightrag/kg/networkx_impl.py index 92d36fa6..9601a35e 100644 --- a/lightrag/kg/networkx_impl.py +++ b/lightrag/kg/networkx_impl.py @@ -237,7 +237,7 @@ class NetworkXStorage(BaseGraphStorage): """ Retrieve a connected subgraph of nodes where the label includes the specified `node_label`. Maximum number of nodes is constrained by the environment variable `MAX_GRAPH_NODES` (default: 1000). - When reducing the number of nodes, the prioritization criteria are as follows: + When reducing the number of nodes, the prioritization criteria are as follows: 1. Label matching nodes take precedence 2. Followed by nodes directly connected to the matching nodes 3. Finally, the degree of the nodes From 11fdb60fe5ef30ec6cb447f6762f8cad1ff67b0b Mon Sep 17 00:00:00 2001 From: yangdx Date: Mon, 3 Mar 2025 01:30:41 +0800 Subject: [PATCH 07/16] Remove Chinese comments --- lightrag/kg/networkx_impl.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/lightrag/kg/networkx_impl.py b/lightrag/kg/networkx_impl.py index 9601a35e..563fc554 100644 --- a/lightrag/kg/networkx_impl.py +++ b/lightrag/kg/networkx_impl.py @@ -283,37 +283,32 @@ class NetworkXStorage(BaseGraphStorage): if len(subgraph.nodes()) > MAX_GRAPH_NODES: origin_nodes = len(subgraph.nodes()) - # 获取节点度数 node_degrees = dict(subgraph.degree()) - # 标记起点节点和直接连接的节点 start_nodes = set() direct_connected_nodes = set() if node_label != "*" and nodes_to_explore: - # 所有在 nodes_to_explore 中的节点都是起点节点 start_nodes = set(nodes_to_explore) - - # 获取与所有起点直接连接的节点 + # Get nodes directly connected to all start nodes for start_node in start_nodes: direct_connected_nodes.update(subgraph.neighbors(start_node)) - # 从直接连接节点中移除起点节点（避免重复） + # Remove start nodes from directly connected nodes (avoid duplicates) direct_connected_nodes -= start_nodes - # 按优先级和度数排序 def priority_key(node_item): node, degree = node_item - # 优先级排序：起点(2) > 直接连接(1) > 其他节点(0) + # Priority order: start(2) > directly connected(1) > other nodes(0) if node in start_nodes: priority = 2 elif node in direct_connected_nodes: priority = 1 else: priority = 0 - return (priority, degree) # 先按优先级，再按度数 + return (priority, degree) - # 排序并选择前MAX_GRAPH_NODES个节点 + # Sort by priority and degree and select top MAX_GRAPH_NODES nodes top_nodes = sorted(node_degrees.items(), key=priority_key, reverse=True)[ :MAX_GRAPH_NODES ] From c21d5744f9f7abb5b2058f8ff4007ef54d7c58e7 Mon Sep 17 00:00:00 2001 From: yangdx Date: Mon, 3 Mar 2025 02:05:54 +0800 Subject: [PATCH 08/16] Remove duplicated run_with_gunicorn.py --- run_with_gunicorn.py | 203 ------------------------------------------- 1 file changed, 203 deletions(-) delete mode 100755 run_with_gunicorn.py diff --git a/run_with_gunicorn.py b/run_with_gunicorn.py deleted file mode 100755 index 2e4e3cf7..00000000 --- a/run_with_gunicorn.py +++ /dev/null @@ -1,203 +0,0 @@ -#!/usr/bin/env python -""" -Start LightRAG server with Gunicorn -""" - -import os -import sys -import signal -import pipmaster as pm -from lightrag.api.utils_api import parse_args, display_splash_screen -from lightrag.kg.shared_storage import initialize_share_data, finalize_share_data - - -def check_and_install_dependencies(): - """Check and install required dependencies""" - required_packages = [ - "gunicorn", - "tiktoken", - "psutil", - # Add other required packages here - ] - - for package in required_packages: - if not pm.is_installed(package): - print(f"Installing {package}...") - pm.install(package) - print(f"{package} installed successfully") - - -# Signal handler for graceful shutdown -def signal_handler(sig, frame): - print("\n\n" + "=" * 80) - print("RECEIVED TERMINATION SIGNAL") - print(f"Process ID: {os.getpid()}") - print("=" * 80 + "\n") - - # Release shared resources - finalize_share_data() - - # Exit with success status - sys.exit(0) - - -def main(): - # Check and install dependencies - check_and_install_dependencies() - - # Register signal handlers for graceful shutdown - signal.signal(signal.SIGINT, signal_handler) # Ctrl+C - signal.signal(signal.SIGTERM, signal_handler) # kill command - - # Parse all arguments using parse_args - args = parse_args(is_uvicorn_mode=False) - - # Display startup information - display_splash_screen(args) - - print("🚀 Starting LightRAG with Gunicorn") - print(f"🔄 Worker management: Gunicorn (workers={args.workers})") - print("🔍 Preloading app: Enabled") - print("📝 Note: Using Gunicorn's preload feature for shared data initialization") - print("\n\n" + "=" * 80) - print("MAIN PROCESS INITIALIZATION") - print(f"Process ID: {os.getpid()}") - print(f"Workers setting: {args.workers}") - print("=" * 80 + "\n") - - # Import Gunicorn's StandaloneApplication - from gunicorn.app.base import BaseApplication - - # Define a custom application class that loads our config - class GunicornApp(BaseApplication): - def __init__(self, app, options=None): - self.options = options or {} - self.application = app - super().__init__() - - def load_config(self): - # Define valid Gunicorn configuration options - valid_options = { - "bind", - "workers", - "worker_class", - "timeout", - "keepalive", - "preload_app", - "errorlog", - "accesslog", - "loglevel", - "certfile", - "keyfile", - "limit_request_line", - "limit_request_fields", - "limit_request_field_size", - "graceful_timeout", - "max_requests", - "max_requests_jitter", - } - - # Special hooks that need to be set separately - special_hooks = { - "on_starting", - "on_reload", - "on_exit", - "pre_fork", - "post_fork", - "pre_exec", - "pre_request", - "post_request", - "worker_init", - "worker_exit", - "nworkers_changed", - "child_exit", - } - - # Import and configure the gunicorn_config module - import gunicorn_config - - # Set configuration variables in gunicorn_config, prioritizing command line arguments - gunicorn_config.workers = ( - args.workers if args.workers else int(os.getenv("WORKERS", 1)) - ) - - # Bind configuration prioritizes command line arguments - host = args.host if args.host != "0.0.0.0" else os.getenv("HOST", "0.0.0.0") - port = args.port if args.port != 9621 else int(os.getenv("PORT", 9621)) - gunicorn_config.bind = f"{host}:{port}" - - # Log level configuration prioritizes command line arguments - gunicorn_config.loglevel = ( - args.log_level.lower() - if args.log_level - else os.getenv("LOG_LEVEL", "info") - ) - - # Timeout configuration prioritizes command line arguments - gunicorn_config.timeout = ( - args.timeout if args.timeout else int(os.getenv("TIMEOUT", 150)) - ) - - # Keepalive configuration - gunicorn_config.keepalive = int(os.getenv("KEEPALIVE", 5)) - - # SSL configuration prioritizes command line arguments - if args.ssl or os.getenv("SSL", "").lower() in ( - "true", - "1", - "yes", - "t", - "on", - ): - gunicorn_config.certfile = ( - args.ssl_certfile - if args.ssl_certfile - else os.getenv("SSL_CERTFILE") - ) - gunicorn_config.keyfile = ( - args.ssl_keyfile if args.ssl_keyfile else os.getenv("SSL_KEYFILE") - ) - - # Set configuration options from the module - for key in dir(gunicorn_config): - if key in valid_options: - value = getattr(gunicorn_config, key) - # Skip functions like on_starting and None values - if not callable(value) and value is not None: - self.cfg.set(key, value) - # Set special hooks - elif key in special_hooks: - value = getattr(gunicorn_config, key) - if callable(value): - self.cfg.set(key, value) - - if hasattr(gunicorn_config, "logconfig_dict"): - self.cfg.set( - "logconfig_dict", getattr(gunicorn_config, "logconfig_dict") - ) - - def load(self): - # Import the application - from lightrag.api.lightrag_server import get_application - - return get_application(args) - - # Create the application - app = GunicornApp("") - - # Force workers to be an integer and greater than 1 for multi-process mode - workers_count = int(args.workers) - if workers_count > 1: - # Set a flag to indicate we're in the main process - os.environ["LIGHTRAG_MAIN_PROCESS"] = "1" - initialize_share_data(workers_count) - else: - initialize_share_data(1) - - # Run the application - print("\nStarting Gunicorn with direct Python API...") - app.run() - - -if __name__ == "__main__": - main() From 887388c317ae4bab933054241f808e2d657461ab Mon Sep 17 00:00:00 2001 From: zrguo Date: Mon, 3 Mar 2025 14:54:28 +0800 Subject: [PATCH 09/16] fix linting --- lightrag/lightrag.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index a76b7199..4f1ad7dc 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -952,11 +952,15 @@ class LightRAG: pipeline_status["latest_message"] = log_message pipeline_status["history_messages"].append(log_message) - def insert_custom_kg(self, custom_kg: dict[str, Any], full_doc_id: str = None) -> None: + def insert_custom_kg( + self, custom_kg: dict[str, Any], full_doc_id: str = None + ) -> None: loop = always_get_an_event_loop() loop.run_until_complete(self.ainsert_custom_kg(custom_kg, full_doc_id)) - async def ainsert_custom_kg(self, custom_kg: dict[str, Any], full_doc_id: str = None) -> None: + async def ainsert_custom_kg( + self, custom_kg: dict[str, Any], full_doc_id: str = None + ) -> None: update_storage = False try: # Insert chunks into vector storage @@ -982,7 +986,9 @@ class LightRAG: "source_id": source_id, "tokens": tokens, "chunk_order_index": chunk_order_index, - "full_doc_id": full_doc_id if full_doc_id is not None else source_id, + "full_doc_id": full_doc_id + if full_doc_id is not None + else source_id, "status": DocStatus.PROCESSED, } all_chunks_data[chunk_id] = chunk_entry From 0ea274a30dce6627986a4f40d51912f7c462b0bf Mon Sep 17 00:00:00 2001 From: MdNazishArmanShorthillsAI Date: Mon, 3 Mar 2025 13:53:45 +0530 Subject: [PATCH 10/16] Improved cashing check --- lightrag/lightrag.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index 4f1ad7dc..04f66adc 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -1150,7 +1150,7 @@ class LightRAG: """ if param.mode in ["local", "global", "hybrid"]: response = await kg_query( - query, + query.strip(), self.chunk_entity_relation_graph, self.entities_vdb, self.relationships_vdb, @@ -1171,7 +1171,7 @@ class LightRAG: ) elif param.mode == "naive": response = await naive_query( - query, + query.strip(), self.chunks_vdb, self.text_chunks, param, @@ -1190,7 +1190,7 @@ class LightRAG: ) elif param.mode == "mix": response = await mix_kg_vector_query( - query, + query.strip(), self.chunk_entity_relation_graph, self.entities_vdb, self.relationships_vdb, From 1611400854f6172d3f56be0f69eddfa7ec07d6c5 Mon Sep 17 00:00:00 2001 From: zrguo Date: Mon, 3 Mar 2025 18:33:42 +0800 Subject: [PATCH 11/16] fix demo --- .gitignore | 2 +- README.md | 225 +++++++++--------- examples/copy_llm_cache_to_another_storage.py | 2 +- examples/lightrag_api_ollama_demo.py | 51 ++-- .../lightrag_api_openai_compatible_demo.py | 39 ++- ...g_api_openai_compatible_demo_simplified.py | 101 -------- examples/lightrag_api_oracle_demo.py | 4 + examples/lightrag_azure_openai_demo.py | 4 + examples/lightrag_bedrock_demo.py | 51 ++-- examples/lightrag_gemini_demo.py | 54 +++-- examples/lightrag_hf_demo.py | 85 ++++--- examples/lightrag_jinaai_demo.py | 115 --------- examples/lightrag_llamaindex_direct_demo.py | 80 ++++--- examples/lightrag_llamaindex_litellm_demo.py | 80 ++++--- examples/lightrag_lmdeploy_demo.py | 89 ++++--- examples/lightrag_nvidia_demo.py | 57 ++--- examples/lightrag_ollama_age_demo.py | 114 +++++---- examples/lightrag_ollama_demo.py | 111 +++++---- examples/lightrag_ollama_gremlin_demo.py | 110 +++++---- ...lightrag_ollama_neo4j_milvus_mongo_demo.py | 82 +++++-- examples/lightrag_openai_compatible_demo.py | 32 ++- ..._openai_compatible_demo_embedding_cache.py | 41 ++-- .../lightrag_openai_compatible_stream_demo.py | 72 +++--- examples/lightrag_openai_demo.py | 64 +++-- .../lightrag_openai_mongodb_graph_demo.py | 43 +++- ...lightrag_openai_neo4j_milvus_redis_demo.py | 63 +++-- examples/lightrag_oracle_demo.py | 72 +++--- examples/lightrag_siliconcloud_demo.py | 64 +++-- examples/lightrag_tidb_demo.py | 54 +++-- examples/lightrag_zhipu_demo.py | 74 +++--- examples/lightrag_zhipu_postgres_demo.py | 14 +- examples/query_keyword_separation_example.py | 34 ++- examples/test.py | 61 +++-- examples/test_chromadb.py | 53 +++-- examples/test_faiss.py | 18 +- examples/test_neo4j.py | 66 +++-- examples/test_split_by_character.ipynb | 37 ++- examples/vram_management_demo.py | 120 +++++----- reproduce/Step_1.py | 19 +- reproduce/Step_1_openai_compatible.py | 31 ++- run_with_gunicorn.py | 203 ---------------- 41 files changed, 1390 insertions(+), 1301 deletions(-) delete mode 100644 examples/lightrag_api_openai_compatible_demo_simplified.py delete mode 100644 examples/lightrag_jinaai_demo.py delete mode 100755 run_with_gunicorn.py diff --git a/.gitignore b/.gitignore index 6deb14d5..a4afe4ea 100644 --- a/.gitignore +++ b/.gitignore @@ -57,7 +57,7 @@ ignore_this.txt *.ignore.* # Project-specific files -dickens/ +dickens*/ book.txt lightrag-dev/ gui/ diff --git a/README.md b/README.md index 86211b35..2b0d9f4b 100644 --- a/README.md +++ b/README.md @@ -102,33 +102,47 @@ Use the below Python snippet (in a script) to initialize LightRAG and perform qu ```python import os +import asyncio from lightrag import LightRAG, QueryParam from lightrag.llm.openai import gpt_4o_mini_complete, gpt_4o_complete, openai_embed +from lightrag.kg.shared_storage import initialize_pipeline_status -rag = LightRAG( - working_dir="your/path", - embedding_func=openai_embed, - llm_model_func=gpt_4o_mini_complete -) +async def initialize_rag(): + rag = LightRAG( + working_dir="your/path", + embedding_func=openai_embed, + llm_model_func=gpt_4o_mini_complete + ) -# Insert text -rag.insert("Your text") + await rag.initialize_storages() + await initialize_pipeline_status() -# Perform naive search -mode="naive" -# Perform local search -mode="local" -# Perform global search -mode="global" -# Perform hybrid search -mode="hybrid" -# Mix mode Integrates knowledge graph and vector retrieval. -mode="mix" + return rag -rag.query( - "What are the top themes in this story?", - param=QueryParam(mode=mode) -) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) + # Insert text + rag.insert("Your text") + + # Perform naive search + mode="naive" + # Perform local search + mode="local" + # Perform global search + mode="global" + # Perform hybrid search + mode="hybrid" + # Mix mode Integrates knowledge graph and vector retrieval. + mode="mix" + + rag.query( + "What are the top themes in this story?", + param=QueryParam(mode=mode) + ) + +if __name__ == "__main__": + main() ``` ### Query Param @@ -190,15 +204,21 @@ async def embedding_func(texts: list[str]) -> np.ndarray: base_url="https://api.upstage.ai/v1/solar" ) -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=4096, - max_token_size=8192, - func=embedding_func +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=llm_model_func, + embedding_func=EmbeddingFunc( + embedding_dim=4096, + max_token_size=8192, + func=embedding_func + ) ) -) + + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag ``` @@ -210,10 +230,6 @@ rag = LightRAG( See `lightrag_hf_demo.py` ```python -from lightrag.llm import hf_model_complete, hf_embed -from transformers import AutoModel, AutoTokenizer -from lightrag.utils import EmbeddingFunc - # Initialize LightRAG with Hugging Face model rag = LightRAG( working_dir=WORKING_DIR, @@ -242,9 +258,6 @@ If you want to use Ollama models, you need to pull model you plan to use and emb Then you only need to set LightRAG as follows: ```python -from lightrag.llm.ollama import ollama_model_complete, ollama_embed -from lightrag.utils import EmbeddingFunc - # Initialize LightRAG with Ollama model rag = LightRAG( working_dir=WORKING_DIR, @@ -325,20 +338,58 @@ LightRAG supports integration with LlamaIndex. ```python # Using LlamaIndex with direct OpenAI access +import asyncio from lightrag import LightRAG from lightrag.llm.llama_index_impl import llama_index_complete_if_cache, llama_index_embed from llama_index.embeddings.openai import OpenAIEmbedding from llama_index.llms.openai import OpenAI +from lightrag.kg.shared_storage import initialize_pipeline_status -rag = LightRAG( - working_dir="your/path", - llm_model_func=llama_index_complete_if_cache, # LlamaIndex-compatible completion function - embedding_func=EmbeddingFunc( # LlamaIndex-compatible embedding function - embedding_dim=1536, - max_token_size=8192, - func=lambda texts: llama_index_embed(texts, embed_model=embed_model) - ), -) +async def initialize_rag(): + rag = LightRAG( + working_dir="your/path", + llm_model_func=llama_index_complete_if_cache, # LlamaIndex-compatible completion function + embedding_func=EmbeddingFunc( # LlamaIndex-compatible embedding function + embedding_dim=1536, + max_token_size=8192, + func=lambda texts: llama_index_embed(texts, embed_model=embed_model) + ), + ) + + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag + +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) + + 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")) + ) + +if __name__ == "__main__": + main() ``` #### For detailed documentation and examples, see: @@ -353,11 +404,6 @@ rag = LightRAG( LightRAG now supports multi-turn dialogue through the conversation history feature. Here's how to use it: ```python -from lightrag import LightRAG, QueryParam - -# Initialize LightRAG -rag = LightRAG(working_dir=WORKING_DIR) - # Create conversation history conversation_history = [ {"role": "user", "content": "What is the main character's attitude towards Christmas?"}, @@ -387,11 +433,6 @@ response = rag.query( LightRAG now supports custom prompts for fine-tuned control over the system's behavior. Here's how to use it: ```python -from lightrag import LightRAG, QueryParam - -# Initialize LightRAG -rag = LightRAG(working_dir=WORKING_DIR) - # Create query parameters query_param = QueryParam( mode="hybrid", # or other mode: "local", "global", "hybrid", "mix" and "naive" @@ -456,16 +497,6 @@ rag.query_with_separate_keyword_extraction( Insert Custom KG ```python -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=embedding_dimension, - max_token_size=8192, - func=embedding_func, - ), -) - custom_kg = { "entities": [ { @@ -534,6 +565,7 @@ rag = LightRAG( "insert_batch_size": 20 # Process 20 documents per batch } ) + rag.insert(["TEXT1", "TEXT2", "TEXT3", ...]) # Documents will be processed in batches of 20 ``` @@ -560,27 +592,6 @@ rag.insert(["TEXT1", "TEXT2",...], ids=["ID_FOR_TEXT1", "ID_FOR_TEXT2"]) -
- Incremental Insert - -```python -# Incremental Insert: Insert new documents into an existing LightRAG instance -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=embedding_dimension, - max_token_size=8192, - func=embedding_func, - ), -) - -with open("./newText.txt") as f: - rag.insert(f.read()) -``` - -
-
Insert using Pipeline @@ -592,6 +603,7 @@ And using a routine to process news documents. ```python rag = LightRAG(..) + await rag.apipeline_enqueue_documents(input) # Your routine in loop await rag.apipeline_process_enqueue_documents(input) @@ -633,8 +645,6 @@ export NEO4J_PASSWORD="password" # Note: Default settings use NetworkX # Initialize LightRAG with Neo4J implementation. -WORKING_DIR = "./local_neo4jWorkDir" - rag = LightRAG( working_dir=WORKING_DIR, llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model @@ -706,26 +716,26 @@ You can also install `faiss-gpu` if you have GPU support. - Here we are using `sentence-transformers` but you can also use `OpenAIEmbedding` model with `3072` dimensions. -``` +```python async def embedding_func(texts: list[str]) -> np.ndarray: model = SentenceTransformer('all-MiniLM-L6-v2') embeddings = model.encode(texts, convert_to_numpy=True) return embeddings # Initialize LightRAG with the LLM model function and embedding function - rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=384, - max_token_size=8192, - func=embedding_func, - ), - vector_storage="FaissVectorDBStorage", - vector_db_storage_cls_kwargs={ - "cosine_better_than_threshold": 0.3 # Your desired threshold - } - ) +rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=llm_model_func, + embedding_func=EmbeddingFunc( + embedding_dim=384, + max_token_size=8192, + func=embedding_func, + ), + vector_storage="FaissVectorDBStorage", + vector_db_storage_cls_kwargs={ + "cosine_better_than_threshold": 0.3 # Your desired threshold + } +) ```
@@ -733,17 +743,6 @@ async def embedding_func(texts: list[str]) -> np.ndarray: ## Delete ```python - -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=embedding_dimension, - max_token_size=8192, - func=embedding_func, - ), -) - # Delete Entity: Deleting entities by their names rag.delete_by_entity("Project Gutenberg") diff --git a/examples/copy_llm_cache_to_another_storage.py b/examples/copy_llm_cache_to_another_storage.py index 5d07ad13..60fa6192 100644 --- a/examples/copy_llm_cache_to_another_storage.py +++ b/examples/copy_llm_cache_to_another_storage.py @@ -10,7 +10,7 @@ import os from dotenv import load_dotenv from lightrag.kg.postgres_impl import PostgreSQLDB, PGKVStorage -from lightrag.storage import JsonKVStorage +from lightrag.kg.json_kv_impl import JsonKVStorage from lightrag.namespace import NameSpace load_dotenv() diff --git a/examples/lightrag_api_ollama_demo.py b/examples/lightrag_api_ollama_demo.py index 079e9935..f1b68795 100644 --- a/examples/lightrag_api_ollama_demo.py +++ b/examples/lightrag_api_ollama_demo.py @@ -1,4 +1,5 @@ from fastapi import FastAPI, HTTPException, File, UploadFile +from contextlib import asynccontextmanager from pydantic import BaseModel import os from lightrag import LightRAG, QueryParam @@ -8,12 +9,12 @@ from typing import Optional import asyncio import nest_asyncio import aiofiles +from lightrag.kg.shared_storage import initialize_pipeline_status # Apply nest_asyncio to solve event loop issues nest_asyncio.apply() DEFAULT_RAG_DIR = "index_default" -app = FastAPI(title="LightRAG API", description="API for RAG operations") DEFAULT_INPUT_FILE = "book.txt" INPUT_FILE = os.environ.get("INPUT_FILE", f"{DEFAULT_INPUT_FILE}") @@ -28,23 +29,41 @@ if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=ollama_model_complete, - llm_model_name="gemma2:9b", - llm_model_max_async=4, - llm_model_max_token_size=8192, - llm_model_kwargs={"host": "http://localhost:11434", "options": {"num_ctx": 8192}}, - embedding_func=EmbeddingFunc( - embedding_dim=768, - max_token_size=8192, - func=lambda texts: ollama_embed( - texts, embed_model="nomic-embed-text", host="http://localhost:11434" +async def init(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=ollama_model_complete, + llm_model_name="gemma2:9b", + llm_model_max_async=4, + llm_model_max_token_size=8192, + llm_model_kwargs={"host": "http://localhost:11434", "options": {"num_ctx": 8192}}, + embedding_func=EmbeddingFunc( + embedding_dim=768, + max_token_size=8192, + func=lambda texts: ollama_embed( + texts, embed_model="nomic-embed-text", host="http://localhost:11434" + ), ), - ), + ) + + # Add initialization code + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag + + +@asynccontextmanager +async def lifespan(app: FastAPI): + global rag + rag = await init() + print("done!") + yield + + +app = FastAPI( + title="LightRAG API", description="API for RAG operations", lifespan=lifespan ) - - # Data models class QueryRequest(BaseModel): query: str diff --git a/examples/lightrag_api_openai_compatible_demo.py b/examples/lightrag_api_openai_compatible_demo.py index 68ccfe95..2206f40d 100644 --- a/examples/lightrag_api_openai_compatible_demo.py +++ b/examples/lightrag_api_openai_compatible_demo.py @@ -1,4 +1,5 @@ from fastapi import FastAPI, HTTPException, File, UploadFile +from contextlib import asynccontextmanager from pydantic import BaseModel import os from lightrag import LightRAG, QueryParam @@ -8,6 +9,7 @@ import numpy as np from typing import Optional import asyncio import nest_asyncio +from lightrag.kg.shared_storage import initialize_pipeline_status # Apply nest_asyncio to solve event loop issues nest_asyncio.apply() @@ -71,16 +73,35 @@ async def get_embedding_dim(): # Initialize RAG instance -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=asyncio.run(get_embedding_dim()), - max_token_size=EMBEDDING_MAX_TOKEN_SIZE, - func=embedding_func, - ), -) +async def init(): + embedding_dimension = await get_embedding_dim() + + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=llm_model_func, + embedding_func=EmbeddingFunc( + embedding_dim=embedding_dimension, + max_token_size=EMBEDDING_MAX_TOKEN_SIZE, + func=embedding_func, + ), + ) + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag + +@asynccontextmanager +async def lifespan(app: FastAPI): + global rag + rag = await init() + print("done!") + yield + + +app = FastAPI( + title="LightRAG API", description="API for RAG operations", lifespan=lifespan +) # Data models diff --git a/examples/lightrag_api_openai_compatible_demo_simplified.py b/examples/lightrag_api_openai_compatible_demo_simplified.py deleted file mode 100644 index fabbb3e2..00000000 --- a/examples/lightrag_api_openai_compatible_demo_simplified.py +++ /dev/null @@ -1,101 +0,0 @@ -import os -from lightrag import LightRAG, QueryParam -from lightrag.llm.openai import openai_complete_if_cache, openai_embed -from lightrag.utils import EmbeddingFunc -import numpy as np -import asyncio -import nest_asyncio - -# Apply nest_asyncio to solve event loop issues -nest_asyncio.apply() - -DEFAULT_RAG_DIR = "index_default" - -# Configure working directory -WORKING_DIR = os.environ.get("RAG_DIR", f"{DEFAULT_RAG_DIR}") -print(f"WORKING_DIR: {WORKING_DIR}") -LLM_MODEL = os.environ.get("LLM_MODEL", "gpt-4o-mini") -print(f"LLM_MODEL: {LLM_MODEL}") -EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "text-embedding-3-small") -print(f"EMBEDDING_MODEL: {EMBEDDING_MODEL}") -EMBEDDING_MAX_TOKEN_SIZE = int(os.environ.get("EMBEDDING_MAX_TOKEN_SIZE", 8192)) -print(f"EMBEDDING_MAX_TOKEN_SIZE: {EMBEDDING_MAX_TOKEN_SIZE}") -BASE_URL = os.environ.get("BASE_URL", "https://api.openai.com/v1") -print(f"BASE_URL: {BASE_URL}") -API_KEY = os.environ.get("API_KEY", "xxxxxxxx") -print(f"API_KEY: {API_KEY}") - -if not os.path.exists(WORKING_DIR): - os.mkdir(WORKING_DIR) - - -# LLM model function - - -async def llm_model_func( - prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs -) -> str: - return await openai_complete_if_cache( - model=LLM_MODEL, - prompt=prompt, - system_prompt=system_prompt, - history_messages=history_messages, - base_url=BASE_URL, - api_key=API_KEY, - **kwargs, - ) - - -# Embedding function - - -async def embedding_func(texts: list[str]) -> np.ndarray: - return await openai_embed( - texts=texts, - model=EMBEDDING_MODEL, - base_url=BASE_URL, - api_key=API_KEY, - ) - - -async def get_embedding_dim(): - test_text = ["This is a test sentence."] - embedding = await embedding_func(test_text) - embedding_dim = embedding.shape[1] - print(f"{embedding_dim=}") - return embedding_dim - - -# Initialize RAG instance -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=asyncio.run(get_embedding_dim()), - max_token_size=EMBEDDING_MAX_TOKEN_SIZE, - func=embedding_func, - ), -) - -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")) -) diff --git a/examples/lightrag_api_oracle_demo.py b/examples/lightrag_api_oracle_demo.py index 3675795e..3a82f479 100644 --- a/examples/lightrag_api_oracle_demo.py +++ b/examples/lightrag_api_oracle_demo.py @@ -16,6 +16,7 @@ from lightrag import LightRAG, QueryParam from lightrag.llm.openai import openai_complete_if_cache, openai_embed from lightrag.utils import EmbeddingFunc import numpy as np +from lightrag.kg.shared_storage import initialize_pipeline_status print(os.getcwd()) @@ -113,6 +114,9 @@ async def init(): vector_storage="OracleVectorDBStorage", ) + await rag.initialize_storages() + await initialize_pipeline_status() + return rag diff --git a/examples/lightrag_azure_openai_demo.py b/examples/lightrag_azure_openai_demo.py index 1c941f93..e0840366 100644 --- a/examples/lightrag_azure_openai_demo.py +++ b/examples/lightrag_azure_openai_demo.py @@ -6,6 +6,7 @@ import numpy as np from dotenv import load_dotenv import logging from openai import AzureOpenAI +from lightrag.kg.shared_storage import initialize_pipeline_status logging.basicConfig(level=logging.INFO) @@ -90,6 +91,9 @@ rag = LightRAG( ), ) +rag.initialize_storages() +initialize_pipeline_status() + book1 = open("./book_1.txt", encoding="utf-8") book2 = open("./book_2.txt", encoding="utf-8") diff --git a/examples/lightrag_bedrock_demo.py b/examples/lightrag_bedrock_demo.py index 6bb6c7d4..700b4391 100644 --- a/examples/lightrag_bedrock_demo.py +++ b/examples/lightrag_bedrock_demo.py @@ -8,6 +8,12 @@ import logging from lightrag import LightRAG, QueryParam from lightrag.llm.bedrock import bedrock_complete, bedrock_embed from lightrag.utils import EmbeddingFunc +from lightrag.kg.shared_storage import initialize_pipeline_status + +import asyncio +import nest_asyncio + +nest_asyncio.apply() logging.getLogger("aiobotocore").setLevel(logging.WARNING) @@ -15,22 +21,31 @@ WORKING_DIR = "./dickens" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=bedrock_complete, - llm_model_name="Anthropic Claude 3 Haiku // Amazon Bedrock", - embedding_func=EmbeddingFunc( - embedding_dim=1024, max_token_size=8192, func=bedrock_embed - ), -) - -with open("./book.txt", "r", encoding="utf-8") as f: - rag.insert(f.read()) - -for mode in ["naive", "local", "global", "hybrid"]: - print("\n+-" + "-" * len(mode) + "-+") - print(f"| {mode.capitalize()} |") - print("+-" + "-" * len(mode) + "-+\n") - print( - rag.query("What are the top themes in this story?", param=QueryParam(mode=mode)) +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=bedrock_complete, + llm_model_name="Anthropic Claude 3 Haiku // Amazon Bedrock", + embedding_func=EmbeddingFunc( + embedding_dim=1024, max_token_size=8192, func=bedrock_embed + ), ) + + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag + +def main(): + rag = asyncio.run(initialize_rag()) + + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) + + for mode in ["naive", "local", "global", "hybrid"]: + print("\n+-" + "-" * len(mode) + "-+") + print(f"| {mode.capitalize()} |") + print("+-" + "-" * len(mode) + "-+\n") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode=mode)) + ) diff --git a/examples/lightrag_gemini_demo.py b/examples/lightrag_gemini_demo.py index 32732ba8..d20c0d0d 100644 --- a/examples/lightrag_gemini_demo.py +++ b/examples/lightrag_gemini_demo.py @@ -8,6 +8,12 @@ from dotenv import load_dotenv from lightrag.utils import EmbeddingFunc from lightrag import LightRAG, QueryParam from sentence_transformers import SentenceTransformer +from lightrag.kg.shared_storage import initialize_pipeline_status + +import asyncio +import nest_asyncio +# Apply nest_asyncio to solve event loop issues +nest_asyncio.apply() load_dotenv() gemini_api_key = os.getenv("GEMINI_API_KEY") @@ -60,25 +66,37 @@ async def embedding_func(texts: list[str]) -> np.ndarray: return embeddings -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=384, - max_token_size=8192, - func=embedding_func, - ), -) +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=llm_model_func, + embedding_func=EmbeddingFunc( + embedding_dim=384, + max_token_size=8192, + func=embedding_func, + ), + ) -file_path = "story.txt" -with open(file_path, "r") as file: - text = file.read() + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag -rag.insert(text) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) + file_path = "story.txt" + with open(file_path, "r") as file: + text = file.read() -response = rag.query( - query="What is the main theme of the story?", - param=QueryParam(mode="hybrid", top_k=5, response_type="single line"), -) + rag.insert(text) -print(response) + response = rag.query( + query="What is the main theme of the story?", + param=QueryParam(mode="hybrid", top_k=5, response_type="single line"), + ) + + print(response) + +if __name__ == "__main__": + main() diff --git a/examples/lightrag_hf_demo.py b/examples/lightrag_hf_demo.py index a5088e54..5cd214fd 100644 --- a/examples/lightrag_hf_demo.py +++ b/examples/lightrag_hf_demo.py @@ -4,51 +4,68 @@ from lightrag import LightRAG, QueryParam from lightrag.llm.hf import hf_model_complete, hf_embed from lightrag.utils import EmbeddingFunc from transformers import AutoModel, AutoTokenizer +from lightrag.kg.shared_storage import initialize_pipeline_status + +import asyncio +import nest_asyncio + +nest_asyncio.apply() WORKING_DIR = "./dickens" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=hf_model_complete, - llm_model_name="meta-llama/Llama-3.1-8B-Instruct", - embedding_func=EmbeddingFunc( - embedding_dim=384, - max_token_size=5000, - func=lambda texts: hf_embed( - texts, - tokenizer=AutoTokenizer.from_pretrained( - "sentence-transformers/all-MiniLM-L6-v2" - ), - embed_model=AutoModel.from_pretrained( - "sentence-transformers/all-MiniLM-L6-v2" +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=hf_model_complete, + llm_model_name="meta-llama/Llama-3.1-8B-Instruct", + embedding_func=EmbeddingFunc( + embedding_dim=384, + max_token_size=5000, + func=lambda texts: hf_embed( + texts, + tokenizer=AutoTokenizer.from_pretrained( + "sentence-transformers/all-MiniLM-L6-v2" + ), + embed_model=AutoModel.from_pretrained( + "sentence-transformers/all-MiniLM-L6-v2" + ), ), ), - ), -) + ) + await rag.initialize_storages() + await initialize_pipeline_status() -with open("./book.txt", "r", encoding="utf-8") as f: - rag.insert(f.read()) + return rag -# Perform naive search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) -) +def main(): + rag = asyncio.run(initialize_rag()) -# Perform local search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) -) + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) -# Perform global search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) -) + # Perform naive search + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + ) -# Perform hybrid search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) -) + # 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")) + ) + +if __name__ == "__main__": + main() diff --git a/examples/lightrag_jinaai_demo.py b/examples/lightrag_jinaai_demo.py deleted file mode 100644 index 0378b61b..00000000 --- a/examples/lightrag_jinaai_demo.py +++ /dev/null @@ -1,115 +0,0 @@ -import numpy as np -from lightrag import LightRAG, QueryParam -from lightrag.utils import EmbeddingFunc -from lightrag.llm.jina import jina_embed -from lightrag.llm.openai import openai_complete_if_cache -import os -import asyncio - - -async def embedding_func(texts: list[str]) -> np.ndarray: - return await jina_embed(texts, api_key="YourJinaAPIKey") - - -WORKING_DIR = "./dickens" - -if not os.path.exists(WORKING_DIR): - os.mkdir(WORKING_DIR) - - -async def llm_model_func( - prompt, system_prompt=None, history_messages=[], **kwargs -) -> str: - return await openai_complete_if_cache( - "solar-mini", - prompt, - system_prompt=system_prompt, - history_messages=history_messages, - api_key=os.getenv("UPSTAGE_API_KEY"), - base_url="https://api.upstage.ai/v1/solar", - **kwargs, - ) - - -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=1024, max_token_size=8192, func=embedding_func - ), -) - - -async def lightraginsert(file_path, semaphore): - async with semaphore: - try: - with open(file_path, "r", encoding="utf-8") as f: - content = f.read() - except UnicodeDecodeError: - # If UTF-8 decoding fails, try other encodings - with open(file_path, "r", encoding="gbk") as f: - content = f.read() - await rag.ainsert(content) - - -async def process_files(directory, concurrency_limit): - semaphore = asyncio.Semaphore(concurrency_limit) - tasks = [] - for root, dirs, files in os.walk(directory): - for f in files: - file_path = os.path.join(root, f) - if f.startswith("."): - continue - tasks.append(lightraginsert(file_path, semaphore)) - await asyncio.gather(*tasks) - - -async def main(): - try: - rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=1024, - max_token_size=8192, - func=embedding_func, - ), - ) - - asyncio.run(process_files(WORKING_DIR, concurrency_limit=4)) - - # Perform naive search - print( - await rag.aquery( - "What are the top themes in this story?", param=QueryParam(mode="naive") - ) - ) - - # Perform local search - print( - await rag.aquery( - "What are the top themes in this story?", param=QueryParam(mode="local") - ) - ) - - # Perform global search - print( - await rag.aquery( - "What are the top themes in this story?", - param=QueryParam(mode="global"), - ) - ) - - # Perform hybrid search - print( - await rag.aquery( - "What are the top themes in this story?", - param=QueryParam(mode="hybrid"), - ) - ) - except Exception as e: - print(f"An error occurred: {e}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/examples/lightrag_llamaindex_direct_demo.py b/examples/lightrag_llamaindex_direct_demo.py index 5db158ce..d1e68233 100644 --- a/examples/lightrag_llamaindex_direct_demo.py +++ b/examples/lightrag_llamaindex_direct_demo.py @@ -8,6 +8,11 @@ from lightrag.utils import EmbeddingFunc from llama_index.llms.openai import OpenAI from llama_index.embeddings.openai import OpenAIEmbedding import asyncio +import nest_asyncio + +nest_asyncio.apply() + +from lightrag.kg.shared_storage import initialize_pipeline_status # Configure working directory WORKING_DIR = "./index_default" @@ -76,38 +81,53 @@ async def get_embedding_dim(): return embedding_dim -# Initialize RAG instance -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=asyncio.run(get_embedding_dim()), - max_token_size=EMBEDDING_MAX_TOKEN_SIZE, - func=embedding_func, - ), -) +async def initialize_rag(): + embedding_dimension = await get_embedding_dim() + + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=llm_model_func, + embedding_func=EmbeddingFunc( + embedding_dim=embedding_dimension, + max_token_size=EMBEDDING_MAX_TOKEN_SIZE, + func=embedding_func, + ), + ) -# Insert example text -with open("./book.txt", "r", encoding="utf-8") as f: - rag.insert(f.read()) + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag -# Test different query modes -print("\nNaive Search:") -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) -) -print("\nLocal Search:") -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) -) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) -print("\nGlobal Search:") -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) -) + # Insert example text + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) -print("\nHybrid Search:") -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) -) + # Test different query modes + print("\nNaive Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + ) + + print("\nLocal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + ) + + print("\nGlobal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + ) + + print("\nHybrid Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) + ) + +if __name__ == "__main__": + main() diff --git a/examples/lightrag_llamaindex_litellm_demo.py b/examples/lightrag_llamaindex_litellm_demo.py index 3511ecf3..d7c609b4 100644 --- a/examples/lightrag_llamaindex_litellm_demo.py +++ b/examples/lightrag_llamaindex_litellm_demo.py @@ -8,6 +8,11 @@ from lightrag.utils import EmbeddingFunc from llama_index.llms.litellm import LiteLLM from llama_index.embeddings.litellm import LiteLLMEmbedding import asyncio +import nest_asyncio + +nest_asyncio.apply() + +from lightrag.kg.shared_storage import initialize_pipeline_status # Configure working directory WORKING_DIR = "./index_default" @@ -79,38 +84,53 @@ async def get_embedding_dim(): return embedding_dim -# Initialize RAG instance -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=asyncio.run(get_embedding_dim()), - max_token_size=EMBEDDING_MAX_TOKEN_SIZE, - func=embedding_func, - ), -) +async def initialize_rag(): + embedding_dimension = await get_embedding_dim() + + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=llm_model_func, + embedding_func=EmbeddingFunc( + embedding_dim=embedding_dimension, + max_token_size=EMBEDDING_MAX_TOKEN_SIZE, + func=embedding_func, + ), + ) -# Insert example text -with open("./book.txt", "r", encoding="utf-8") as f: - rag.insert(f.read()) + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag -# Test different query modes -print("\nNaive Search:") -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) -) -print("\nLocal Search:") -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) -) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) -print("\nGlobal Search:") -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) -) + # Insert example text + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) -print("\nHybrid Search:") -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) -) + # Test different query modes + print("\nNaive Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + ) + + print("\nLocal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + ) + + print("\nGlobal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + ) + + print("\nHybrid Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) + ) + +if __name__ == "__main__": + main() diff --git a/examples/lightrag_lmdeploy_demo.py b/examples/lightrag_lmdeploy_demo.py index d12eb564..e640a613 100644 --- a/examples/lightrag_lmdeploy_demo.py +++ b/examples/lightrag_lmdeploy_demo.py @@ -5,6 +5,12 @@ from lightrag.llm.lmdeploy import lmdeploy_model_if_cache from lightrag.llm.hf import hf_embed from lightrag.utils import EmbeddingFunc from transformers import AutoModel, AutoTokenizer +from lightrag.kg.shared_storage import initialize_pipeline_status + +import asyncio +import nest_asyncio + +nest_asyncio.apply() WORKING_DIR = "./dickens" @@ -35,46 +41,59 @@ async def lmdeploy_model_complete( **kwargs, ) - -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=lmdeploy_model_complete, - llm_model_name="meta-llama/Llama-3.1-8B-Instruct", # please use definite path for local model - embedding_func=EmbeddingFunc( - embedding_dim=384, - max_token_size=5000, - func=lambda texts: hf_embed( - texts, - tokenizer=AutoTokenizer.from_pretrained( - "sentence-transformers/all-MiniLM-L6-v2" - ), - embed_model=AutoModel.from_pretrained( - "sentence-transformers/all-MiniLM-L6-v2" +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=lmdeploy_model_complete, + llm_model_name="meta-llama/Llama-3.1-8B-Instruct", # please use definite path for local model + embedding_func=EmbeddingFunc( + embedding_dim=384, + max_token_size=5000, + func=lambda texts: hf_embed( + texts, + tokenizer=AutoTokenizer.from_pretrained( + "sentence-transformers/all-MiniLM-L6-v2" + ), + embed_model=AutoModel.from_pretrained( + "sentence-transformers/all-MiniLM-L6-v2" + ), ), ), - ), -) + ) + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag -with open("./book.txt", "r", encoding="utf-8") as f: - rag.insert(f.read()) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) -# Perform naive search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) -) + # Insert example text + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) -# Perform local search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) -) + # Test different query modes + print("\nNaive Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + ) -# Perform global search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) -) + print("\nLocal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + ) -# Perform hybrid search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) -) + print("\nGlobal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + ) + + print("\nHybrid Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) + ) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/lightrag_nvidia_demo.py b/examples/lightrag_nvidia_demo.py index da4b46ff..9137c1b6 100644 --- a/examples/lightrag_nvidia_demo.py +++ b/examples/lightrag_nvidia_demo.py @@ -1,5 +1,9 @@ import os import asyncio +import nest_asyncio + +nest_asyncio.apply() + from lightrag import LightRAG, QueryParam from lightrag.llm import ( openai_complete_if_cache, @@ -7,10 +11,12 @@ from lightrag.llm import ( ) from lightrag.utils import EmbeddingFunc import numpy as np +from lightrag.kg.shared_storage import initialize_pipeline_status # for custom llm_model_func from lightrag.utils import locate_json_string_body_from_string + WORKING_DIR = "./dickens" if not os.path.exists(WORKING_DIR): @@ -91,42 +97,37 @@ async def test_funcs(): # asyncio.run(test_funcs()) +async def initialize_rag(): + embedding_dimension = await get_embedding_dim() + print(f"Detected embedding dimension: {embedding_dimension}") + # lightRAG class during indexing + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=llm_model_func, + # llm_model_name="meta/llama3-70b-instruct", #un comment if + embedding_func=EmbeddingFunc( + embedding_dim=embedding_dimension, + max_token_size=512, # maximum token size, somehow it's still exceed maximum number of token + # so truncate (trunc) parameter on embedding_func will handle it and try to examine the tokenizer used in LightRAG + # so you can adjust to be able to fit the NVIDIA model (future work) + func=indexing_embedding_func, + ), + ) + + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag async def main(): try: - embedding_dimension = await get_embedding_dim() - print(f"Detected embedding dimension: {embedding_dimension}") - - # lightRAG class during indexing - rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - # llm_model_name="meta/llama3-70b-instruct", #un comment if - embedding_func=EmbeddingFunc( - embedding_dim=embedding_dimension, - max_token_size=512, # maximum token size, somehow it's still exceed maximum number of token - # so truncate (trunc) parameter on embedding_func will handle it and try to examine the tokenizer used in LightRAG - # so you can adjust to be able to fit the NVIDIA model (future work) - func=indexing_embedding_func, - ), - ) + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) # reading file with open("./book.txt", "r", encoding="utf-8") as f: await rag.ainsert(f.read()) - # redefine rag to change embedding into query type - rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - # llm_model_name="meta/llama3-70b-instruct", #un comment if - embedding_func=EmbeddingFunc( - embedding_dim=embedding_dimension, - max_token_size=512, - func=query_embedding_func, - ), - ) - # Perform naive search print("==============Naive===============") print( diff --git a/examples/lightrag_ollama_age_demo.py b/examples/lightrag_ollama_age_demo.py index d394ded4..22e42190 100644 --- a/examples/lightrag_ollama_age_demo.py +++ b/examples/lightrag_ollama_age_demo.py @@ -1,4 +1,8 @@ import asyncio +import nest_asyncio + +nest_asyncio.apply() + import inspect import logging import os @@ -6,6 +10,7 @@ import os from lightrag import LightRAG, QueryParam from lightrag.llm.ollama import ollama_embed, ollama_model_complete from lightrag.utils import EmbeddingFunc +from lightrag.kg.shared_storage import initialize_pipeline_status WORKING_DIR = "./dickens_age" @@ -22,59 +27,72 @@ 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_embed( - texts, embed_model="nomic-embed-text", host="http://localhost:11434" +async def initialize_rag(): + 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_embed( + 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), -) + graph_storage="AGEStorage", + ) + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag async def print_stream(stream): async for chunk in stream: print(chunk, end="", flush=True) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) -if inspect.isasyncgen(resp): - asyncio.run(print_stream(resp)) -else: - print(resp) + # Insert example text + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) + + # Test different query modes + print("\nNaive Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + ) + + print("\nLocal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + ) + + print("\nGlobal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + ) + + print("\nHybrid 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), + ) + + if inspect.isasyncgen(resp): + asyncio.run(print_stream(resp)) + else: + print(resp) + +if __name__ == "__main__": + main() diff --git a/examples/lightrag_ollama_demo.py b/examples/lightrag_ollama_demo.py index 95856fa2..6715ea72 100644 --- a/examples/lightrag_ollama_demo.py +++ b/examples/lightrag_ollama_demo.py @@ -1,10 +1,14 @@ import asyncio +import nest_asyncio + +nest_asyncio.apply() import os import inspect import logging from lightrag import LightRAG, QueryParam from lightrag.llm.ollama import ollama_model_complete, ollama_embed from lightrag.utils import EmbeddingFunc +from lightrag.kg.shared_storage import initialize_pipeline_status WORKING_DIR = "./dickens" @@ -13,58 +17,71 @@ logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO) if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=ollama_model_complete, - llm_model_name="gemma2:2b", - 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_embed( - texts, embed_model="nomic-embed-text", host="http://localhost:11434" +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=ollama_model_complete, + llm_model_name="gemma2:2b", + 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_embed( + texts, embed_model="nomic-embed-text", host="http://localhost:11434" + ), ), - ), -) - -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), -) + ) + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag async def print_stream(stream): async for chunk in stream: print(chunk, end="", flush=True) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) -if inspect.isasyncgen(resp): - asyncio.run(print_stream(resp)) -else: - print(resp) + # Insert example text + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) + + # Test different query modes + print("\nNaive Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + ) + + print("\nLocal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + ) + + print("\nGlobal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + ) + + print("\nHybrid 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), + ) + + if inspect.isasyncgen(resp): + asyncio.run(print_stream(resp)) + else: + print(resp) + +if __name__ == "__main__": + main() diff --git a/examples/lightrag_ollama_gremlin_demo.py b/examples/lightrag_ollama_gremlin_demo.py index fa7d4fb5..4d657afa 100644 --- a/examples/lightrag_ollama_gremlin_demo.py +++ b/examples/lightrag_ollama_gremlin_demo.py @@ -12,6 +12,7 @@ import os from lightrag import LightRAG, QueryParam from lightrag.llm.ollama import ollama_embed, ollama_model_complete from lightrag.utils import EmbeddingFunc +from lightrag.kg.shared_storage import initialize_pipeline_status WORKING_DIR = "./dickens_gremlin" @@ -31,59 +32,72 @@ os.environ["GREMLIN_TRAVERSE_SOURCE"] = "g" os.environ["GREMLIN_USER"] = "" os.environ["GREMLIN_PASSWORD"] = "" -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_embed( - texts, embed_model="nomic-embed-text", host="http://localhost:11434" +async def initialize_rag(): + 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_embed( + texts, embed_model="nomic-embed-text", host="http://localhost:11434" + ), ), - ), - graph_storage="GremlinStorage", -) - -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), -) + graph_storage="GremlinStorage", + ) + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag async def print_stream(stream): async for chunk in stream: print(chunk, end="", flush=True) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) -if inspect.isasyncgen(resp): - asyncio.run(print_stream(resp)) -else: - print(resp) + # Insert example text + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) + + # Test different query modes + print("\nNaive Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + ) + + print("\nLocal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + ) + + print("\nGlobal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + ) + + print("\nHybrid 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), + ) + + if inspect.isasyncgen(resp): + asyncio.run(print_stream(resp)) + else: + print(resp) + +if __name__ == "__main__": + main() diff --git a/examples/lightrag_ollama_neo4j_milvus_mongo_demo.py b/examples/lightrag_ollama_neo4j_milvus_mongo_demo.py index 7e2d1026..e5d4064d 100644 --- a/examples/lightrag_ollama_neo4j_milvus_mongo_demo.py +++ b/examples/lightrag_ollama_neo4j_milvus_mongo_demo.py @@ -2,6 +2,11 @@ import os from lightrag import LightRAG, QueryParam from lightrag.llm.ollama import ollama_model_complete, ollama_embed from lightrag.utils import EmbeddingFunc +import asyncio +import nest_asyncio + +nest_asyncio.apply() +from lightrag.kg.shared_storage import initialize_pipeline_status # WorkingDir ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) @@ -27,30 +32,59 @@ os.environ["MILVUS_USER"] = "root" os.environ["MILVUS_PASSWORD"] = "root" os.environ["MILVUS_DB_NAME"] = "lightrag" - -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=ollama_model_complete, - llm_model_name="qwen2.5:14b", - llm_model_max_async=4, - llm_model_max_token_size=32768, - llm_model_kwargs={"host": "http://127.0.0.1:11434", "options": {"num_ctx": 32768}}, - embedding_func=EmbeddingFunc( - embedding_dim=1024, - max_token_size=8192, - func=lambda texts: ollama_embed( - texts=texts, embed_model="bge-m3:latest", host="http://127.0.0.1:11434" +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=ollama_model_complete, + llm_model_name="qwen2.5:14b", + llm_model_max_async=4, + llm_model_max_token_size=32768, + llm_model_kwargs={"host": "http://127.0.0.1:11434", "options": {"num_ctx": 32768}}, + embedding_func=EmbeddingFunc( + embedding_dim=1024, + max_token_size=8192, + func=lambda texts: ollama_embed( + texts=texts, embed_model="bge-m3:latest", host="http://127.0.0.1:11434" + ), ), - ), - kv_storage="MongoKVStorage", - graph_storage="Neo4JStorage", - vector_storage="MilvusVectorDBStorage", -) + kv_storage="MongoKVStorage", + graph_storage="Neo4JStorage", + vector_storage="MilvusVectorDBStorage", + ) -file = "./book.txt" -with open(file, "r") as f: - rag.insert(f.read()) + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) -) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) + + # Insert example text + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) + + # Test different query modes + print("\nNaive Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + ) + + print("\nLocal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + ) + + print("\nGlobal Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + ) + + print("\nHybrid Search:") + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) + ) + +if __name__ == "__main__": + main() diff --git a/examples/lightrag_openai_compatible_demo.py b/examples/lightrag_openai_compatible_demo.py index 09673dd8..f4af7be6 100644 --- a/examples/lightrag_openai_compatible_demo.py +++ b/examples/lightrag_openai_compatible_demo.py @@ -4,6 +4,7 @@ from lightrag import LightRAG, QueryParam from lightrag.llm.openai import openai_complete_if_cache, openai_embed from lightrag.utils import EmbeddingFunc import numpy as np +from lightrag.kg.shared_storage import initialize_pipeline_status WORKING_DIR = "./dickens" @@ -52,21 +53,28 @@ async def test_funcs(): # asyncio.run(test_funcs()) +async def initialize_rag(): + embedding_dimension = await get_embedding_dim() + print(f"Detected embedding dimension: {embedding_dimension}") + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=llm_model_func, + embedding_func=EmbeddingFunc( + embedding_dim=embedding_dimension, + max_token_size=8192, + func=embedding_func, + ), + ) + + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag async def main(): try: - embedding_dimension = await get_embedding_dim() - print(f"Detected embedding dimension: {embedding_dimension}") - - rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=embedding_dimension, - max_token_size=8192, - func=embedding_func, - ), - ) + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) with open("./book.txt", "r", encoding="utf-8") as f: await rag.ainsert(f.read()) diff --git a/examples/lightrag_openai_compatible_demo_embedding_cache.py b/examples/lightrag_openai_compatible_demo_embedding_cache.py index d696ce25..fcdd1ef3 100644 --- a/examples/lightrag_openai_compatible_demo_embedding_cache.py +++ b/examples/lightrag_openai_compatible_demo_embedding_cache.py @@ -4,6 +4,7 @@ from lightrag import LightRAG, QueryParam from lightrag.llm.openai import openai_complete_if_cache, openai_embed from lightrag.utils import EmbeddingFunc import numpy as np +from lightrag.kg.shared_storage import initialize_pipeline_status WORKING_DIR = "./dickens" @@ -52,25 +53,33 @@ async def test_funcs(): # asyncio.run(test_funcs()) +async def initialize_rag(): + embedding_dimension = await get_embedding_dim() + print(f"Detected embedding dimension: {embedding_dimension}") + + rag = LightRAG( + working_dir=WORKING_DIR, + embedding_cache_config={ + "enabled": True, + "similarity_threshold": 0.90, + }, + llm_model_func=llm_model_func, + embedding_func=EmbeddingFunc( + embedding_dim=embedding_dimension, + max_token_size=8192, + func=embedding_func, + ), + ) + + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag async def main(): try: - embedding_dimension = await get_embedding_dim() - print(f"Detected embedding dimension: {embedding_dimension}") - - rag = LightRAG( - working_dir=WORKING_DIR, - embedding_cache_config={ - "enabled": True, - "similarity_threshold": 0.90, - }, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=embedding_dimension, - max_token_size=8192, - func=embedding_func, - ), - ) + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) with open("./book.txt", "r", encoding="utf-8") as f: await rag.ainsert(f.read()) diff --git a/examples/lightrag_openai_compatible_stream_demo.py b/examples/lightrag_openai_compatible_stream_demo.py index a974ca14..b3f237e5 100644 --- a/examples/lightrag_openai_compatible_stream_demo.py +++ b/examples/lightrag_openai_compatible_stream_demo.py @@ -1,9 +1,11 @@ import inspect import os +import asyncio from lightrag import LightRAG from lightrag.llm import openai_complete, openai_embed from lightrag.utils import EmbeddingFunc, always_get_an_event_loop from lightrag import QueryParam +from lightrag.kg.shared_storage import initialize_pipeline_status # WorkingDir ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) @@ -13,42 +15,54 @@ if not os.path.exists(WORKING_DIR): print(f"WorkingDir: {WORKING_DIR}") api_key = "empty" -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=openai_complete, - llm_model_name="qwen2.5-14b-instruct@4bit", - llm_model_max_async=4, - llm_model_max_token_size=32768, - llm_model_kwargs={"base_url": "http://127.0.0.1:1234/v1", "api_key": api_key}, - embedding_func=EmbeddingFunc( - embedding_dim=1024, - max_token_size=8192, - func=lambda texts: openai_embed( - texts=texts, - model="text-embedding-bge-m3", - base_url="http://127.0.0.1:1234/v1", - api_key=api_key, +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=openai_complete, + llm_model_name="qwen2.5-14b-instruct@4bit", + llm_model_max_async=4, + llm_model_max_token_size=32768, + llm_model_kwargs={"base_url": "http://127.0.0.1:1234/v1", "api_key": api_key}, + embedding_func=EmbeddingFunc( + embedding_dim=1024, + max_token_size=8192, + func=lambda texts: openai_embed( + texts=texts, + model="text-embedding-bge-m3", + base_url="http://127.0.0.1:1234/v1", + api_key=api_key, + ), ), - ), -) + ) -with open("./book.txt", "r", encoding="utf-8") as f: - rag.insert(f.read()) - -resp = rag.query( - "What are the top themes in this story?", - param=QueryParam(mode="hybrid", stream=True), -) + await rag.initialize_storages() + await initialize_pipeline_status() + return rag async def print_stream(stream): async for chunk in stream: if chunk: print(chunk, end="", flush=True) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) + + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) + + resp = rag.query( + "What are the top themes in this story?", + param=QueryParam(mode="hybrid", stream=True), + ) + + loop = always_get_an_event_loop() + if inspect.isasyncgen(resp): + loop.run_until_complete(print_stream(resp)) + else: + print(resp) + +if __name__ == "__main__": + main() -loop = always_get_an_event_loop() -if inspect.isasyncgen(resp): - loop.run_until_complete(print_stream(resp)) -else: - print(resp) diff --git a/examples/lightrag_openai_demo.py b/examples/lightrag_openai_demo.py index c5393fc8..f5f47ee2 100644 --- a/examples/lightrag_openai_demo.py +++ b/examples/lightrag_openai_demo.py @@ -1,40 +1,54 @@ import os - +import asyncio from lightrag import LightRAG, QueryParam from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed +from lightrag.kg.shared_storage import initialize_pipeline_status WORKING_DIR = "./dickens" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) -rag = LightRAG( - working_dir=WORKING_DIR, - embedding_func=openai_embed, - llm_model_func=gpt_4o_mini_complete, - # llm_model_func=gpt_4o_complete -) +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + embedding_func=openai_embed, + llm_model_func=gpt_4o_mini_complete, + # llm_model_func=gpt_4o_complete + ) + await rag.initialize_storages() + await initialize_pipeline_status() -with open("./book.txt", "r", encoding="utf-8") as f: - rag.insert(f.read()) + return rag -# Perform naive search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) -) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) -# Perform local search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) -) + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) -# Perform global search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) -) + # 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")) + ) + +if __name__ == "__main__": + main() -# Perform hybrid search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) -) diff --git a/examples/lightrag_openai_mongodb_graph_demo.py b/examples/lightrag_openai_mongodb_graph_demo.py index 775eb296..ddf5ca63 100644 --- a/examples/lightrag_openai_mongodb_graph_demo.py +++ b/examples/lightrag_openai_mongodb_graph_demo.py @@ -4,6 +4,7 @@ from lightrag import LightRAG, QueryParam from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed from lightrag.utils import EmbeddingFunc import numpy as np +from lightrag.kg.shared_storage import initialize_pipeline_status ######### # Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert() @@ -52,7 +53,7 @@ async def create_embedding_function_instance(): async def initialize_rag(): embedding_func_instance = await create_embedding_function_instance() - return LightRAG( + rag = LightRAG( working_dir=WORKING_DIR, llm_model_func=gpt_4o_mini_complete, embedding_func=embedding_func_instance, @@ -60,14 +61,38 @@ async def initialize_rag(): log_level="DEBUG", ) + await rag.initialize_storages() + await initialize_pipeline_status() -# Run the initialization -rag = asyncio.run(initialize_rag()) + return rag -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")) -) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) + + 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")) + ) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/lightrag_openai_neo4j_milvus_redis_demo.py b/examples/lightrag_openai_neo4j_milvus_redis_demo.py index da5c5a8f..adf87691 100644 --- a/examples/lightrag_openai_neo4j_milvus_redis_demo.py +++ b/examples/lightrag_openai_neo4j_milvus_redis_demo.py @@ -1,7 +1,9 @@ import os +import asyncio from lightrag import LightRAG, QueryParam from lightrag.llm.ollama import ollama_embed, openai_complete_if_cache from lightrag.utils import EmbeddingFunc +from lightrag.kg.shared_storage import initialize_pipeline_status # WorkingDir ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) @@ -48,23 +50,52 @@ embedding_func = EmbeddingFunc( texts, embed_model="shaw/dmeta-embedding-zh", host="http://117.50.173.35:11434" ), ) +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=llm_model_func, + llm_model_max_token_size=32768, + embedding_func=embedding_func, + chunk_token_size=512, + chunk_overlap_token_size=256, + kv_storage="RedisKVStorage", + graph_storage="Neo4JStorage", + vector_storage="MilvusVectorDBStorage", + doc_status_storage="RedisKVStorage", + ) -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - llm_model_max_token_size=32768, - embedding_func=embedding_func, - chunk_token_size=512, - chunk_overlap_token_size=256, - kv_storage="RedisKVStorage", - graph_storage="Neo4JStorage", - vector_storage="MilvusVectorDBStorage", - doc_status_storage="RedisKVStorage", -) + await rag.initialize_storages() + await initialize_pipeline_status() -file = "../book.txt" -with open(file, "r", encoding="utf-8") as f: - rag.insert(f.read()) + return rag -print(rag.query("谁会3D建模 ？", param=QueryParam(mode="mix"))) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) + + 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")) + ) + +if __name__ == "__main__": + main() diff --git a/examples/lightrag_oracle_demo.py b/examples/lightrag_oracle_demo.py index c6121840..53139220 100644 --- a/examples/lightrag_oracle_demo.py +++ b/examples/lightrag_oracle_demo.py @@ -6,6 +6,7 @@ from lightrag import LightRAG, QueryParam from lightrag.llm.openai import openai_complete_if_cache, openai_embed from lightrag.utils import EmbeddingFunc import numpy as np +from lightrag.kg.shared_storage import initialize_pipeline_status print(os.getcwd()) script_directory = Path(__file__).resolve().parent.parent @@ -63,41 +64,48 @@ async def get_embedding_dim(): embedding_dim = embedding.shape[1] return embedding_dim +async def initialize_rag(): + # Detect embedding dimension + embedding_dimension = await get_embedding_dim() + print(f"Detected embedding dimension: {embedding_dimension}") + + # Initialize LightRAG + # We use Oracle DB as the KV/vector/graph storage + # You can add `addon_params={"example_number": 1, "language": "Simplfied Chinese"}` to control the prompt + rag = LightRAG( + # log_level="DEBUG", + working_dir=WORKING_DIR, + entity_extract_max_gleaning=1, + enable_llm_cache=True, + enable_llm_cache_for_entity_extract=True, + embedding_cache_config=None, # {"enabled": True,"similarity_threshold": 0.90}, + chunk_token_size=CHUNK_TOKEN_SIZE, + llm_model_max_token_size=MAX_TOKENS, + llm_model_func=llm_model_func, + embedding_func=EmbeddingFunc( + embedding_dim=embedding_dimension, + max_token_size=500, + func=embedding_func, + ), + graph_storage="OracleGraphStorage", + kv_storage="OracleKVStorage", + vector_storage="OracleVectorDBStorage", + addon_params={ + "example_number": 1, + "language": "Simplfied Chinese", + "entity_types": ["organization", "person", "geo", "event"], + "insert_batch_size": 2, + }, + ) + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag async def main(): try: - # Detect embedding dimension - embedding_dimension = await get_embedding_dim() - print(f"Detected embedding dimension: {embedding_dimension}") - - # Initialize LightRAG - # We use Oracle DB as the KV/vector/graph storage - # You can add `addon_params={"example_number": 1, "language": "Simplfied Chinese"}` to control the prompt - rag = LightRAG( - # log_level="DEBUG", - working_dir=WORKING_DIR, - entity_extract_max_gleaning=1, - enable_llm_cache=True, - enable_llm_cache_for_entity_extract=True, - embedding_cache_config=None, # {"enabled": True,"similarity_threshold": 0.90}, - chunk_token_size=CHUNK_TOKEN_SIZE, - llm_model_max_token_size=MAX_TOKENS, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=embedding_dimension, - max_token_size=500, - func=embedding_func, - ), - graph_storage="OracleGraphStorage", - kv_storage="OracleKVStorage", - vector_storage="OracleVectorDBStorage", - addon_params={ - "example_number": 1, - "language": "Simplfied Chinese", - "entity_types": ["organization", "person", "geo", "event"], - "insert_batch_size": 2, - }, - ) + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) # Extract and Insert into LightRAG storage with open(WORKING_DIR + "/docs.txt", "r", encoding="utf-8") as f: diff --git a/examples/lightrag_siliconcloud_demo.py b/examples/lightrag_siliconcloud_demo.py index 5f4f86a1..1deb5a66 100644 --- a/examples/lightrag_siliconcloud_demo.py +++ b/examples/lightrag_siliconcloud_demo.py @@ -5,6 +5,7 @@ from lightrag.llm.openai import openai_complete_if_cache from lightrag.llm.siliconcloud import siliconcloud_embedding from lightrag.utils import EmbeddingFunc import numpy as np +from lightrag.kg.shared_storage import initialize_pipeline_status WORKING_DIR = "./dickens" @@ -46,35 +47,48 @@ async def test_funcs(): asyncio.run(test_funcs()) +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=llm_model_func, + embedding_func=EmbeddingFunc( + embedding_dim=768, max_token_size=512, func=embedding_func + ), + ) -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=768, max_token_size=512, func=embedding_func - ), -) + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag -with open("./book.txt") as f: - rag.insert(f.read()) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) -# Perform naive search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) -) + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) -# Perform local search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) -) + # Perform naive search + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + ) -# Perform global search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) -) + # 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")) + ) + +if __name__ == "__main__": + main() -# Perform hybrid search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) -) diff --git a/examples/lightrag_tidb_demo.py b/examples/lightrag_tidb_demo.py index f2ee9ad8..c3f8fd19 100644 --- a/examples/lightrag_tidb_demo.py +++ b/examples/lightrag_tidb_demo.py @@ -6,6 +6,7 @@ import numpy as np from lightrag import LightRAG, QueryParam from lightrag.llm import siliconcloud_embedding, openai_complete_if_cache from lightrag.utils import EmbeddingFunc +from lightrag.kg.shared_storage import initialize_pipeline_status WORKING_DIR = "./dickens" @@ -54,33 +55,40 @@ async def get_embedding_dim(): embedding_dim = embedding.shape[1] return embedding_dim +async def initialize_rag(): + # Detect embedding dimension + embedding_dimension = await get_embedding_dim() + print(f"Detected embedding dimension: {embedding_dimension}") + + # Initialize LightRAG + # We use TiDB DB as the KV/vector + rag = LightRAG( + enable_llm_cache=False, + working_dir=WORKING_DIR, + chunk_token_size=512, + llm_model_func=llm_model_func, + embedding_func=EmbeddingFunc( + embedding_dim=embedding_dimension, + max_token_size=512, + func=embedding_func, + ), + kv_storage="TiDBKVStorage", + vector_storage="TiDBVectorDBStorage", + graph_storage="TiDBGraphStorage", + ) + + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag async def main(): try: - # Detect embedding dimension - embedding_dimension = await get_embedding_dim() - print(f"Detected embedding dimension: {embedding_dimension}") + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) - # Initialize LightRAG - # We use TiDB DB as the KV/vector - rag = LightRAG( - enable_llm_cache=False, - working_dir=WORKING_DIR, - chunk_token_size=512, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=embedding_dimension, - max_token_size=512, - func=embedding_func, - ), - kv_storage="TiDBKVStorage", - vector_storage="TiDBVectorDBStorage", - graph_storage="TiDBGraphStorage", - ) - - # Extract and Insert into LightRAG storage - with open("./dickens/demo.txt", "r", encoding="utf-8") as f: - await rag.ainsert(f.read()) + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) # Perform search in different modes modes = ["naive", "local", "global", "hybrid"] diff --git a/examples/lightrag_zhipu_demo.py b/examples/lightrag_zhipu_demo.py index 97a5042e..30b7316f 100644 --- a/examples/lightrag_zhipu_demo.py +++ b/examples/lightrag_zhipu_demo.py @@ -1,10 +1,12 @@ import os import logging +import asyncio from lightrag import LightRAG, QueryParam from lightrag.llm.zhipu import zhipu_complete, zhipu_embedding from lightrag.utils import EmbeddingFunc +from lightrag.kg.shared_storage import initialize_pipeline_status WORKING_DIR = "./dickens" @@ -17,39 +19,51 @@ api_key = os.environ.get("ZHIPUAI_API_KEY") if api_key is None: raise Exception("Please set ZHIPU_API_KEY in your environment") +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=zhipu_complete, + llm_model_name="glm-4-flashx", # Using the most cost/performance balance model, but you can change it here. + llm_model_max_async=4, + llm_model_max_token_size=32768, + embedding_func=EmbeddingFunc( + embedding_dim=2048, # Zhipu embedding-3 dimension + max_token_size=8192, + func=lambda texts: zhipu_embedding(texts), + ), + ) -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=zhipu_complete, - llm_model_name="glm-4-flashx", # Using the most cost/performance balance model, but you can change it here. - llm_model_max_async=4, - llm_model_max_token_size=32768, - embedding_func=EmbeddingFunc( - embedding_dim=2048, # Zhipu embedding-3 dimension - max_token_size=8192, - func=lambda texts: zhipu_embedding(texts), - ), -) + await rag.initialize_storages() + await initialize_pipeline_status() -with open("./book.txt", "r", encoding="utf-8") as f: - rag.insert(f.read()) + return rag -# Perform naive search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) -) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) -# Perform local search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) -) + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) -# Perform global search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) -) + # Perform naive search + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + ) -# Perform hybrid search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) -) + # 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")) + ) + +if __name__ == "__main__": + main() diff --git a/examples/lightrag_zhipu_postgres_demo.py b/examples/lightrag_zhipu_postgres_demo.py index 8f40690e..310786a5 100644 --- a/examples/lightrag_zhipu_postgres_demo.py +++ b/examples/lightrag_zhipu_postgres_demo.py @@ -8,6 +8,7 @@ from lightrag import LightRAG, QueryParam from lightrag.llm.zhipu import zhipu_complete from lightrag.llm.ollama import ollama_embedding from lightrag.utils import EmbeddingFunc +from lightrag.kg.shared_storage import initialize_pipeline_status load_dotenv() ROOT_DIR = os.environ.get("ROOT_DIR") @@ -27,8 +28,7 @@ os.environ["POSTGRES_USER"] = "rag" os.environ["POSTGRES_PASSWORD"] = "rag" os.environ["POSTGRES_DATABASE"] = "rag" - -async def main(): +async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, llm_model_func=zhipu_complete, @@ -50,9 +50,17 @@ async def main(): auto_manage_storages_states=False, ) + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag + +async def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) + # add embedding_func for graph database, it's deleted in commit 5661d76860436f7bf5aef2e50d9ee4a59660146c rag.chunk_entity_relation_graph.embedding_func = rag.embedding_func - await rag.initialize_storages() with open(f"{ROOT_DIR}/book.txt", "r", encoding="utf-8") as f: await rag.ainsert(f.read()) diff --git a/examples/query_keyword_separation_example.py b/examples/query_keyword_separation_example.py index f11ce8c1..de106de6 100644 --- a/examples/query_keyword_separation_example.py +++ b/examples/query_keyword_separation_example.py @@ -6,6 +6,7 @@ import numpy as np from dotenv import load_dotenv import logging from openai import AzureOpenAI +from lightrag.kg.shared_storage import initialize_pipeline_status logging.basicConfig(level=logging.INFO) @@ -79,25 +80,32 @@ async def test_funcs(): asyncio.run(test_funcs()) embedding_dimension = 3072 +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=llm_model_func, + embedding_func=EmbeddingFunc( + embedding_dim=embedding_dimension, + max_token_size=8192, + func=embedding_func, + ), + ) -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=embedding_dimension, - max_token_size=8192, - func=embedding_func, - ), -) + await rag.initialize_storages() + await initialize_pipeline_status() -book1 = open("./book_1.txt", encoding="utf-8") -book2 = open("./book_2.txt", encoding="utf-8") - -rag.insert([book1.read(), book2.read()]) + return rag # Example function demonstrating the new query_with_separate_keyword_extraction usage async def run_example(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) + + book1 = open("./book_1.txt", encoding="utf-8") + book2 = open("./book_2.txt", encoding="utf-8") + + rag.insert([book1.read(), book2.read()]) query = "What are the top themes in this story?" prompt = "Please simplify the response for a young audience." diff --git a/examples/test.py b/examples/test.py index 67ee22eb..dc186bda 100644 --- a/examples/test.py +++ b/examples/test.py @@ -1,6 +1,7 @@ import os from lightrag import LightRAG, QueryParam from lightrag.llm.openai import gpt_4o_mini_complete +from lightrag.kg.shared_storage import initialize_pipeline_status ######### # Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert() # import nest_asyncio @@ -12,31 +13,45 @@ WORKING_DIR = "./dickens" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model - # llm_model_func=gpt_4o_complete # Optionally, use a stronger model -) +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model + # llm_model_func=gpt_4o_complete # Optionally, use a stronger model + ) -with open("./dickens/book.txt", "r", encoding="utf-8") as f: - rag.insert(f.read()) + await rag.initialize_storages() + await initialize_pipeline_status() -# Perform naive search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) -) + return rag -# 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")) -) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) -# Perform hybrid search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) -) + 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")) + ) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/test_chromadb.py b/examples/test_chromadb.py index 99090a6d..10d69cc1 100644 --- a/examples/test_chromadb.py +++ b/examples/test_chromadb.py @@ -4,6 +4,7 @@ from lightrag import LightRAG, QueryParam from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed from lightrag.utils import EmbeddingFunc import numpy as np +from lightrag.kg.shared_storage import initialize_pipeline_status ######### # Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert() @@ -67,7 +68,7 @@ async def create_embedding_function_instance(): async def initialize_rag(): embedding_func_instance = await create_embedding_function_instance() if CHROMADB_USE_LOCAL_PERSISTENT: - return LightRAG( + rag = LightRAG( working_dir=WORKING_DIR, llm_model_func=gpt_4o_mini_complete, embedding_func=embedding_func_instance, @@ -87,7 +88,7 @@ async def initialize_rag(): }, ) else: - return LightRAG( + rag = LightRAG( working_dir=WORKING_DIR, llm_model_func=gpt_4o_mini_complete, embedding_func=embedding_func_instance, @@ -112,28 +113,36 @@ async def initialize_rag(): ) -# Run the initialization -rag = asyncio.run(initialize_rag()) + await rag.initialize_storages() + await initialize_pipeline_status() -# with open("./dickens/book.txt", "r", encoding="utf-8") as f: -# rag.insert(f.read()) + return rag -# Perform naive search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) -) + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) -# Perform local search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) -) + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) -# Perform global search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) -) + # Perform naive search + print( + rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + ) -# Perform hybrid search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) -) + # 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")) + ) + +if __name__ == "__main__": + main() diff --git a/examples/test_faiss.py b/examples/test_faiss.py index c3ac6f47..34991f47 100644 --- a/examples/test_faiss.py +++ b/examples/test_faiss.py @@ -1,5 +1,6 @@ import os import logging +import asyncio import numpy as np from dotenv import load_dotenv @@ -8,7 +9,9 @@ from sentence_transformers import SentenceTransformer from openai import AzureOpenAI from lightrag import LightRAG, QueryParam from lightrag.utils import EmbeddingFunc +from lightrag.kg.shared_storage import initialize_pipeline_status +WORKING_DIR = "./dickens" # Configure Logging logging.basicConfig(level=logging.INFO) @@ -55,11 +58,7 @@ async def embedding_func(texts: list[str]) -> np.ndarray: embeddings = model.encode(texts, convert_to_numpy=True) return embeddings - -def main(): - WORKING_DIR = "./dickens" - - # Initialize LightRAG with the LLM model function and embedding function +async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, llm_model_func=llm_model_func, @@ -74,6 +73,15 @@ def main(): }, ) + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag + +def main(): + + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) # Insert the custom chunks into LightRAG book1 = open("./book_1.txt", encoding="utf-8") book2 = open("./book_2.txt", encoding="utf-8") diff --git a/examples/test_neo4j.py b/examples/test_neo4j.py index ac5f7fb7..2d1d527a 100644 --- a/examples/test_neo4j.py +++ b/examples/test_neo4j.py @@ -1,7 +1,8 @@ import os +import asyncio from lightrag import LightRAG, QueryParam from lightrag.llm.openai import gpt_4o_mini_complete - +from lightrag.kg.shared_storage import initialize_pipeline_status ######### # Uncomment the below two lines if running in a jupyter notebook to handle the async nature of rag.insert() @@ -14,33 +15,46 @@ WORKING_DIR = "./local_neo4jWorkDir" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model - graph_storage="Neo4JStorage", - log_level="INFO", - # llm_model_func=gpt_4o_complete # Optionally, use a stronger model -) +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model + graph_storage="Neo4JStorage", + log_level="INFO", + # llm_model_func=gpt_4o_complete # Optionally, use a stronger model + ) -with open("./book.txt") as f: - rag.insert(f.read()) + await rag.initialize_storages() + await initialize_pipeline_status() -# Perform naive search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) -) + return rag -# Perform local search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) -) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) -# Perform global search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) -) + with open("./book.txt", "r", encoding="utf-8") as f: + rag.insert(f.read()) -# Perform hybrid search -print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) -) + # 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")) + ) + +if __name__ == "__main__": + main() diff --git a/examples/test_split_by_character.ipynb b/examples/test_split_by_character.ipynb index df5d938d..f70f9f34 100644 --- a/examples/test_split_by_character.ipynb +++ b/examples/test_split_by_character.ipynb @@ -18,6 +18,7 @@ "from lightrag import LightRAG, QueryParam\n", "from lightrag.llm.openai import openai_complete_if_cache, openai_embed\n", "from lightrag.utils import EmbeddingFunc\n", + "from lightrag.kg.shared_storage import initialize_pipeline_status\n", "import nest_asyncio" ] }, @@ -25,7 +26,9 @@ "cell_type": "markdown", "id": "dd17956ec322b361", "metadata": {}, - "source": "#### split by character" + "source": [ + "#### split by character" + ] }, { "cell_type": "code", @@ -109,14 +112,26 @@ } ], "source": [ - "rag = LightRAG(\n", - " working_dir=WORKING_DIR,\n", - " llm_model_func=llm_model_func,\n", - " embedding_func=EmbeddingFunc(\n", - " embedding_dim=4096, max_token_size=8192, func=embedding_func\n", - " ),\n", - " chunk_token_size=512,\n", - ")" + "import asyncio\n", + "import nest_asyncio\n", + "\n", + "nest_asyncio.apply()\n", + "\n", + "async def initialize_rag():\n", + " rag = LightRAG(\n", + " working_dir=WORKING_DIR,\n", + " llm_model_func=llm_model_func,\n", + " embedding_func=EmbeddingFunc(\n", + " embedding_dim=4096, max_token_size=8192, func=embedding_func\n", + " ),\n", + " chunk_token_size=512,\n", + " )\n", + " await rag.initialize_storages()\n", + " await initialize_pipeline_status()\n", + "\n", + " return rag\n", + "\n", + "rag = asyncio.run(initialize_rag())" ] }, { @@ -908,7 +923,9 @@ "cell_type": "markdown", "id": "4e5bfad24cb721a8", "metadata": {}, - "source": "#### split by character only" + "source": [ + "#### split by character only" + ] }, { "cell_type": "code", diff --git a/examples/vram_management_demo.py b/examples/vram_management_demo.py index b8d0872e..f4d46ab4 100644 --- a/examples/vram_management_demo.py +++ b/examples/vram_management_demo.py @@ -1,8 +1,10 @@ import os import time +import asyncio from lightrag import LightRAG, QueryParam from lightrag.llm.ollama import ollama_model_complete, ollama_embed from lightrag.utils import EmbeddingFunc +from lightrag.kg.shared_storage import initialize_pipeline_status # Working directory and the directory path for text files WORKING_DIR = "./dickens" @@ -12,17 +14,22 @@ TEXT_FILES_DIR = "/llm/mt" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) -# Initialize LightRAG -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=ollama_model_complete, - llm_model_name="qwen2.5:3b-instruct-max-context", - embedding_func=EmbeddingFunc( - embedding_dim=768, - max_token_size=8192, - func=lambda texts: ollama_embed(texts, embed_model="nomic-embed-text"), - ), -) +async def initialize_rag(): + # Initialize LightRAG + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=ollama_model_complete, + llm_model_name="qwen2.5:3b-instruct-max-context", + embedding_func=EmbeddingFunc( + embedding_dim=768, + max_token_size=8192, + func=lambda texts: ollama_embed(texts, embed_model="nomic-embed-text"), + ), + ) + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag # Read all .txt files from the TEXT_FILES_DIR directory texts = [] @@ -47,58 +54,65 @@ def insert_texts_with_retry(rag, texts, retries=3, delay=5): raise RuntimeError("Failed to insert texts after multiple retries.") -insert_texts_with_retry(rag, texts) +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) -# Perform different types of queries and handle potential errors -try: - print( - rag.query( - "What are the top themes in this story?", param=QueryParam(mode="naive") + insert_texts_with_retry(rag, texts) + + # Perform different types of queries and handle potential errors + try: + print( + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="naive") + ) ) - ) -except Exception as e: - print(f"Error performing naive search: {e}") + except Exception as e: + print(f"Error performing naive search: {e}") -try: - print( - rag.query( - "What are the top themes in this story?", param=QueryParam(mode="local") + try: + print( + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="local") + ) ) - ) -except Exception as e: - print(f"Error performing local search: {e}") + except Exception as e: + print(f"Error performing local search: {e}") -try: - print( - rag.query( - "What are the top themes in this story?", param=QueryParam(mode="global") + try: + print( + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="global") + ) ) - ) -except Exception as e: - print(f"Error performing global search: {e}") + except Exception as e: + print(f"Error performing global search: {e}") -try: - print( - rag.query( - "What are the top themes in this story?", param=QueryParam(mode="hybrid") + try: + print( + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) - ) -except Exception as e: - print(f"Error performing hybrid search: {e}") + except Exception as e: + print(f"Error performing hybrid search: {e}") -# Function to clear VRAM resources -def clear_vram(): - os.system("sudo nvidia-smi --gpu-reset") + # Function to clear VRAM resources + def clear_vram(): + os.system("sudo nvidia-smi --gpu-reset") -# Regularly clear VRAM to prevent overflow -clear_vram_interval = 3600 # Clear once every hour -start_time = time.time() + # Regularly clear VRAM to prevent overflow + clear_vram_interval = 3600 # Clear once every hour + start_time = time.time() -while True: - current_time = time.time() - if current_time - start_time > clear_vram_interval: - clear_vram() - start_time = current_time - time.sleep(60) # Check the time every minute + while True: + current_time = time.time() + if current_time - start_time > clear_vram_interval: + clear_vram() + start_time = current_time + time.sleep(60) # Check the time every minute + +if __name__ == "__main__": + main() diff --git a/reproduce/Step_1.py b/reproduce/Step_1.py index e318c145..6df00b8a 100644 --- a/reproduce/Step_1.py +++ b/reproduce/Step_1.py @@ -1,8 +1,10 @@ import os import json import time +import asyncio from lightrag import LightRAG +from lightrag.kg.shared_storage import initialize_pipeline_status def insert_text(rag, file_path): @@ -29,6 +31,19 @@ WORKING_DIR = f"../{cls}" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) -rag = LightRAG(working_dir=WORKING_DIR) +async def initialize_rag(): + rag = LightRAG(working_dir=WORKING_DIR) -insert_text(rag, f"../datasets/unique_contexts/{cls}_unique_contexts.json") + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag + +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) + insert_text(rag, f"../datasets/unique_contexts/{cls}_unique_contexts.json") + + +if __name__ == "__main__": + main() diff --git a/reproduce/Step_1_openai_compatible.py b/reproduce/Step_1_openai_compatible.py index 09ca78c5..7e4139b8 100644 --- a/reproduce/Step_1_openai_compatible.py +++ b/reproduce/Step_1_openai_compatible.py @@ -1,11 +1,13 @@ import os import json import time +import asyncio import numpy as np from lightrag import LightRAG from lightrag.utils import EmbeddingFunc from lightrag.llm.openai import openai_complete_if_cache, openai_embed +from lightrag.kg.shared_storage import initialize_pipeline_status ## For Upstage API @@ -60,12 +62,25 @@ WORKING_DIR = f"../{cls}" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) -rag = LightRAG( - working_dir=WORKING_DIR, - llm_model_func=llm_model_func, - embedding_func=EmbeddingFunc( - embedding_dim=4096, max_token_size=8192, func=embedding_func - ), -) +async def initialize_rag(): + rag = LightRAG( + working_dir=WORKING_DIR, + llm_model_func=llm_model_func, + embedding_func=EmbeddingFunc( + embedding_dim=4096, max_token_size=8192, func=embedding_func + ), + ) -insert_text(rag, f"../datasets/unique_contexts/{cls}_unique_contexts.json") + await rag.initialize_storages() + await initialize_pipeline_status() + + return rag + +def main(): + # Initialize RAG instance + rag = asyncio.run(initialize_rag()) + insert_text(rag, f"../datasets/unique_contexts/{cls}_unique_contexts.json") + + +if __name__ == "__main__": + main() diff --git a/run_with_gunicorn.py b/run_with_gunicorn.py deleted file mode 100755 index 2e4e3cf7..00000000 --- a/run_with_gunicorn.py +++ /dev/null @@ -1,203 +0,0 @@ -#!/usr/bin/env python -""" -Start LightRAG server with Gunicorn -""" - -import os -import sys -import signal -import pipmaster as pm -from lightrag.api.utils_api import parse_args, display_splash_screen -from lightrag.kg.shared_storage import initialize_share_data, finalize_share_data - - -def check_and_install_dependencies(): - """Check and install required dependencies""" - required_packages = [ - "gunicorn", - "tiktoken", - "psutil", - # Add other required packages here - ] - - for package in required_packages: - if not pm.is_installed(package): - print(f"Installing {package}...") - pm.install(package) - print(f"{package} installed successfully") - - -# Signal handler for graceful shutdown -def signal_handler(sig, frame): - print("\n\n" + "=" * 80) - print("RECEIVED TERMINATION SIGNAL") - print(f"Process ID: {os.getpid()}") - print("=" * 80 + "\n") - - # Release shared resources - finalize_share_data() - - # Exit with success status - sys.exit(0) - - -def main(): - # Check and install dependencies - check_and_install_dependencies() - - # Register signal handlers for graceful shutdown - signal.signal(signal.SIGINT, signal_handler) # Ctrl+C - signal.signal(signal.SIGTERM, signal_handler) # kill command - - # Parse all arguments using parse_args - args = parse_args(is_uvicorn_mode=False) - - # Display startup information - display_splash_screen(args) - - print("🚀 Starting LightRAG with Gunicorn") - print(f"🔄 Worker management: Gunicorn (workers={args.workers})") - print("🔍 Preloading app: Enabled") - print("📝 Note: Using Gunicorn's preload feature for shared data initialization") - print("\n\n" + "=" * 80) - print("MAIN PROCESS INITIALIZATION") - print(f"Process ID: {os.getpid()}") - print(f"Workers setting: {args.workers}") - print("=" * 80 + "\n") - - # Import Gunicorn's StandaloneApplication - from gunicorn.app.base import BaseApplication - - # Define a custom application class that loads our config - class GunicornApp(BaseApplication): - def __init__(self, app, options=None): - self.options = options or {} - self.application = app - super().__init__() - - def load_config(self): - # Define valid Gunicorn configuration options - valid_options = { - "bind", - "workers", - "worker_class", - "timeout", - "keepalive", - "preload_app", - "errorlog", - "accesslog", - "loglevel", - "certfile", - "keyfile", - "limit_request_line", - "limit_request_fields", - "limit_request_field_size", - "graceful_timeout", - "max_requests", - "max_requests_jitter", - } - - # Special hooks that need to be set separately - special_hooks = { - "on_starting", - "on_reload", - "on_exit", - "pre_fork", - "post_fork", - "pre_exec", - "pre_request", - "post_request", - "worker_init", - "worker_exit", - "nworkers_changed", - "child_exit", - } - - # Import and configure the gunicorn_config module - import gunicorn_config - - # Set configuration variables in gunicorn_config, prioritizing command line arguments - gunicorn_config.workers = ( - args.workers if args.workers else int(os.getenv("WORKERS", 1)) - ) - - # Bind configuration prioritizes command line arguments - host = args.host if args.host != "0.0.0.0" else os.getenv("HOST", "0.0.0.0") - port = args.port if args.port != 9621 else int(os.getenv("PORT", 9621)) - gunicorn_config.bind = f"{host}:{port}" - - # Log level configuration prioritizes command line arguments - gunicorn_config.loglevel = ( - args.log_level.lower() - if args.log_level - else os.getenv("LOG_LEVEL", "info") - ) - - # Timeout configuration prioritizes command line arguments - gunicorn_config.timeout = ( - args.timeout if args.timeout else int(os.getenv("TIMEOUT", 150)) - ) - - # Keepalive configuration - gunicorn_config.keepalive = int(os.getenv("KEEPALIVE", 5)) - - # SSL configuration prioritizes command line arguments - if args.ssl or os.getenv("SSL", "").lower() in ( - "true", - "1", - "yes", - "t", - "on", - ): - gunicorn_config.certfile = ( - args.ssl_certfile - if args.ssl_certfile - else os.getenv("SSL_CERTFILE") - ) - gunicorn_config.keyfile = ( - args.ssl_keyfile if args.ssl_keyfile else os.getenv("SSL_KEYFILE") - ) - - # Set configuration options from the module - for key in dir(gunicorn_config): - if key in valid_options: - value = getattr(gunicorn_config, key) - # Skip functions like on_starting and None values - if not callable(value) and value is not None: - self.cfg.set(key, value) - # Set special hooks - elif key in special_hooks: - value = getattr(gunicorn_config, key) - if callable(value): - self.cfg.set(key, value) - - if hasattr(gunicorn_config, "logconfig_dict"): - self.cfg.set( - "logconfig_dict", getattr(gunicorn_config, "logconfig_dict") - ) - - def load(self): - # Import the application - from lightrag.api.lightrag_server import get_application - - return get_application(args) - - # Create the application - app = GunicornApp("") - - # Force workers to be an integer and greater than 1 for multi-process mode - workers_count = int(args.workers) - if workers_count > 1: - # Set a flag to indicate we're in the main process - os.environ["LIGHTRAG_MAIN_PROCESS"] = "1" - initialize_share_data(workers_count) - else: - initialize_share_data(1) - - # Run the application - print("\nStarting Gunicorn with direct Python API...") - app.run() - - -if __name__ == "__main__": - main() From ef2a5ad1910e26cb8d6f5a156561f24816f1d299 Mon Sep 17 00:00:00 2001 From: zrguo Date: Mon, 3 Mar 2025 18:40:03 +0800 Subject: [PATCH 12/16] fix linting --- examples/lightrag_api_ollama_demo.py | 7 +- .../lightrag_api_openai_compatible_demo.py | 5 +- examples/lightrag_bedrock_demo.py | 8 +- examples/lightrag_gemini_demo.py | 5 +- examples/lightrag_hf_demo.py | 19 +- examples/lightrag_llamaindex_direct_demo.py | 21 +- examples/lightrag_llamaindex_litellm_demo.py | 21 +- examples/lightrag_lmdeploy_demo.py | 23 +- examples/lightrag_nvidia_demo.py | 5 +- examples/lightrag_ollama_age_demo.py | 27 +- examples/lightrag_ollama_demo.py | 27 +- examples/lightrag_ollama_gremlin_demo.py | 27 +- ...lightrag_ollama_neo4j_milvus_mongo_demo.py | 26 +- examples/lightrag_openai_compatible_demo.py | 3 + ..._openai_compatible_demo_embedding_cache.py | 2 + .../lightrag_openai_compatible_stream_demo.py | 6 +- examples/lightrag_openai_demo.py | 20 +- .../lightrag_openai_mongodb_graph_demo.py | 19 +- ...lightrag_openai_neo4j_milvus_redis_demo.py | 19 +- examples/lightrag_oracle_demo.py | 2 + examples/lightrag_siliconcloud_demo.py | 19 +- examples/lightrag_tidb_demo.py | 2 + examples/lightrag_zhipu_demo.py | 19 +- examples/lightrag_zhipu_postgres_demo.py | 4 +- examples/query_keyword_separation_example.py | 4 +- examples/test.py | 21 +- examples/test_chromadb.py | 20 +- examples/test_faiss.py | 3 +- examples/test_neo4j.py | 19 +- examples/test_split_by_character.ipynb | 1313 ----------------- examples/vram_management_demo.py | 11 +- reproduce/Step_1.py | 2 + reproduce/Step_1_openai_compatible.py | 2 + 33 files changed, 320 insertions(+), 1411 deletions(-) delete mode 100644 examples/test_split_by_character.ipynb diff --git a/examples/lightrag_api_ollama_demo.py b/examples/lightrag_api_ollama_demo.py index f1b68795..dad2a2e0 100644 --- a/examples/lightrag_api_ollama_demo.py +++ b/examples/lightrag_api_ollama_demo.py @@ -36,7 +36,10 @@ async def init(): llm_model_name="gemma2:9b", llm_model_max_async=4, llm_model_max_token_size=8192, - llm_model_kwargs={"host": "http://localhost:11434", "options": {"num_ctx": 8192}}, + llm_model_kwargs={ + "host": "http://localhost:11434", + "options": {"num_ctx": 8192}, + }, embedding_func=EmbeddingFunc( embedding_dim=768, max_token_size=8192, @@ -64,6 +67,8 @@ async def lifespan(app: FastAPI): app = FastAPI( title="LightRAG API", description="API for RAG operations", lifespan=lifespan ) + + # Data models class QueryRequest(BaseModel): query: str diff --git a/examples/lightrag_api_openai_compatible_demo.py b/examples/lightrag_api_openai_compatible_demo.py index 2206f40d..312be872 100644 --- a/examples/lightrag_api_openai_compatible_demo.py +++ b/examples/lightrag_api_openai_compatible_demo.py @@ -75,7 +75,7 @@ async def get_embedding_dim(): # Initialize RAG instance async def init(): embedding_dimension = await get_embedding_dim() - + rag = LightRAG( working_dir=WORKING_DIR, llm_model_func=llm_model_func, @@ -88,9 +88,10 @@ async def init(): await rag.initialize_storages() await initialize_pipeline_status() - + return rag + @asynccontextmanager async def lifespan(app: FastAPI): global rag diff --git a/examples/lightrag_bedrock_demo.py b/examples/lightrag_bedrock_demo.py index 700b4391..68e9f962 100644 --- a/examples/lightrag_bedrock_demo.py +++ b/examples/lightrag_bedrock_demo.py @@ -21,6 +21,7 @@ WORKING_DIR = "./dickens" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -33,9 +34,10 @@ async def initialize_rag(): await rag.initialize_storages() await initialize_pipeline_status() - + return rag + def main(): rag = asyncio.run(initialize_rag()) @@ -47,5 +49,7 @@ def main(): print(f"| {mode.capitalize()} |") print("+-" + "-" * len(mode) + "-+\n") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode=mode)) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode=mode) + ) ) diff --git a/examples/lightrag_gemini_demo.py b/examples/lightrag_gemini_demo.py index d20c0d0d..cd2bb579 100644 --- a/examples/lightrag_gemini_demo.py +++ b/examples/lightrag_gemini_demo.py @@ -12,6 +12,7 @@ from lightrag.kg.shared_storage import initialize_pipeline_status import asyncio import nest_asyncio + # Apply nest_asyncio to solve event loop issues nest_asyncio.apply() @@ -79,9 +80,10 @@ async def initialize_rag(): await rag.initialize_storages() await initialize_pipeline_status() - + return rag + def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) @@ -98,5 +100,6 @@ def main(): print(response) + if __name__ == "__main__": main() diff --git a/examples/lightrag_hf_demo.py b/examples/lightrag_hf_demo.py index 5cd214fd..f2abbb2f 100644 --- a/examples/lightrag_hf_demo.py +++ b/examples/lightrag_hf_demo.py @@ -16,6 +16,7 @@ WORKING_DIR = "./dickens" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -41,6 +42,7 @@ async def initialize_rag(): return rag + def main(): rag = asyncio.run(initialize_rag()) @@ -49,23 +51,32 @@ def main(): # Perform naive search print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + 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")) + 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")) + 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")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) + if __name__ == "__main__": main() diff --git a/examples/lightrag_llamaindex_direct_demo.py b/examples/lightrag_llamaindex_direct_demo.py index d1e68233..d5e3f617 100644 --- a/examples/lightrag_llamaindex_direct_demo.py +++ b/examples/lightrag_llamaindex_direct_demo.py @@ -83,7 +83,7 @@ async def get_embedding_dim(): async def initialize_rag(): embedding_dimension = await get_embedding_dim() - + rag = LightRAG( working_dir=WORKING_DIR, llm_model_func=llm_model_func, @@ -96,7 +96,7 @@ async def initialize_rag(): await rag.initialize_storages() await initialize_pipeline_status() - + return rag @@ -111,23 +111,32 @@ def main(): # Test different query modes print("\nNaive Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="naive") + ) ) print("\nLocal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="local") + ) ) print("\nGlobal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="global") + ) ) print("\nHybrid Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) + if __name__ == "__main__": main() diff --git a/examples/lightrag_llamaindex_litellm_demo.py b/examples/lightrag_llamaindex_litellm_demo.py index d7c609b4..6e738628 100644 --- a/examples/lightrag_llamaindex_litellm_demo.py +++ b/examples/lightrag_llamaindex_litellm_demo.py @@ -86,7 +86,7 @@ async def get_embedding_dim(): async def initialize_rag(): embedding_dimension = await get_embedding_dim() - + rag = LightRAG( working_dir=WORKING_DIR, llm_model_func=llm_model_func, @@ -99,7 +99,7 @@ async def initialize_rag(): await rag.initialize_storages() await initialize_pipeline_status() - + return rag @@ -114,23 +114,32 @@ def main(): # Test different query modes print("\nNaive Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="naive") + ) ) print("\nLocal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="local") + ) ) print("\nGlobal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="global") + ) ) print("\nHybrid Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) + if __name__ == "__main__": main() diff --git a/examples/lightrag_lmdeploy_demo.py b/examples/lightrag_lmdeploy_demo.py index e640a613..ba118fc9 100644 --- a/examples/lightrag_lmdeploy_demo.py +++ b/examples/lightrag_lmdeploy_demo.py @@ -41,6 +41,7 @@ async def lmdeploy_model_complete( **kwargs, ) + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -63,9 +64,10 @@ async def initialize_rag(): await rag.initialize_storages() await initialize_pipeline_status() - + return rag + def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) @@ -77,23 +79,32 @@ def main(): # Test different query modes print("\nNaive Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="naive") + ) ) print("\nLocal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="local") + ) ) print("\nGlobal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="global") + ) ) print("\nHybrid Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/lightrag_nvidia_demo.py b/examples/lightrag_nvidia_demo.py index 9137c1b6..6de0814c 100644 --- a/examples/lightrag_nvidia_demo.py +++ b/examples/lightrag_nvidia_demo.py @@ -97,6 +97,7 @@ async def test_funcs(): # asyncio.run(test_funcs()) + async def initialize_rag(): embedding_dimension = await get_embedding_dim() print(f"Detected embedding dimension: {embedding_dimension}") @@ -117,8 +118,10 @@ async def initialize_rag(): await rag.initialize_storages() await initialize_pipeline_status() - + return rag + + async def main(): try: # Initialize RAG instance diff --git a/examples/lightrag_ollama_age_demo.py b/examples/lightrag_ollama_age_demo.py index 22e42190..b1b4607a 100644 --- a/examples/lightrag_ollama_age_demo.py +++ b/examples/lightrag_ollama_age_demo.py @@ -27,6 +27,7 @@ os.environ["AGE_POSTGRES_HOST"] = "localhost" os.environ["AGE_POSTGRES_PORT"] = "5455" os.environ["AGE_GRAPH_NAME"] = "dickens" + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -34,7 +35,10 @@ async def initialize_rag(): 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}}, + llm_model_kwargs={ + "host": "http://localhost:11434", + "options": {"num_ctx": 32768}, + }, embedding_func=EmbeddingFunc( embedding_dim=768, max_token_size=8192, @@ -47,13 +51,15 @@ async def initialize_rag(): await rag.initialize_storages() await initialize_pipeline_status() - + return rag + async def print_stream(stream): async for chunk in stream: print(chunk, end="", flush=True) + def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) @@ -65,22 +71,30 @@ def main(): # Test different query modes print("\nNaive Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="naive") + ) ) print("\nLocal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="local") + ) ) print("\nGlobal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="global") + ) ) print("\nHybrid Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) # stream response @@ -94,5 +108,6 @@ def main(): else: print(resp) + if __name__ == "__main__": main() diff --git a/examples/lightrag_ollama_demo.py b/examples/lightrag_ollama_demo.py index 6715ea72..cf43aa4a 100644 --- a/examples/lightrag_ollama_demo.py +++ b/examples/lightrag_ollama_demo.py @@ -17,6 +17,7 @@ logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO) if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -24,7 +25,10 @@ async def initialize_rag(): llm_model_name="gemma2:2b", llm_model_max_async=4, llm_model_max_token_size=32768, - llm_model_kwargs={"host": "http://localhost:11434", "options": {"num_ctx": 32768}}, + llm_model_kwargs={ + "host": "http://localhost:11434", + "options": {"num_ctx": 32768}, + }, embedding_func=EmbeddingFunc( embedding_dim=768, max_token_size=8192, @@ -36,13 +40,15 @@ async def initialize_rag(): await rag.initialize_storages() await initialize_pipeline_status() - + return rag + async def print_stream(stream): async for chunk in stream: print(chunk, end="", flush=True) + def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) @@ -54,22 +60,30 @@ def main(): # Test different query modes print("\nNaive Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="naive") + ) ) print("\nLocal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="local") + ) ) print("\nGlobal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="global") + ) ) print("\nHybrid Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) # stream response @@ -83,5 +97,6 @@ def main(): else: print(resp) + if __name__ == "__main__": main() diff --git a/examples/lightrag_ollama_gremlin_demo.py b/examples/lightrag_ollama_gremlin_demo.py index 4d657afa..893b5606 100644 --- a/examples/lightrag_ollama_gremlin_demo.py +++ b/examples/lightrag_ollama_gremlin_demo.py @@ -32,6 +32,7 @@ os.environ["GREMLIN_TRAVERSE_SOURCE"] = "g" os.environ["GREMLIN_USER"] = "" os.environ["GREMLIN_PASSWORD"] = "" + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -39,7 +40,10 @@ async def initialize_rag(): 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}}, + llm_model_kwargs={ + "host": "http://localhost:11434", + "options": {"num_ctx": 32768}, + }, embedding_func=EmbeddingFunc( embedding_dim=768, max_token_size=8192, @@ -52,13 +56,15 @@ async def initialize_rag(): await rag.initialize_storages() await initialize_pipeline_status() - + return rag + async def print_stream(stream): async for chunk in stream: print(chunk, end="", flush=True) + def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) @@ -70,22 +76,30 @@ def main(): # Test different query modes print("\nNaive Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="naive") + ) ) print("\nLocal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="local") + ) ) print("\nGlobal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="global") + ) ) print("\nHybrid Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) # stream response @@ -99,5 +113,6 @@ def main(): else: print(resp) + if __name__ == "__main__": main() diff --git a/examples/lightrag_ollama_neo4j_milvus_mongo_demo.py b/examples/lightrag_ollama_neo4j_milvus_mongo_demo.py index e5d4064d..b6cc931c 100644 --- a/examples/lightrag_ollama_neo4j_milvus_mongo_demo.py +++ b/examples/lightrag_ollama_neo4j_milvus_mongo_demo.py @@ -32,6 +32,7 @@ os.environ["MILVUS_USER"] = "root" os.environ["MILVUS_PASSWORD"] = "root" os.environ["MILVUS_DB_NAME"] = "lightrag" + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -39,7 +40,10 @@ async def initialize_rag(): llm_model_name="qwen2.5:14b", llm_model_max_async=4, llm_model_max_token_size=32768, - llm_model_kwargs={"host": "http://127.0.0.1:11434", "options": {"num_ctx": 32768}}, + llm_model_kwargs={ + "host": "http://127.0.0.1:11434", + "options": {"num_ctx": 32768}, + }, embedding_func=EmbeddingFunc( embedding_dim=1024, max_token_size=8192, @@ -54,9 +58,10 @@ async def initialize_rag(): await rag.initialize_storages() await initialize_pipeline_status() - + return rag + def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) @@ -68,23 +73,32 @@ def main(): # Test different query modes print("\nNaive Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="naive") + ) ) print("\nLocal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="local")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="local") + ) ) print("\nGlobal Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="global")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="global") + ) ) print("\nHybrid Search:") print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) + if __name__ == "__main__": main() diff --git a/examples/lightrag_openai_compatible_demo.py b/examples/lightrag_openai_compatible_demo.py index f4af7be6..1c4a7a92 100644 --- a/examples/lightrag_openai_compatible_demo.py +++ b/examples/lightrag_openai_compatible_demo.py @@ -53,6 +53,7 @@ async def test_funcs(): # asyncio.run(test_funcs()) + async def initialize_rag(): embedding_dimension = await get_embedding_dim() print(f"Detected embedding dimension: {embedding_dimension}") @@ -71,6 +72,8 @@ async def initialize_rag(): await initialize_pipeline_status() return rag + + async def main(): try: # Initialize RAG instance diff --git a/examples/lightrag_openai_compatible_demo_embedding_cache.py b/examples/lightrag_openai_compatible_demo_embedding_cache.py index fcdd1ef3..85408f3b 100644 --- a/examples/lightrag_openai_compatible_demo_embedding_cache.py +++ b/examples/lightrag_openai_compatible_demo_embedding_cache.py @@ -53,6 +53,7 @@ async def test_funcs(): # asyncio.run(test_funcs()) + async def initialize_rag(): embedding_dimension = await get_embedding_dim() print(f"Detected embedding dimension: {embedding_dimension}") @@ -76,6 +77,7 @@ async def initialize_rag(): return rag + async def main(): try: # Initialize RAG instance diff --git a/examples/lightrag_openai_compatible_stream_demo.py b/examples/lightrag_openai_compatible_stream_demo.py index b3f237e5..ab3e73a5 100644 --- a/examples/lightrag_openai_compatible_stream_demo.py +++ b/examples/lightrag_openai_compatible_stream_demo.py @@ -15,6 +15,8 @@ if not os.path.exists(WORKING_DIR): print(f"WorkingDir: {WORKING_DIR}") api_key = "empty" + + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -40,11 +42,13 @@ async def initialize_rag(): return rag + async def print_stream(stream): async for chunk in stream: if chunk: print(chunk, end="", flush=True) + def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) @@ -63,6 +67,6 @@ def main(): else: print(resp) + if __name__ == "__main__": main() - diff --git a/examples/lightrag_openai_demo.py b/examples/lightrag_openai_demo.py index f5f47ee2..138b31a2 100644 --- a/examples/lightrag_openai_demo.py +++ b/examples/lightrag_openai_demo.py @@ -9,6 +9,7 @@ WORKING_DIR = "./dickens" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -22,6 +23,7 @@ async def initialize_rag(): return rag + def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) @@ -31,24 +33,32 @@ def main(): # Perform naive search print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + 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")) + 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")) + 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")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) + if __name__ == "__main__": main() - diff --git a/examples/lightrag_openai_mongodb_graph_demo.py b/examples/lightrag_openai_mongodb_graph_demo.py index ddf5ca63..67c51892 100644 --- a/examples/lightrag_openai_mongodb_graph_demo.py +++ b/examples/lightrag_openai_mongodb_graph_demo.py @@ -76,23 +76,32 @@ def main(): # Perform naive search print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + 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")) + 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")) + 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")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/lightrag_openai_neo4j_milvus_redis_demo.py b/examples/lightrag_openai_neo4j_milvus_redis_demo.py index adf87691..88a61246 100644 --- a/examples/lightrag_openai_neo4j_milvus_redis_demo.py +++ b/examples/lightrag_openai_neo4j_milvus_redis_demo.py @@ -50,6 +50,8 @@ embedding_func = EmbeddingFunc( texts, embed_model="shaw/dmeta-embedding-zh", host="http://117.50.173.35:11434" ), ) + + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -79,23 +81,32 @@ def main(): # Perform naive search print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + 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")) + 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")) + 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")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) + if __name__ == "__main__": main() diff --git a/examples/lightrag_oracle_demo.py b/examples/lightrag_oracle_demo.py index 53139220..420f1af0 100644 --- a/examples/lightrag_oracle_demo.py +++ b/examples/lightrag_oracle_demo.py @@ -64,6 +64,7 @@ async def get_embedding_dim(): embedding_dim = embedding.shape[1] return embedding_dim + async def initialize_rag(): # Detect embedding dimension embedding_dimension = await get_embedding_dim() @@ -102,6 +103,7 @@ async def initialize_rag(): return rag + async def main(): try: # Initialize RAG instance diff --git a/examples/lightrag_siliconcloud_demo.py b/examples/lightrag_siliconcloud_demo.py index 1deb5a66..7a414aca 100644 --- a/examples/lightrag_siliconcloud_demo.py +++ b/examples/lightrag_siliconcloud_demo.py @@ -47,6 +47,7 @@ async def test_funcs(): asyncio.run(test_funcs()) + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -71,24 +72,32 @@ def main(): # Perform naive search print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + 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")) + 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")) + 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")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) + if __name__ == "__main__": main() - diff --git a/examples/lightrag_tidb_demo.py b/examples/lightrag_tidb_demo.py index c3f8fd19..f167e9cc 100644 --- a/examples/lightrag_tidb_demo.py +++ b/examples/lightrag_tidb_demo.py @@ -55,6 +55,7 @@ async def get_embedding_dim(): embedding_dim = embedding.shape[1] return embedding_dim + async def initialize_rag(): # Detect embedding dimension embedding_dimension = await get_embedding_dim() @@ -82,6 +83,7 @@ async def initialize_rag(): return rag + async def main(): try: # Initialize RAG instance diff --git a/examples/lightrag_zhipu_demo.py b/examples/lightrag_zhipu_demo.py index 30b7316f..fdc37c9c 100644 --- a/examples/lightrag_zhipu_demo.py +++ b/examples/lightrag_zhipu_demo.py @@ -19,6 +19,7 @@ api_key = os.environ.get("ZHIPUAI_API_KEY") if api_key is None: raise Exception("Please set ZHIPU_API_KEY in your environment") + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -38,6 +39,7 @@ async def initialize_rag(): return rag + def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) @@ -47,23 +49,32 @@ def main(): # Perform naive search print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + 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")) + 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")) + 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")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) + if __name__ == "__main__": main() diff --git a/examples/lightrag_zhipu_postgres_demo.py b/examples/lightrag_zhipu_postgres_demo.py index 310786a5..304c5f2c 100644 --- a/examples/lightrag_zhipu_postgres_demo.py +++ b/examples/lightrag_zhipu_postgres_demo.py @@ -28,6 +28,7 @@ os.environ["POSTGRES_USER"] = "rag" os.environ["POSTGRES_PASSWORD"] = "rag" os.environ["POSTGRES_DATABASE"] = "rag" + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -55,8 +56,9 @@ async def initialize_rag(): return rag + async def main(): - # Initialize RAG instance + # Initialize RAG instance rag = asyncio.run(initialize_rag()) # add embedding_func for graph database, it's deleted in commit 5661d76860436f7bf5aef2e50d9ee4a59660146c diff --git a/examples/query_keyword_separation_example.py b/examples/query_keyword_separation_example.py index de106de6..cbfdd930 100644 --- a/examples/query_keyword_separation_example.py +++ b/examples/query_keyword_separation_example.py @@ -80,6 +80,8 @@ async def test_funcs(): asyncio.run(test_funcs()) embedding_dimension = 3072 + + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -101,7 +103,7 @@ async def initialize_rag(): async def run_example(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) - + book1 = open("./book_1.txt", encoding="utf-8") book2 = open("./book_2.txt", encoding="utf-8") diff --git a/examples/test.py b/examples/test.py index dc186bda..f2456436 100644 --- a/examples/test.py +++ b/examples/test.py @@ -1,4 +1,5 @@ import os +import asyncio from lightrag import LightRAG, QueryParam from lightrag.llm.openai import gpt_4o_mini_complete from lightrag.kg.shared_storage import initialize_pipeline_status @@ -13,6 +14,7 @@ WORKING_DIR = "./dickens" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -35,23 +37,32 @@ def main(): # Perform naive search print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + 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")) + 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")) + 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")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/test_chromadb.py b/examples/test_chromadb.py index 10d69cc1..e4e9b698 100644 --- a/examples/test_chromadb.py +++ b/examples/test_chromadb.py @@ -112,12 +112,13 @@ async def initialize_rag(): }, ) - await rag.initialize_storages() await initialize_pipeline_status() return rag + +def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) @@ -126,23 +127,32 @@ async def initialize_rag(): # Perform naive search print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + 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")) + 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")) + 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")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) + if __name__ == "__main__": main() diff --git a/examples/test_faiss.py b/examples/test_faiss.py index 34991f47..febdce14 100644 --- a/examples/test_faiss.py +++ b/examples/test_faiss.py @@ -58,6 +58,7 @@ async def embedding_func(texts: list[str]) -> np.ndarray: embeddings = model.encode(texts, convert_to_numpy=True) return embeddings + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -78,8 +79,8 @@ async def initialize_rag(): return rag + def main(): - # Initialize RAG instance rag = asyncio.run(initialize_rag()) # Insert the custom chunks into LightRAG diff --git a/examples/test_neo4j.py b/examples/test_neo4j.py index 2d1d527a..7f620acc 100644 --- a/examples/test_neo4j.py +++ b/examples/test_neo4j.py @@ -15,6 +15,7 @@ WORKING_DIR = "./local_neo4jWorkDir" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -29,6 +30,7 @@ async def initialize_rag(): return rag + def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) @@ -38,23 +40,32 @@ def main(): # Perform naive search print( - rag.query("What are the top themes in this story?", param=QueryParam(mode="naive")) + 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")) + 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")) + 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")) + rag.query( + "What are the top themes in this story?", param=QueryParam(mode="hybrid") + ) ) + if __name__ == "__main__": main() diff --git a/examples/test_split_by_character.ipynb b/examples/test_split_by_character.ipynb deleted file mode 100644 index f70f9f34..00000000 --- a/examples/test_split_by_character.ipynb +++ /dev/null @@ -1,1313 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 1, - "id": "4b5690db12e34685", - "metadata": { - "ExecuteTime": { - "end_time": "2025-01-09T03:40:58.307102Z", - "start_time": "2025-01-09T03:40:51.935233Z" - } - }, - "outputs": [], - "source": [ - "import os\n", - "import logging\n", - "import numpy as np\n", - "from lightrag import LightRAG, QueryParam\n", - "from lightrag.llm.openai import openai_complete_if_cache, openai_embed\n", - "from lightrag.utils import EmbeddingFunc\n", - "from lightrag.kg.shared_storage import initialize_pipeline_status\n", - "import nest_asyncio" - ] - }, - { - "cell_type": "markdown", - "id": "dd17956ec322b361", - "metadata": {}, - "source": [ - "#### split by character" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "8c8ee7c061bf9159", - "metadata": { - "ExecuteTime": { - "end_time": "2025-01-09T03:41:13.961167Z", - "start_time": "2025-01-09T03:41:13.958357Z" - } - }, - "outputs": [], - "source": [ - "nest_asyncio.apply()\n", - "WORKING_DIR = \"../../llm_rag/paper_db/R000088_test1\"\n", - "logging.basicConfig(format=\"%(levelname)s:%(message)s\", level=logging.INFO)\n", - "if not os.path.exists(WORKING_DIR):\n", - " os.mkdir(WORKING_DIR)\n", - "API = os.environ.get(\"DOUBAO_API_KEY\")" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "a5009d16e0851dca", - "metadata": { - "ExecuteTime": { - "end_time": "2025-01-09T03:41:16.862036Z", - "start_time": "2025-01-09T03:41:16.859306Z" - } - }, - "outputs": [], - "source": [ - "async def llm_model_func(\n", - " prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs\n", - ") -> str:\n", - " return await openai_complete_if_cache(\n", - " \"ep-20241218114828-2tlww\",\n", - " prompt,\n", - " system_prompt=system_prompt,\n", - " history_messages=history_messages,\n", - " api_key=API,\n", - " base_url=\"https://ark.cn-beijing.volces.com/api/v3\",\n", - " **kwargs,\n", - " )\n", - "\n", - "\n", - "async def embedding_func(texts: list[str]) -> np.ndarray:\n", - " return await openai_embed(\n", - " texts,\n", - " model=\"ep-20241231173413-pgjmk\",\n", - " api_key=API,\n", - " base_url=\"https://ark.cn-beijing.volces.com/api/v3\",\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "397fcad24ce4d0ed", - "metadata": { - "ExecuteTime": { - "end_time": "2025-01-09T03:41:24.950307Z", - "start_time": "2025-01-09T03:41:24.940353Z" - } - }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "INFO:lightrag:Logger initialized for working directory: ../../llm_rag/paper_db/R000088_test1\n", - "INFO:lightrag:Load KV llm_response_cache with 0 data\n", - "INFO:lightrag:Load KV full_docs with 0 data\n", - "INFO:lightrag:Load KV text_chunks with 0 data\n", - "INFO:nano-vectordb:Init {'embedding_dim': 4096, 'metric': 'cosine', 'storage_file': '../../llm_rag/paper_db/R000088_test1/vdb_entities.json'} 0 data\n", - "INFO:nano-vectordb:Init {'embedding_dim': 4096, 'metric': 'cosine', 'storage_file': '../../llm_rag/paper_db/R000088_test1/vdb_relationships.json'} 0 data\n", - "INFO:nano-vectordb:Init {'embedding_dim': 4096, 'metric': 'cosine', 'storage_file': '../../llm_rag/paper_db/R000088_test1/vdb_chunks.json'} 0 data\n", - "INFO:lightrag:Loaded document status storage with 0 records\n" - ] - } - ], - "source": [ - "import asyncio\n", - "import nest_asyncio\n", - "\n", - "nest_asyncio.apply()\n", - "\n", - "async def initialize_rag():\n", - " rag = LightRAG(\n", - " working_dir=WORKING_DIR,\n", - " llm_model_func=llm_model_func,\n", - " embedding_func=EmbeddingFunc(\n", - " embedding_dim=4096, max_token_size=8192, func=embedding_func\n", - " ),\n", - " chunk_token_size=512,\n", - " )\n", - " await rag.initialize_storages()\n", - " await initialize_pipeline_status()\n", - "\n", - " return rag\n", - "\n", - "rag = asyncio.run(initialize_rag())" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "1dc3603677f7484d", - "metadata": { - "ExecuteTime": { - "end_time": "2025-01-09T03:41:37.947456Z", - "start_time": "2025-01-09T03:41:37.941901Z" - } - }, - "outputs": [], - "source": [ - "with open(\n", - " \"../../llm_rag/example/R000088/auto/R000088_full_txt.md\", \"r\", encoding=\"utf-8\"\n", - ") as f:\n", - " content = f.read()\n", - "\n", - "\n", - "async def embedding_func(texts: list[str]) -> np.ndarray:\n", - " return await openai_embed(\n", - " texts,\n", - " model=\"ep-20241231173413-pgjmk\",\n", - " api_key=API,\n", - " base_url=\"https://ark.cn-beijing.volces.com/api/v3\",\n", - " )\n", - "\n", - "\n", - "async def get_embedding_dim():\n", - " test_text = [\"This is a test sentence.\"]\n", - " embedding = await embedding_func(test_text)\n", - " embedding_dim = embedding.shape[1]\n", - " return embedding_dim" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "6844202606acfbe5", - "metadata": { - "ExecuteTime": { - "end_time": "2025-01-09T03:41:39.608541Z", - "start_time": "2025-01-09T03:41:39.165057Z" - } - }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n" - ] - } - ], - "source": [ - "embedding_dimension = await get_embedding_dim()" - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "id": "d6273839d9681403", - "metadata": { - "ExecuteTime": { - "end_time": "2025-01-09T03:44:34.295345Z", - "start_time": "2025-01-09T03:41:48.324171Z" - } - }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "INFO:lightrag:Processing 1 new unique documents\n", - "Processing batch 1: 0%| | 0/1 [00:00标签中，针对每个问题详细分析你的思考过程。然后在<回答>标签中给出所有问题的最终答案。\"\"\"" - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "id": "7a6491385b050095", - "metadata": { - "ExecuteTime": { - "end_time": "2025-01-09T03:45:40.829111Z", - "start_time": "2025-01-09T03:45:13.530298Z" - } - }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n", - "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n", - "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n", - "INFO:lightrag:Local query uses 5 entites, 12 relations, 3 text units\n", - "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/embeddings \"HTTP/1.1 200 OK\"\n", - "INFO:lightrag:Global query uses 8 entites, 5 relations, 4 text units\n", - "INFO:httpx:HTTP Request: POST https://ark.cn-beijing.volces.com/api/v3/chat/completions \"HTTP/1.1 200 OK\"\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "<分析>\n", - "1. **该文献主要研究的问题是什么？**\n", - " - 思考过程：通过浏览论文内容，查找作者明确阐述研究目的的部分。文中多处提及“Our study was performed to explore whether folic acid treatment was associated with cancer outcomes and all-cause mortality after extended follow-up”，表明作者旨在探究叶酸治疗与癌症结局及全因死亡率之间的关系，尤其是在经过长期随访后。\n", - "2. **该文献采用什么方法进行分析？**\n", - " - 思考过程：寻找描述研究方法和数据分析过程的段落。文中提到“Survival curves were constructed using the Kaplan-Meier method and differences in survival between groups were analyzed using the log-rank test. Estimates of hazard ratios (HRs) with 95% CIs were obtained by using Cox proportional hazards regression models stratified by trial”，可以看出作者使用了Kaplan-Meier法构建生存曲线、log-rank检验分析组间生存差异以及Cox比例风险回归模型估计风险比等方法。\n", - "3. **该文献的主要结论是什么？**\n", - " - 思考过程：定位到论文中总结结论的部分，如“Conclusion Treatment with folic acid plus vitamin $\\mathsf{B}_{12}$ was associated with increased cancer outcomes and all-cause mortality in patients with ischemic heart disease in Norway, where there is no folic acid fortification of foods”，可知作者得出叶酸加维生素$\\mathsf{B}_{12}$治疗与癌症结局和全因死亡率增加有关的结论。\n", - "<回答>\n", - "1. 该文献主要研究的问题是：叶酸治疗与癌症结局及全因死亡率之间的关系，尤其是在经过长期随访后，叶酸治疗是否与癌症结局和全因死亡率相关。\n", - "2. 该文献采用的分析方法包括：使用Kaplan-Meier法构建生存曲线、log-rank检验分析组间生存差异、Cox比例风险回归模型估计风险比等。\n", - "3. 该文献的主要结论是：在挪威没有叶酸强化食品的情况下，叶酸加维生素$\\mathsf{B}_{12}$治疗与缺血性心脏病患者的癌症结局和全因死亡率增加有关。\n", - "\n", - "**参考文献**\n", - "- [VD] In2Norwegianhomocysteine-lowering trialsamongpatientswithischemicheart disease, there was a statistically nonsignificantincreaseincancerincidenceinthe groupsassignedtofolicacidtreatment.15,16 Our study was performed to explore whetherfolicacidtreatmentwasassociatedwithcanceroutcomesandall-cause mortality after extended follow-up.\n", - "- [VD] Survivalcurveswereconstructedusing theKaplan-Meiermethodanddifferences insurvivalbetweengroupswereanalyzed usingthelog-ranktest.Estimatesofhazard ratios (HRs) with $95\\%$ CIs were obtainedbyusingCoxproportionalhazards regressionmodelsstratifiedbytrial.\n", - "- [VD] Conclusion Treatment with folic acid plus vitamin $\\mathsf{B}_{12}$ was associated with increased cancer outcomes and all-cause mortality in patients with ischemic heart disease in Norway, where there is no folic acid fortification of foods.\n" - ] - } - ], - "source": [ - "resp = rag.query(prompt1, param=QueryParam(mode=\"mix\", top_k=5))\n", - "print(resp)" - ] - }, - { - "cell_type": "markdown", - "id": "4e5bfad24cb721a8", - "metadata": {}, - "source": [ - "#### split by character only" - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "id": "44e2992dc95f8ce0", - "metadata": { - "ExecuteTime": { - "end_time": "2025-01-09T03:47:40.988796Z", - "start_time": "2025-01-09T03:47:40.982648Z" - } - }, - "outputs": [], - "source": [ - "WORKING_DIR = \"../../llm_rag/paper_db/R000088_test2\"\n", - "if not os.path.exists(WORKING_DIR):\n", - " os.mkdir(WORKING_DIR)" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "id": "62c63385d2d973d5", - "metadata": { - "ExecuteTime": { - "end_time": "2025-01-09T03:51:39.951329Z", - "start_time": "2025-01-09T03:49:15.218976Z" - } - }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "INFO:lightrag:Logger initialized for working directory: ../../llm_rag/paper_db/R000088_test2\n", - "INFO:lightrag:Load KV llm_response_cache with 0 data\n", - "INFO:lightrag:Load KV full_docs with 0 data\n", - "INFO:lightrag:Load KV text_chunks with 0 data\n", - "INFO:nano-vectordb:Init {'embedding_dim': 4096, 'metric': 'cosine', 'storage_file': '../../llm_rag/paper_db/R000088_test2/vdb_entities.json'} 0 data\n", - "INFO:nano-vectordb:Init {'embedding_dim': 4096, 'metric': 'cosine', 'storage_file': '../../llm_rag/paper_db/R000088_test2/vdb_relationships.json'} 0 data\n", - "INFO:nano-vectordb:Init {'embedding_dim': 4096, 'metric': 'cosine', 'storage_file': '../../llm_rag/paper_db/R000088_test2/vdb_chunks.json'} 0 data\n", - "INFO:lightrag:Loaded document status storage with 0 records\n", - "INFO:lightrag:Processing 1 new unique documents\n", - "Processing batch 1: 0%| | 0/1 [00:00\n", - "- **该文献主要研究的问题是什么？**\n", - " - **思考过程**：通过浏览论文的标题、摘要、引言等部分，寻找关于研究目的和问题的描述。论文标题为“Cancer Incidence and Mortality After Treatment With Folic Acid and Vitamin B12”，摘要中的“Objective”部分明确指出研究目的是“To evaluate effects of treatment with B vitamins on cancer outcomes and all-cause mortality in 2 randomized controlled trials”。因此，可以确定该文献主要研究的问题是评估B族维生素治疗对两项随机对照试验中癌症结局和全因死亡率的影响。\n", - "- **该文献采用什么方法进行分析？**\n", - " - **思考过程**：在论文的“METHODS”部分详细描述了研究方法。文中提到这是一个对两项随机、双盲、安慰剂对照临床试验（Norwegian Vitamin [NORVIT] trial和Western Norway B Vitamin Intervention Trial [WENBIT]）数据的联合分析，并进行了观察性的试验后随访。具体包括对参与者进行分组干预（不同剂量的叶酸、维生素B12、维生素B6或安慰剂），收集临床信息和血样，分析循环B族维生素、同型半胱氨酸和可替宁等指标，并进行基因分型等，还涉及到多种统计分析方法，如计算预期癌症发生率、构建生存曲线、进行Cox比例风险回归模型分析等。\n", - "- **该文献的主要结论是什么？**\n", - " - **思考过程**：在论文的“Results”和“Conclusion”部分寻找主要结论。研究结果表明，在治疗期间，接受叶酸加维生素B12治疗的参与者血清叶酸浓度显著增加，且在后续随访中，该组癌症发病率、癌症死亡率和全因死亡率均有所上升，主要是肺癌发病率增加，而维生素B6治疗未显示出显著影响。结论部分明确指出“Treatment with folic acid plus vitamin $\\mathsf{B}_{12}$ was associated with increased cancer outcomes and all-cause mortality in patients with ischemic heart disease in Norway, where there is no folic acid fortification of foods”。\n", - "\n", - "\n", - "<回答>\n", - "- **主要研究问题**：评估B族维生素治疗对两项随机对照试验中癌症结局和全因死亡率的影响。\n", - "- **研究方法**：采用对两项随机、双盲、安慰剂对照临床试验（Norwegian Vitamin [NORVIT] trial和Western Norway B Vitamin Intervention Trial [WENBIT]）数据的联合分析，并进行观察性的试验后随访，涉及分组干预、多种指标检测以及多种统计分析方法。\n", - "- **主要结论**：在挪威（食品中未添加叶酸），对于缺血性心脏病患者，叶酸加维生素B12治疗与癌症结局和全因死亡率的增加有关，而维生素B6治疗未显示出显著影响。\n", - "\n", - "**参考文献**\n", - "- [VD] Cancer Incidence and Mortality After Treatment With Folic Acid and Vitamin B12\n", - "- [VD] METHODS Study Design, Participants, and Study Intervention\n", - "- [VD] RESULTS\n", - "- [VD] Conclusion\n", - "- [VD] Objective To evaluate effects of treatment with B vitamins on cancer outcomes and all-cause mortality in 2 randomized controlled trials.\n" - ] - } - ], - "source": [ - "resp = rag.query(prompt1, param=QueryParam(mode=\"mix\", top_k=5))\n", - "print(resp)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7ba6fa79a2550d10", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 2 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython2", - "version": "2.7.6" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/examples/vram_management_demo.py b/examples/vram_management_demo.py index f4d46ab4..36eb5468 100644 --- a/examples/vram_management_demo.py +++ b/examples/vram_management_demo.py @@ -14,6 +14,7 @@ TEXT_FILES_DIR = "/llm/mt" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) + async def initialize_rag(): # Initialize LightRAG rag = LightRAG( @@ -31,6 +32,7 @@ async def initialize_rag(): return rag + # Read all .txt files from the TEXT_FILES_DIR directory texts = [] for filename in os.listdir(TEXT_FILES_DIR): @@ -82,7 +84,8 @@ def main(): try: print( rag.query( - "What are the top themes in this story?", param=QueryParam(mode="global") + "What are the top themes in this story?", + param=QueryParam(mode="global"), ) ) except Exception as e: @@ -91,18 +94,17 @@ def main(): try: print( rag.query( - "What are the top themes in this story?", param=QueryParam(mode="hybrid") + "What are the top themes in this story?", + param=QueryParam(mode="hybrid"), ) ) except Exception as e: print(f"Error performing hybrid search: {e}") - # Function to clear VRAM resources def clear_vram(): os.system("sudo nvidia-smi --gpu-reset") - # Regularly clear VRAM to prevent overflow clear_vram_interval = 3600 # Clear once every hour start_time = time.time() @@ -114,5 +116,6 @@ def main(): start_time = current_time time.sleep(60) # Check the time every minute + if __name__ == "__main__": main() diff --git a/reproduce/Step_1.py b/reproduce/Step_1.py index 6df00b8a..c94015ad 100644 --- a/reproduce/Step_1.py +++ b/reproduce/Step_1.py @@ -31,6 +31,7 @@ WORKING_DIR = f"../{cls}" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) + async def initialize_rag(): rag = LightRAG(working_dir=WORKING_DIR) @@ -39,6 +40,7 @@ async def initialize_rag(): return rag + def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) diff --git a/reproduce/Step_1_openai_compatible.py b/reproduce/Step_1_openai_compatible.py index 7e4139b8..3b9944eb 100644 --- a/reproduce/Step_1_openai_compatible.py +++ b/reproduce/Step_1_openai_compatible.py @@ -62,6 +62,7 @@ WORKING_DIR = f"../{cls}" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) + async def initialize_rag(): rag = LightRAG( working_dir=WORKING_DIR, @@ -76,6 +77,7 @@ async def initialize_rag(): return rag + def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) From a27cf5a09e1af9240e723ecfb7e986e852079f84 Mon Sep 17 00:00:00 2001 From: zrguo Date: Mon, 3 Mar 2025 19:17:34 +0800 Subject: [PATCH 13/16] fix delete_by_doc_id --- lightrag/lightrag.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index 4f1ad7dc..5470d2e2 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -395,6 +395,7 @@ class LightRAG: namespace=make_namespace( self.namespace_prefix, NameSpace.KV_STORE_LLM_RESPONSE_CACHE ), + global_config=asdict(self), embedding_func=self.embedding_func, ) @@ -1421,17 +1422,19 @@ class LightRAG: # 3. Before deleting, check the related entities and relationships for these chunks for chunk_id in chunk_ids: # Check entities + entities_storage = await self.entities_vdb.client_storage entities = [ dp - for dp in self.entities_vdb.client_storage["data"] + for dp in entities_storage["data"] if chunk_id in dp.get("source_id") ] logger.debug(f"Chunk {chunk_id} has {len(entities)} related entities") # Check relationships + relationships_storage = await self.relationships_vdb.client_storage relations = [ dp - for dp in self.relationships_vdb.client_storage["data"] + for dp in relationships_storage["data"] if chunk_id in dp.get("source_id") ] logger.debug(f"Chunk {chunk_id} has {len(relations)} related relations") @@ -1495,7 +1498,9 @@ class LightRAG: for entity in entities_to_delete: await self.entities_vdb.delete_entity(entity) logger.debug(f"Deleted entity {entity} from vector DB") - self.chunk_entity_relation_graph.remove_nodes(list(entities_to_delete)) + await self.chunk_entity_relation_graph.remove_nodes( + list(entities_to_delete) + ) logger.debug(f"Deleted {len(entities_to_delete)} entities from graph") # Update entities @@ -1514,7 +1519,7 @@ class LightRAG: rel_id_1 = compute_mdhash_id(tgt + src, prefix="rel-") await self.relationships_vdb.delete([rel_id_0, rel_id_1]) logger.debug(f"Deleted relationship {src}-{tgt} from vector DB") - self.chunk_entity_relation_graph.remove_edges( + await self.chunk_entity_relation_graph.remove_edges( list(relationships_to_delete) ) logger.debug( @@ -1545,9 +1550,10 @@ class LightRAG: async def process_data(data_type, vdb, chunk_id): # Check data (entities or relationships) + storage = await vdb.client_storage data_with_chunk = [ dp - for dp in vdb.client_storage["data"] + for dp in storage["data"] if chunk_id in (dp.get("source_id") or "").split(GRAPH_FIELD_SEP) ] From a81428515569ce08b832aaae1fdf00bdc3796fcd Mon Sep 17 00:00:00 2001 From: zrguo Date: Mon, 3 Mar 2025 21:09:45 +0800 Subject: [PATCH 14/16] Add a feature that allows modifying nodes and relationships. --- README.md | 76 +++++++ lightrag/lightrag.py | 492 +++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 555 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 2b0d9f4b..a3f2152f 100644 --- a/README.md +++ b/README.md @@ -750,6 +750,82 @@ rag.delete_by_entity("Project Gutenberg") rag.delete_by_doc_id("doc_id") ``` +## Edit Entities and Relations + +LightRAG now supports comprehensive knowledge graph management capabilities, allowing you to create, edit, and delete entities and relationships within your knowledge graph. + +### Create Entities and Relations + +```python +import asyncio +from lightrag import LightRAG + +# Initialize LightRAG +rag = LightRAG( + working_dir="your_working_dir", + embedding_func=your_embedding_function, + llm_model_func=your_llm_function +) + +# Create new entity +entity = rag.create_entity("Google", { + "description": "Google is a multinational technology company specializing in internet-related services and products.", + "entity_type": "company" +}) + +# Create another entity +product = rag.create_entity("Gmail", { + "description": "Gmail is an email service developed by Google.", + "entity_type": "product" +}) + +# Create relation between entities +relation = rag.create_relation("Google", "Gmail", { + "description": "Google develops and operates Gmail.", + "keywords": "develops operates service", + "weight": 2.0 +}) +``` + +### Edit Entities and Relations + +```python +# Edit an existing entity +updated_entity = rag.edit_entity("Google", { + "description": "Google is a subsidiary of Alphabet Inc., founded in 1998.", + "entity_type": "tech_company" +}) + +# Rename an entity (with all its relationships properly migrated) +renamed_entity = rag.edit_entity("Gmail", { + "entity_name": "Google Mail", + "description": "Google Mail (formerly Gmail) is an email service." +}) + +# Edit a relation between entities +updated_relation = rag.edit_relation("Google", "Google Mail", { + "description": "Google created and maintains Google Mail service.", + "keywords": "creates maintains email service", + "weight": 3.0 +}) +``` + +All operations are available in both synchronous and asynchronous versions. The asynchronous versions have the prefix "a" (e.g., `acreate_entity`, `aedit_relation`). + +#### Entity Operations + +- **create_entity**: Creates a new entity with specified attributes +- **edit_entity**: Updates an existing entity's attributes or renames it +- **delete_entity**: Removes an entity and all its relationships + +#### Relation Operations + +- **create_relation**: Creates a new relation between existing entities +- **edit_relation**: Updates an existing relation's attributes +- **delete_relation**: Removes a relation between entities + +These operations maintain data consistency across both the graph database and vector database components, ensuring your knowledge graph remains coherent. + ## Cache
diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index 5470d2e2..8bcb6b51 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -967,7 +967,7 @@ class LightRAG: # Insert chunks into vector storage all_chunks_data: dict[str, dict[str, str]] = {} chunk_to_source_map: dict[str, str] = {} - for chunk_data in custom_kg.get("chunks", {}): + for chunk_data in custom_kg.get("chunks", []): chunk_content = self.clean_text(chunk_data["content"]) source_id = chunk_data["source_id"] tokens = len( @@ -997,9 +997,10 @@ class LightRAG: update_storage = True if all_chunks_data: - await self.chunks_vdb.upsert(all_chunks_data) - if all_chunks_data: - await self.text_chunks.upsert(all_chunks_data) + await asyncio.gather( + self.chunks_vdb.upsert(all_chunks_data), + self.text_chunks.upsert(all_chunks_data), + ) # Insert entities into knowledge graph all_entities_data: list[dict[str, str]] = [] @@ -1007,7 +1008,6 @@ class LightRAG: entity_name = entity_data["entity_name"] entity_type = entity_data.get("entity_type", "UNKNOWN") description = entity_data.get("description", "No description provided") - # source_id = entity_data["source_id"] source_chunk_id = entity_data.get("source_id", "UNKNOWN") source_id = chunk_to_source_map.get(source_chunk_id, "UNKNOWN") @@ -1039,7 +1039,6 @@ class LightRAG: description = relationship_data["description"] keywords = relationship_data["keywords"] weight = relationship_data.get("weight", 1.0) - # source_id = relationship_data["source_id"] source_chunk_id = relationship_data.get("source_id", "UNKNOWN") source_id = chunk_to_source_map.get(source_chunk_id, "UNKNOWN") @@ -1079,34 +1078,43 @@ class LightRAG: "tgt_id": tgt_id, "description": description, "keywords": keywords, + "source_id": source_id, + "weight": weight, } all_relationships_data.append(edge_data) update_storage = True - # Insert entities into vector storage if needed + # Insert entities into vector storage with consistent format data_for_vdb = { compute_mdhash_id(dp["entity_name"], prefix="ent-"): { - "content": dp["entity_name"] + dp["description"], + "content": dp["entity_name"] + "\n" + dp["description"], "entity_name": dp["entity_name"], + "source_id": dp["source_id"], + "description": dp["description"], + "entity_type": dp["entity_type"], } for dp in all_entities_data } await self.entities_vdb.upsert(data_for_vdb) - # Insert relationships into vector storage if needed + # Insert relationships into vector storage with consistent format data_for_vdb = { compute_mdhash_id(dp["src_id"] + dp["tgt_id"], prefix="rel-"): { "src_id": dp["src_id"], "tgt_id": dp["tgt_id"], - "content": dp["keywords"] - + dp["src_id"] - + dp["tgt_id"] - + dp["description"], + "source_id": dp["source_id"], + "content": f"{dp['keywords']}\t{dp['src_id']}\n{dp['tgt_id']}\n{dp['description']}", + "keywords": dp["keywords"], + "description": dp["description"], + "weight": dp["weight"], } for dp in all_relationships_data } await self.relationships_vdb.upsert(data_for_vdb) + except Exception as e: + logger.error(f"Error in ainsert_custom_kg: {e}") + raise finally: if update_storage: await self._insert_done() @@ -1759,3 +1767,461 @@ class LightRAG: def clear_cache(self, modes: list[str] | None = None) -> None: """Synchronous version of aclear_cache.""" return always_get_an_event_loop().run_until_complete(self.aclear_cache(modes)) + + async def aedit_entity( + self, entity_name: str, updated_data: dict[str, str], allow_rename: bool = True + ) -> dict[str, Any]: + """Asynchronously edit entity information. + + Updates entity information in the knowledge graph and re-embeds the entity in the vector database. + + Args: + entity_name: Name of the entity to edit + updated_data: Dictionary containing updated attributes, e.g. {"description": "new description", "entity_type": "new type"} + allow_rename: Whether to allow entity renaming, defaults to True + + Returns: + Dictionary containing updated entity information + """ + try: + # 1. Get current entity information + node_data = await self.chunk_entity_relation_graph.get_node(entity_name) + if not node_data: + raise ValueError(f"Entity '{entity_name}' does not exist") + + # Check if entity is being renamed + new_entity_name = updated_data.get("entity_name", entity_name) + is_renaming = new_entity_name != entity_name + + # If renaming, check if new name already exists + if is_renaming: + if not allow_rename: + raise ValueError( + "Entity renaming is not allowed. Set allow_rename=True to enable this feature" + ) + + existing_node = await self.chunk_entity_relation_graph.get_node( + new_entity_name + ) + if existing_node: + raise ValueError( + f"Entity name '{new_entity_name}' already exists, cannot rename" + ) + + # 2. Update entity information in the graph + new_node_data = {**node_data, **updated_data} + if "entity_name" in new_node_data: + del new_node_data[ + "entity_name" + ] # Node data should not contain entity_name field + + # If renaming entity + if is_renaming: + logger.info(f"Renaming entity '{entity_name}' to '{new_entity_name}'") + + # Create new entity + await self.chunk_entity_relation_graph.upsert_node( + new_entity_name, new_node_data + ) + + # Get all edges related to the original entity + edges = await self.chunk_entity_relation_graph.get_node_edges( + entity_name + ) + if edges: + # Recreate edges for the new entity + for source, target in edges: + edge_data = await self.chunk_entity_relation_graph.get_edge( + source, target + ) + if edge_data: + if source == entity_name: + await self.chunk_entity_relation_graph.upsert_edge( + new_entity_name, target, edge_data + ) + else: # target == entity_name + await self.chunk_entity_relation_graph.upsert_edge( + source, new_entity_name, edge_data + ) + + # Delete old entity + await self.chunk_entity_relation_graph.delete_node(entity_name) + + # Delete old entity record from vector database + old_entity_id = compute_mdhash_id(entity_name, prefix="ent-") + await self.entities_vdb.delete([old_entity_id]) + + # Update working entity name to new name + entity_name = new_entity_name + else: + # If not renaming, directly update node data + await self.chunk_entity_relation_graph.upsert_node( + entity_name, new_node_data + ) + + # 3. Recalculate entity's vector representation and update vector database + description = new_node_data.get("description", "") + source_id = new_node_data.get("source_id", "") + entity_type = new_node_data.get("entity_type", "") + content = entity_name + "\n" + description + + # Calculate entity ID + entity_id = compute_mdhash_id(entity_name, prefix="ent-") + + # Prepare data for vector database update + entity_data = { + entity_id: { + "content": content, + "entity_name": entity_name, + "source_id": source_id, + "description": description, + "entity_type": entity_type, + } + } + + # Update vector database + await self.entities_vdb.upsert(entity_data) + + # 4. Save changes + await self._edit_entity_done() + + logger.info(f"Entity '{entity_name}' successfully updated") + return await self.get_entity_info(entity_name, include_vector_data=True) + except Exception as e: + logger.error(f"Error while editing entity '{entity_name}': {e}") + raise + + def edit_entity( + self, entity_name: str, updated_data: dict[str, str], allow_rename: bool = True + ) -> dict[str, Any]: + """Synchronously edit entity information. + + Updates entity information in the knowledge graph and re-embeds the entity in the vector database. + + Args: + entity_name: Name of the entity to edit + updated_data: Dictionary containing updated attributes, e.g. {"description": "new description", "entity_type": "new type"} + allow_rename: Whether to allow entity renaming, defaults to True + + Returns: + Dictionary containing updated entity information + """ + loop = always_get_an_event_loop() + return loop.run_until_complete( + self.aedit_entity(entity_name, updated_data, allow_rename) + ) + + async def _edit_entity_done(self) -> None: + """Callback after entity editing is complete, ensures updates are persisted""" + await asyncio.gather( + *[ + cast(StorageNameSpace, storage_inst).index_done_callback() + for storage_inst in [ # type: ignore + self.entities_vdb, + self.chunk_entity_relation_graph, + ] + ] + ) + + async def aedit_relation( + self, source_entity: str, target_entity: str, updated_data: dict[str, Any] + ) -> dict[str, Any]: + """Asynchronously edit relation information. + + Updates relation (edge) information in the knowledge graph and re-embeds the relation in the vector database. + + Args: + source_entity: Name of the source entity + target_entity: Name of the target entity + updated_data: Dictionary containing updated attributes, e.g. {"description": "new description", "keywords": "new keywords"} + + Returns: + Dictionary containing updated relation information + """ + try: + # 1. Get current relation information + edge_data = await self.chunk_entity_relation_graph.get_edge( + source_entity, target_entity + ) + if not edge_data: + raise ValueError( + f"Relation from '{source_entity}' to '{target_entity}' does not exist" + ) + + # 2. Update relation information in the graph + new_edge_data = {**edge_data, **updated_data} + await self.chunk_entity_relation_graph.upsert_edge( + source_entity, target_entity, new_edge_data + ) + + # 3. Recalculate relation's vector representation and update vector database + description = new_edge_data.get("description", "") + keywords = new_edge_data.get("keywords", "") + source_id = new_edge_data.get("source_id", "") + weight = float(new_edge_data.get("weight", 1.0)) + + # Create content for embedding + content = f"{keywords}\t{source_entity}\n{target_entity}\n{description}" + + # Calculate relation ID + relation_id = compute_mdhash_id( + source_entity + target_entity, prefix="rel-" + ) + + # Prepare data for vector database update + relation_data = { + relation_id: { + "content": content, + "src_id": source_entity, + "tgt_id": target_entity, + "source_id": source_id, + "description": description, + "keywords": keywords, + "weight": weight, + } + } + + # Update vector database + await self.relationships_vdb.upsert(relation_data) + + # 4. Save changes + await self._edit_relation_done() + + logger.info( + f"Relation from '{source_entity}' to '{target_entity}' successfully updated" + ) + return await self.get_relation_info( + source_entity, target_entity, include_vector_data=True + ) + except Exception as e: + logger.error( + f"Error while editing relation from '{source_entity}' to '{target_entity}': {e}" + ) + raise + + def edit_relation( + self, source_entity: str, target_entity: str, updated_data: dict[str, Any] + ) -> dict[str, Any]: + """Synchronously edit relation information. + + Updates relation (edge) information in the knowledge graph and re-embeds the relation in the vector database. + + Args: + source_entity: Name of the source entity + target_entity: Name of the target entity + updated_data: Dictionary containing updated attributes, e.g. {"description": "new description", "keywords": "keywords"} + + Returns: + Dictionary containing updated relation information + """ + loop = always_get_an_event_loop() + return loop.run_until_complete( + self.aedit_relation(source_entity, target_entity, updated_data) + ) + + async def _edit_relation_done(self) -> None: + """Callback after relation editing is complete, ensures updates are persisted""" + await asyncio.gather( + *[ + cast(StorageNameSpace, storage_inst).index_done_callback() + for storage_inst in [ # type: ignore + self.relationships_vdb, + self.chunk_entity_relation_graph, + ] + ] + ) + + async def acreate_entity( + self, entity_name: str, entity_data: dict[str, Any] + ) -> dict[str, Any]: + """Asynchronously create a new entity. + + Creates a new entity in the knowledge graph and adds it to the vector database. + + Args: + entity_name: Name of the new entity + entity_data: Dictionary containing entity attributes, e.g. {"description": "description", "entity_type": "type"} + + Returns: + Dictionary containing created entity information + """ + try: + # Check if entity already exists + existing_node = await self.chunk_entity_relation_graph.get_node(entity_name) + if existing_node: + raise ValueError(f"Entity '{entity_name}' already exists") + + # Prepare node data with defaults if missing + node_data = { + "entity_type": entity_data.get("entity_type", "UNKNOWN"), + "description": entity_data.get("description", ""), + "source_id": entity_data.get("source_id", "manual"), + } + + # Add entity to knowledge graph + await self.chunk_entity_relation_graph.upsert_node(entity_name, node_data) + + # Prepare content for entity + description = node_data.get("description", "") + source_id = node_data.get("source_id", "") + entity_type = node_data.get("entity_type", "") + content = entity_name + "\n" + description + + # Calculate entity ID + entity_id = compute_mdhash_id(entity_name, prefix="ent-") + + # Prepare data for vector database update + entity_data_for_vdb = { + entity_id: { + "content": content, + "entity_name": entity_name, + "source_id": source_id, + "description": description, + "entity_type": entity_type, + } + } + + # Update vector database + await self.entities_vdb.upsert(entity_data_for_vdb) + + # Save changes + await self._edit_entity_done() + + logger.info(f"Entity '{entity_name}' successfully created") + return await self.get_entity_info(entity_name, include_vector_data=True) + except Exception as e: + logger.error(f"Error while creating entity '{entity_name}': {e}") + raise + + def create_entity( + self, entity_name: str, entity_data: dict[str, Any] + ) -> dict[str, Any]: + """Synchronously create a new entity. + + Creates a new entity in the knowledge graph and adds it to the vector database. + + Args: + entity_name: Name of the new entity + entity_data: Dictionary containing entity attributes, e.g. {"description": "description", "entity_type": "type"} + + Returns: + Dictionary containing created entity information + """ + loop = always_get_an_event_loop() + return loop.run_until_complete(self.acreate_entity(entity_name, entity_data)) + + async def acreate_relation( + self, source_entity: str, target_entity: str, relation_data: dict[str, Any] + ) -> dict[str, Any]: + """Asynchronously create a new relation between entities. + + Creates a new relation (edge) in the knowledge graph and adds it to the vector database. + + Args: + source_entity: Name of the source entity + target_entity: Name of the target entity + relation_data: Dictionary containing relation attributes, e.g. {"description": "description", "keywords": "keywords"} + + Returns: + Dictionary containing created relation information + """ + try: + # Check if both entities exist + source_exists = await self.chunk_entity_relation_graph.has_node( + source_entity + ) + target_exists = await self.chunk_entity_relation_graph.has_node( + target_entity + ) + + if not source_exists: + raise ValueError(f"Source entity '{source_entity}' does not exist") + if not target_exists: + raise ValueError(f"Target entity '{target_entity}' does not exist") + + # Check if relation already exists + existing_edge = await self.chunk_entity_relation_graph.get_edge( + source_entity, target_entity + ) + if existing_edge: + raise ValueError( + f"Relation from '{source_entity}' to '{target_entity}' already exists" + ) + + # Prepare edge data with defaults if missing + edge_data = { + "description": relation_data.get("description", ""), + "keywords": relation_data.get("keywords", ""), + "source_id": relation_data.get("source_id", "manual"), + "weight": float(relation_data.get("weight", 1.0)), + } + + # Add relation to knowledge graph + await self.chunk_entity_relation_graph.upsert_edge( + source_entity, target_entity, edge_data + ) + + # Prepare content for embedding + description = edge_data.get("description", "") + keywords = edge_data.get("keywords", "") + source_id = edge_data.get("source_id", "") + weight = edge_data.get("weight", 1.0) + + # Create content for embedding + content = f"{keywords}\t{source_entity}\n{target_entity}\n{description}" + + # Calculate relation ID + relation_id = compute_mdhash_id( + source_entity + target_entity, prefix="rel-" + ) + + # Prepare data for vector database update + relation_data_for_vdb = { + relation_id: { + "content": content, + "src_id": source_entity, + "tgt_id": target_entity, + "source_id": source_id, + "description": description, + "keywords": keywords, + "weight": weight, + } + } + + # Update vector database + await self.relationships_vdb.upsert(relation_data_for_vdb) + + # Save changes + await self._edit_relation_done() + + logger.info( + f"Relation from '{source_entity}' to '{target_entity}' successfully created" + ) + return await self.get_relation_info( + source_entity, target_entity, include_vector_data=True + ) + except Exception as e: + logger.error( + f"Error while creating relation from '{source_entity}' to '{target_entity}': {e}" + ) + raise + + def create_relation( + self, source_entity: str, target_entity: str, relation_data: dict[str, Any] + ) -> dict[str, Any]: + """Synchronously create a new relation between entities. + + Creates a new relation (edge) in the knowledge graph and adds it to the vector database. + + Args: + source_entity: Name of the source entity + target_entity: Name of the target entity + relation_data: Dictionary containing relation attributes, e.g. {"description": "description", "keywords": "keywords"} + + Returns: + Dictionary containing created relation information + """ + loop = always_get_an_event_loop() + return loop.run_until_complete( + self.acreate_relation(source_entity, target_entity, relation_data) + ) From 78151d045b865d54b8a744fa8989b6f3e630d9fa Mon Sep 17 00:00:00 2001 From: zrguo Date: Mon, 3 Mar 2025 21:13:43 +0800 Subject: [PATCH 15/16] fix README.md --- README.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/README.md b/README.md index a3f2152f..4215a94a 100644 --- a/README.md +++ b/README.md @@ -816,13 +816,11 @@ All operations are available in both synchronous and asynchronous versions. The - **create_entity**: Creates a new entity with specified attributes - **edit_entity**: Updates an existing entity's attributes or renames it -- **delete_entity**: Removes an entity and all its relationships #### Relation Operations - **create_relation**: Creates a new relation between existing entities - **edit_relation**: Updates an existing relation's attributes -- **delete_relation**: Removes a relation between entities These operations maintain data consistency across both the graph database and vector database components, ensuring your knowledge graph remains coherent. From 0ac482f6d3bc6dc68de55851b575372ffbee6c98 Mon Sep 17 00:00:00 2001 From: zrguo Date: Mon, 3 Mar 2025 21:15:34 +0800 Subject: [PATCH 16/16] Update README.md --- README.md | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/README.md b/README.md index 4215a94a..abc2f8b3 100644 --- a/README.md +++ b/README.md @@ -757,16 +757,6 @@ LightRAG now supports comprehensive knowledge graph management capabilities, all ### Create Entities and Relations ```python -import asyncio -from lightrag import LightRAG - -# Initialize LightRAG -rag = LightRAG( - working_dir="your_working_dir", - embedding_func=your_embedding_function, - llm_model_func=your_llm_function -) - # Create new entity entity = rag.create_entity("Google", { "description": "Google is a multinational technology company specializing in internet-related services and products.",