Merge pull request #768 from ArnoChenFx/fix-backends

Fix various runtime errors

config.ini.example

@@ -6,7 +6,6 @@ password = your-password
 [mongodb]
 uri = mongodb+srv://name:password@your-cluster-address
 database = lightrag
-graph = false
 
 [redis]
 uri=redis://localhost:6379/1

lightrag/kg/milvus_impl.py

@@ -32,8 +32,8 @@ class MilvusVectorDBStorage(BaseVectorStorage):
         )
 
     def __post_init__(self):
-        config = self.global_config.get("vector_db_storage_cls_kwargs", {})
-        cosine_threshold = config.get("cosine_better_than_threshold")
+        kwargs = self.global_config.get("vector_db_storage_cls_kwargs", {})
+        cosine_threshold = kwargs.get("cosine_better_than_threshold")
         if cosine_threshold is None:
             raise ValueError(
                 "cosine_better_than_threshold must be specified in vector_db_storage_cls_kwargs"

lightrag/kg/mongo_impl.py

@@ -44,7 +44,7 @@ class MongoKVStorage(BaseKVStorage):
         database = client.get_database(
             os.environ.get(
                 "MONGO_DATABASE",
-                mongo_database=config.get("mongodb", "database", fallback="LightRAG"),
+                config.get("mongodb", "database", fallback="LightRAG"),
             )
         )
         self._data = database.get_collection(self.namespace)

lightrag/kg/qdrant_impl.py

@@ -15,7 +15,6 @@ if not pm.is_installed("qdrant_client"):
 
 from qdrant_client import QdrantClient, models
 
-
 config = configparser.ConfigParser()
 config.read("config.ini", "utf-8")
 
@@ -61,8 +60,8 @@ class QdrantVectorDBStorage(BaseVectorStorage):
         client.create_collection(collection_name, **kwargs)
 
     def __post_init__(self):
-        config = self.global_config.get("vector_db_storage_cls_kwargs", {})
-        cosine_threshold = config.get("cosine_better_than_threshold")
+        kwargs = self.global_config.get("vector_db_storage_cls_kwargs", {})
+        cosine_threshold = kwargs.get("cosine_better_than_threshold")
         if cosine_threshold is None:
             raise ValueError(
                 "cosine_better_than_threshold must be specified in vector_db_storage_cls_kwargs"
@@ -138,12 +137,9 @@ class QdrantVectorDBStorage(BaseVectorStorage):
             query_vector=embedding[0],
             limit=top_k,
             with_payload=True,
+            score_threshold=self.cosine_better_than_threshold,
         )
+
         logger.debug(f"query result: {results}")
-        # 添加余弦相似度过滤
-        filtered_results = [
-            dp for dp in results if dp.score >= self.cosine_better_than_threshold
-        ]
-        return [
-            {**dp.payload, "id": dp.id, "distance": dp.score} for dp in filtered_results
-        ]
+
+        return [{**dp.payload, "id": dp.id, "distance": dp.score} for dp in results]

lightrag/lightrag.py

@@ -80,7 +80,12 @@ STORAGE_IMPLEMENTATIONS = {
         "required_methods": ["query", "upsert"],
     },
     "DOC_STATUS_STORAGE": {
-        "implementations": ["JsonDocStatusStorage", "PGDocStatusStorage"],
+        "implementations": [
+            "JsonDocStatusStorage",
+            "PGDocStatusStorage",
+            "PGDocStatusStorage",
+            "MongoDocStatusStorage",
+        ],
         "required_methods": ["get_pending_docs"],
     },
 }
@@ -421,7 +426,7 @@ class LightRAG:
             # Verify storage implementation compatibility
             self.verify_storage_implementation(storage_type, storage_name)
             # Check environment variables
-            self.check_storage_env_vars(storage_name)
+            # self.check_storage_env_vars(storage_name)
 
         # Ensure vector_db_storage_cls_kwargs has required fields
         default_vector_db_kwargs = {