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_dotenv()
async def main():
"""Example usage of LightRAG with FalkorDB"""
@ -33,8 +34,8 @@ async def main():
rag = LightRAG(
working_dir="./falkordb_example",
llm_model_func=gpt_4o_mini_complete, # Updated function name
embedding_func=openai_embed, # Updated function name
graph_storage="FalkorDBStorage", # Specify FalkorDB backend
embedding_func=openai_embed, # Updated function name
graph_storage="FalkorDBStorage", # Specify FalkorDB backend
)
# Initialize storage connections
@ -66,7 +67,7 @@ async def main():
questions = [
"What is FalkorDB and how does it relate to LightRAG?",
"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):
@ -75,31 +76,30 @@ async def main():
try:
response = await rag.aquery(
question,
param=QueryParam(mode="hybrid", top_k=3)
question, param=QueryParam(mode="hybrid", top_k=3)
)
print(f"A: {response}")
except Exception as e:
print(f"Error querying: {e}")
# Show some graph statistics
print(f"\n--- Graph Statistics ---")
print("\n--- Graph Statistics ---")
try:
all_labels = await storage.get_all_labels()
print(f"Unique entities: {len(all_labels)}")
if nodes:
print(f"Sample entities:")
print("Sample entities:")
for i, node in enumerate(nodes[:3]):
entity_id = node.get('entity_id', 'Unknown')
entity_type = node.get('entity_type', 'Unknown')
entity_id = node.get("entity_id", "Unknown")
entity_type = node.get("entity_type", "Unknown")
print(f" {i+1}. {entity_id} ({entity_type})")
if edges:
print(f"Sample relationships:")
print("Sample relationships:")
for i, edge in enumerate(edges[:2]):
source = edge.get('source', 'Unknown')
target = edge.get('target', 'Unknown')
source = edge.get("source", "Unknown")
target = edge.get("target", "Unknown")
print(f" {i+1}. {source} → {target}")
except Exception as e:
@ -110,7 +110,9 @@ if __name__ == "__main__":
print("LightRAG with FalkorDB Example")
print("==============================")
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()
# Check OpenAI API key

31
examples/graph_visual_with_falkordb.py
View file

@ -1,5 +1,4 @@
import os
import json
import xml.etree.ElementTree as ET
import falkordb
@ -130,12 +129,18 @@ def insert_nodes_and_edges_to_falkordb(data):
# Print some statistics
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
edge_count = edge_count_result.result_set[0][0] if edge_count_result.result_set else 0
node_count = (
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"- Edges in database: {edge_count}")
@ -153,7 +158,9 @@ def query_graph_data():
print("\n=== Sample Graph Data ===")
# 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)
print("\nSample Nodes:")
@ -172,7 +179,9 @@ def query_graph_data():
print("\nSample Edges:")
if 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
query = """
@ -212,8 +221,12 @@ def main():
xml_file = os.path.join(WORKING_DIR, "graph_chunk_entity_relation.graphml")
if not os.path.exists(xml_file):
print(f"Error: File {xml_file} not found. Please ensure the GraphML file exists.")
print("This file is typically generated by LightRAG after processing documents.")
print(
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
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"
async def initialize(self):
HOST = os.environ.get("FALKORDB_HOST", config.get("falkordb", "host", fallback="localhost"))
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))
HOST = os.environ.get(
"FALKORDB_HOST", config.get("falkordb", "host", fallback="localhost")
)
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(
"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:
@ -93,9 +107,13 @@ class FalkorDBStorage(BaseGraphStorage):
# Create index for workspace nodes on entity_id if it doesn't exist
workspace_label = self._get_workspace_label()
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)
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:
# Index may already exist, which is not an error
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"""
loop = asyncio.get_event_loop()
return await loop.run_in_executor(
self._executor,
lambda: self._graph.query(query, params or {})
self._executor, lambda: self._graph.query(query, params or {})
)
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}}) "
"RETURN COUNT(r) > 0 AS edgeExists"
)
result = await self._run_query(query, {
"source_entity_id": source_node_id,
"target_entity_id": target_node_id,
})
result = await self._run_query(
query,
{
"source_entity_id": source_node_id,
"target_entity_id": target_node_id,
},
)
return result.result_set[0][0] if result.result_set else False
except Exception as e:
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}})
RETURN properties(r) as edge_properties
"""
result = await self._run_query(query, {
"source_entity_id": source_node_id,
"target_entity_id": target_node_id,
})
result = await self._run_query(
query,
{
"source_entity_id": source_node_id,
"target_entity_id": target_node_id,
},
)
if result.result_set and len(result.result_set) > 0:
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)
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 = []
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)
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 = []
if result.result_set:
@ -624,7 +651,9 @@ class FalkorDBStorage(BaseGraphStorage):
properties = node_data
entity_type = properties["entity_type"]
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:
query = f"""
@ -632,7 +661,9 @@ class FalkorDBStorage(BaseGraphStorage):
SET n += $properties
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:
logger.error(f"Error during upsert: {str(e)}")
raise
@ -669,11 +700,14 @@ class FalkorDBStorage(BaseGraphStorage):
SET r += $properties
RETURN r, source, target
"""
await self._run_query(query, {
"source_entity_id": source_node_id,
"target_entity_id": target_node_id,
"properties": edge_properties,
})
await self._run_query(
query,
{
"source_entity_id": source_node_id,
"target_entity_id": target_node_id,
"properties": edge_properties,
},
)
except Exception as e:
logger.error(f"Error during edge upsert: {str(e)}")
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}})
DELETE r
"""
await self._run_query(query, {
"source_entity_id": source,
"target_entity_id": target
})
await self._run_query(
query, {"source_entity_id": source, "target_entity_id": target}
)
logger.debug(f"Deleted edge from '{source}' to '{target}'")
except Exception as e:
logger.error(f"Error during edge deletion: {str(e)}")
@ -948,7 +981,5 @@ class FalkorDBStorage(BaseGraphStorage):
"message": f"workspace '{workspace_label}' data dropped",
}
except Exception as e:
logger.error(
f"Error dropping FalkorDB workspace '{workspace_label}': {e}"
)
logger.error(f"Error dropping FalkorDB workspace '{workspace_label}': {e}")
return {"status": "error", "message": str(e)}