gzdaniel commited on
Commit
ddc7121
·
1 Parent(s): 51f0667

Fix linting

Browse files
Files changed (1) hide show
  1. lightrag/kg/mongo_impl.py +7 -6
lightrag/kg/mongo_impl.py CHANGED
@@ -527,7 +527,9 @@ class MongoGraphStorage(BaseGraphStorage):
527
  {"$group": {"_id": "$source_node_id", "degree": {"$sum": 1}}},
528
  ]
529
 
530
- cursor = await self.edge_collection.aggregate(outbound_pipeline, allowDiskUse=True)
 
 
531
  async for doc in cursor:
532
  merged_results[doc.get("_id")] = doc.get("degree")
533
 
@@ -537,7 +539,9 @@ class MongoGraphStorage(BaseGraphStorage):
537
  {"$group": {"_id": "$target_node_id", "degree": {"$sum": 1}}},
538
  ]
539
 
540
- cursor = await self.edge_collection.aggregate(inbound_pipeline, allowDiskUse=True)
 
 
541
  async for doc in cursor:
542
  merged_results[doc.get("_id")] = merged_results.get(
543
  doc.get("_id"), 0
@@ -698,10 +702,7 @@ class MongoGraphStorage(BaseGraphStorage):
698
  """
699
 
700
  # Use aggregation with allowDiskUse for large datasets
701
- pipeline = [
702
- {"$project": {"_id": 1}},
703
- {"$sort": {"_id": 1}}
704
- ]
705
  cursor = await self.collection.aggregate(pipeline, allowDiskUse=True)
706
  labels = []
707
  async for doc in cursor:
 
527
  {"$group": {"_id": "$source_node_id", "degree": {"$sum": 1}}},
528
  ]
529
 
530
+ cursor = await self.edge_collection.aggregate(
531
+ outbound_pipeline, allowDiskUse=True
532
+ )
533
  async for doc in cursor:
534
  merged_results[doc.get("_id")] = doc.get("degree")
535
 
 
539
  {"$group": {"_id": "$target_node_id", "degree": {"$sum": 1}}},
540
  ]
541
 
542
+ cursor = await self.edge_collection.aggregate(
543
+ inbound_pipeline, allowDiskUse=True
544
+ )
545
  async for doc in cursor:
546
  merged_results[doc.get("_id")] = merged_results.get(
547
  doc.get("_id"), 0
 
702
  """
703
 
704
  # Use aggregation with allowDiskUse for large datasets
705
+ pipeline = [{"$project": {"_id": 1}}, {"$sort": {"_id": 1}}]
 
 
 
706
  cursor = await self.collection.aggregate(pipeline, allowDiskUse=True)
707
  labels = []
708
  async for doc in cursor: