gmakstutis/LightRAG
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix code formatting with pre-commit hooks

Browse source
This commit is contained in:
Naseem Ali 2025-08-26 13:03:42 +03:00
parent e00b2a1334
commit c5ef544f8b
4 changed files with 156 additions and 110 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
examples/falkordb_example.py
View file

@ -20,6 +20,7 @@ from lightrag.kg.shared_storage import initialize_pipeline_status
# Load environment variables # Load environment variables
load_dotenv() load_dotenv()
async def main(): async def main():
"""Example usage of LightRAG with FalkorDB""" """Example usage of LightRAG with FalkorDB"""
@ -33,8 +34,8 @@ async def main():
rag = LightRAG( rag = LightRAG(
working_dir="./falkordb_example", working_dir="./falkordb_example",
llm_model_func=gpt_4o_mini_complete, # Updated function name llm_model_func=gpt_4o_mini_complete, # Updated function name
embedding_func=openai_embed, # Updated function name embedding_func=openai_embed, # Updated function name
graph_storage="FalkorDBStorage", # Specify FalkorDB backend graph_storage="FalkorDBStorage", # Specify FalkorDB backend
) )
# Initialize storage connections # Initialize storage connections
@ -66,7 +67,7 @@ async def main():
questions = [ questions = [
"What is FalkorDB and how does it relate to LightRAG?", "What is FalkorDB and how does it relate to LightRAG?",
"What are the benefits of using Redis with graph databases?", "What are the benefits of using Redis with graph databases?",
"How does FalkorDB support OpenCypher queries?" "How does FalkorDB support OpenCypher queries?",
] ]
for i, question in enumerate(questions, 1): for i, question in enumerate(questions, 1):
@ -75,31 +76,30 @@ async def main():
try: try:
response = await rag.aquery( response = await rag.aquery(
question, question, param=QueryParam(mode="hybrid", top_k=3)
param=QueryParam(mode="hybrid", top_k=3)
) )
print(f"A: {response}") print(f"A: {response}")
except Exception as e: except Exception as e:
print(f"Error querying: {e}") print(f"Error querying: {e}")
# Show some graph statistics # Show some graph statistics
print(f"\n--- Graph Statistics ---") print("\n--- Graph Statistics ---")
try: try:
all_labels = await storage.get_all_labels() all_labels = await storage.get_all_labels()
print(f"Unique entities: {len(all_labels)}") print(f"Unique entities: {len(all_labels)}")
if nodes: if nodes:
print(f"Sample entities:") print("Sample entities:")
for i, node in enumerate(nodes[:3]): for i, node in enumerate(nodes[:3]):
entity_id = node.get('entity_id', 'Unknown') entity_id = node.get("entity_id", "Unknown")
entity_type = node.get('entity_type', 'Unknown') entity_type = node.get("entity_type", "Unknown")
print(f" {i+1}. {entity_id} ({entity_type})") print(f" {i+1}. {entity_id} ({entity_type})")
if edges: if edges:
print(f"Sample relationships:") print("Sample relationships:")
for i, edge in enumerate(edges[:2]): for i, edge in enumerate(edges[:2]):
source = edge.get('source', 'Unknown') source = edge.get("source", "Unknown")
target = edge.get('target', 'Unknown') target = edge.get("target", "Unknown")
print(f" {i+1}. {source} → {target}") print(f" {i+1}. {source} → {target}")
except Exception as e: except Exception as e:
@ -110,7 +110,9 @@ if __name__ == "__main__":
print("LightRAG with FalkorDB Example") print("LightRAG with FalkorDB Example")
print("==============================") print("==============================")
print("Note: This requires FalkorDB running on localhost:6379") print("Note: This requires FalkorDB running on localhost:6379")
print("You can start FalkorDB with: docker run -p 6379:6379 falkordb/falkordb:latest") print(
"You can start FalkorDB with: docker run -p 6379:6379 falkordb/falkordb:latest"
)
print() print()
# Check OpenAI API key # Check OpenAI API key

31
examples/graph_visual_with_falkordb.py
View file

@ -1,5 +1,4 @@
import os import os
import json
import xml.etree.ElementTree as ET import xml.etree.ElementTree as ET
import falkordb import falkordb
@ -130,12 +129,18 @@ def insert_nodes_and_edges_to_falkordb(data):
# Print some statistics # Print some statistics
node_count_result = graph.query("MATCH (n:Entity) RETURN count(n) AS count") node_count_result = graph.query("MATCH (n:Entity) RETURN count(n) AS count")
edge_count_result = graph.query("MATCH ()-[r:DIRECTED]-() RETURN count(r) AS count") edge_count_result = graph.query(
"MATCH ()-[r:DIRECTED]-() RETURN count(r) AS count"
)
node_count = node_count_result.result_set[0][0] if node_count_result.result_set else 0 node_count = (
edge_count = edge_count_result.result_set[0][0] if edge_count_result.result_set else 0 node_count_result.result_set[0][0] if node_count_result.result_set else 0
)
edge_count = (
edge_count_result.result_set[0][0] if edge_count_result.result_set else 0
)
print(f"Final statistics:") print("Final statistics:")
print(f"- Nodes in database: {node_count}") print(f"- Nodes in database: {node_count}")
print(f"- Edges in database: {edge_count}") print(f"- Edges in database: {edge_count}")
@ -153,7 +158,9 @@ def query_graph_data():
print("\n=== Sample Graph Data ===") print("\n=== Sample Graph Data ===")
# Get some sample nodes # Get some sample nodes
query = "MATCH (n:Entity) RETURN n.entity_id, n.entity_type, n.description LIMIT 5" query = (
"MATCH (n:Entity) RETURN n.entity_id, n.entity_type, n.description LIMIT 5"
)
result = graph.query(query) result = graph.query(query)
print("\nSample Nodes:") print("\nSample Nodes:")
@ -172,7 +179,9 @@ def query_graph_data():
print("\nSample Edges:") print("\nSample Edges:")
if result.result_set: if result.result_set:
for record in result.result_set: for record in result.result_set:
print(f"- {record[0]} -> {record[1]} (weight: {record[2]}): {record[3][:100]}...") print(
f"- {record[0]} -> {record[1]} (weight: {record[2]}): {record[3][:100]}..."
)
# Get node degree statistics # Get node degree statistics
query = """ query = """
@ -212,8 +221,12 @@ def main():
xml_file = os.path.join(WORKING_DIR, "graph_chunk_entity_relation.graphml") xml_file = os.path.join(WORKING_DIR, "graph_chunk_entity_relation.graphml")
if not os.path.exists(xml_file): if not os.path.exists(xml_file):
print(f"Error: File {xml_file} not found. Please ensure the GraphML file exists.") print(
print("This file is typically generated by LightRAG after processing documents.") f"Error: File {xml_file} not found. Please ensure the GraphML file exists."
)
print(
"This file is typically generated by LightRAG after processing documents."
)
return return
print("FalkorDB Graph Visualization Example") print("FalkorDB Graph Visualization Example")

97
lightrag/kg/falkordb_impl.py
View file

@ -66,13 +66,27 @@ class FalkorDBStorage(BaseGraphStorage):
return workspace if workspace else "base" return workspace if workspace else "base"
async def initialize(self): async def initialize(self):
HOST = os.environ.get("FALKORDB_HOST", config.get("falkordb", "host", fallback="localhost")) HOST = os.environ.get(
PORT = int(os.environ.get("FALKORDB_PORT", config.get("falkordb", "port", fallback=6379))) "FALKORDB_HOST", config.get("falkordb", "host", fallback="localhost")
PASSWORD = os.environ.get("FALKORDB_PASSWORD", config.get("falkordb", "password", fallback=None)) )
USERNAME = os.environ.get("FALKORDB_USERNAME", config.get("falkordb", "username", fallback=None)) PORT = int(
os.environ.get(
"FALKORDB_PORT", config.get("falkordb", "port", fallback=6379)
)
)
PASSWORD = os.environ.get(
"FALKORDB_PASSWORD", config.get("falkordb", "password", fallback=None)
)
USERNAME = os.environ.get(
"FALKORDB_USERNAME", config.get("falkordb", "username", fallback=None)
)
GRAPH_NAME = os.environ.get( GRAPH_NAME = os.environ.get(
"FALKORDB_GRAPH_NAME", "FALKORDB_GRAPH_NAME",
config.get("falkordb", "graph_name", fallback=re.sub(r"[^a-zA-Z0-9-]", "-", self.namespace)) config.get(
"falkordb",
"graph_name",
fallback=re.sub(r"[^a-zA-Z0-9-]", "-", self.namespace),
),
) )
try: try:
@ -93,9 +107,13 @@ class FalkorDBStorage(BaseGraphStorage):
# Create index for workspace nodes on entity_id if it doesn't exist # Create index for workspace nodes on entity_id if it doesn't exist
workspace_label = self._get_workspace_label() workspace_label = self._get_workspace_label()
try: try:
index_query = f"CREATE INDEX FOR (n:`{workspace_label}`) ON (n.entity_id)" index_query = (
f"CREATE INDEX FOR (n:`{workspace_label}`) ON (n.entity_id)"
)
await self._run_query(index_query) await self._run_query(index_query)
logger.info(f"Created index for {workspace_label} nodes on entity_id in FalkorDB") logger.info(
f"Created index for {workspace_label} nodes on entity_id in FalkorDB"
)
except Exception as e: except Exception as e:
# Index may already exist, which is not an error # Index may already exist, which is not an error
logger.debug(f"Index creation may have failed or already exists: {e}") logger.debug(f"Index creation may have failed or already exists: {e}")
@ -124,8 +142,7 @@ class FalkorDBStorage(BaseGraphStorage):
"""Run a query asynchronously using thread pool""" """Run a query asynchronously using thread pool"""
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
return await loop.run_in_executor( return await loop.run_in_executor(
self._executor, self._executor, lambda: self._graph.query(query, params or {})
lambda: self._graph.query(query, params or {})
) )
async def index_done_callback(self) -> None: async def index_done_callback(self) -> None:
@ -176,10 +193,13 @@ class FalkorDBStorage(BaseGraphStorage):
f"MATCH (a:`{workspace_label}` {{entity_id: $source_entity_id}})-[r]-(b:`{workspace_label}` {{entity_id: $target_entity_id}}) " f"MATCH (a:`{workspace_label}` {{entity_id: $source_entity_id}})-[r]-(b:`{workspace_label}` {{entity_id: $target_entity_id}}) "
"RETURN COUNT(r) > 0 AS edgeExists" "RETURN COUNT(r) > 0 AS edgeExists"
) )
result = await self._run_query(query, { result = await self._run_query(
"source_entity_id": source_node_id, query,
"target_entity_id": target_node_id, {
}) "source_entity_id": source_node_id,
"target_entity_id": target_node_id,
},
)
return result.result_set[0][0] if result.result_set else False return result.result_set[0][0] if result.result_set else False
except Exception as e: except Exception as e:
logger.error( logger.error(
@ -380,10 +400,13 @@ class FalkorDBStorage(BaseGraphStorage):
MATCH (start:`{workspace_label}` {{entity_id: $source_entity_id}})-[r]-(end:`{workspace_label}` {{entity_id: $target_entity_id}}) MATCH (start:`{workspace_label}` {{entity_id: $source_entity_id}})-[r]-(end:`{workspace_label}` {{entity_id: $target_entity_id}})
RETURN properties(r) as edge_properties RETURN properties(r) as edge_properties
""" """
result = await self._run_query(query, { result = await self._run_query(
"source_entity_id": source_node_id, query,
"target_entity_id": target_node_id, {
}) "source_entity_id": source_node_id,
"target_entity_id": target_node_id,
},
)
if result.result_set and len(result.result_set) > 0: if result.result_set and len(result.result_set) > 0:
edge_result = result.result_set[0][0] # Get properties dict edge_result = result.result_set[0][0] # Get properties dict
@ -574,7 +597,9 @@ class FalkorDBStorage(BaseGraphStorage):
WHERE n.source_id IS NOT NULL AND chunk_id IN split(n.source_id, $sep) WHERE n.source_id IS NOT NULL AND chunk_id IN split(n.source_id, $sep)
RETURN DISTINCT n RETURN DISTINCT n
""" """
result = await self._run_query(query, {"chunk_ids": chunk_ids, "sep": GRAPH_FIELD_SEP}) result = await self._run_query(
query, {"chunk_ids": chunk_ids, "sep": GRAPH_FIELD_SEP}
)
nodes = [] nodes = []
if result.result_set: if result.result_set:
@ -595,7 +620,9 @@ class FalkorDBStorage(BaseGraphStorage):
WHERE r.source_id IS NOT NULL AND chunk_id IN split(r.source_id, $sep) WHERE r.source_id IS NOT NULL AND chunk_id IN split(r.source_id, $sep)
RETURN DISTINCT a.entity_id AS source, b.entity_id AS target, properties(r) AS properties RETURN DISTINCT a.entity_id AS source, b.entity_id AS target, properties(r) AS properties
""" """
result = await self._run_query(query, {"chunk_ids": chunk_ids, "sep": GRAPH_FIELD_SEP}) result = await self._run_query(
query, {"chunk_ids": chunk_ids, "sep": GRAPH_FIELD_SEP}
)
edges = [] edges = []
if result.result_set: if result.result_set:
@ -624,7 +651,9 @@ class FalkorDBStorage(BaseGraphStorage):
properties = node_data properties = node_data
entity_type = properties["entity_type"] entity_type = properties["entity_type"]
if "entity_id" not in properties: if "entity_id" not in properties:
raise ValueError("FalkorDB: node properties must contain an 'entity_id' field") raise ValueError(
"FalkorDB: node properties must contain an 'entity_id' field"
)
try: try:
query = f""" query = f"""
@ -632,7 +661,9 @@ class FalkorDBStorage(BaseGraphStorage):
SET n += $properties SET n += $properties
SET n:`{entity_type}` SET n:`{entity_type}`
""" """
await self._run_query(query, {"entity_id": node_id, "properties": properties}) await self._run_query(
query, {"entity_id": node_id, "properties": properties}
)
except Exception as e: except Exception as e:
logger.error(f"Error during upsert: {str(e)}") logger.error(f"Error during upsert: {str(e)}")
raise raise
@ -669,11 +700,14 @@ class FalkorDBStorage(BaseGraphStorage):
SET r += $properties SET r += $properties
RETURN r, source, target RETURN r, source, target
""" """
await self._run_query(query, { await self._run_query(
"source_entity_id": source_node_id, query,
"target_entity_id": target_node_id, {
"properties": edge_properties, "source_entity_id": source_node_id,
}) "target_entity_id": target_node_id,
"properties": edge_properties,
},
)
except Exception as e: except Exception as e:
logger.error(f"Error during edge upsert: {str(e)}") logger.error(f"Error during edge upsert: {str(e)}")
raise raise
@ -868,10 +902,9 @@ class FalkorDBStorage(BaseGraphStorage):
MATCH (source:`{workspace_label}` {{entity_id: $source_entity_id}})-[r]-(target:`{workspace_label}` {{entity_id: $target_entity_id}}) MATCH (source:`{workspace_label}` {{entity_id: $source_entity_id}})-[r]-(target:`{workspace_label}` {{entity_id: $target_entity_id}})
DELETE r DELETE r
""" """
await self._run_query(query, { await self._run_query(
"source_entity_id": source, query, {"source_entity_id": source, "target_entity_id": target}
"target_entity_id": target )
})
logger.debug(f"Deleted edge from '{source}' to '{target}'") logger.debug(f"Deleted edge from '{source}' to '{target}'")
except Exception as e: except Exception as e:
logger.error(f"Error during edge deletion: {str(e)}") logger.error(f"Error during edge deletion: {str(e)}")
@ -948,7 +981,5 @@ class FalkorDBStorage(BaseGraphStorage):
"message": f"workspace '{workspace_label}' data dropped", "message": f"workspace '{workspace_label}' data dropped",
} }
except Exception as e: except Exception as e:
logger.error( logger.error(f"Error dropping FalkorDB workspace '{workspace_label}': {e}")
f"Error dropping FalkorDB workspace '{workspace_label}': {e}"
)
return {"status": "error", "message": str(e)} return {"status": "error", "message": str(e)}