From 8413537ad2879c830d6769704db22c1d2987f7ef Mon Sep 17 00:00:00 2001
From: Yannick Stephan <stephan.yannick@me.com>
Date: Thu, 20 Feb 2025 13:44:17 +0100
Subject: [PATCH] cleanup

---
 lightrag/kg/__init__.py | 35 +++++++++++++++++------------------
 lightrag/lightrag.py    | 16 ++++++++++------
 2 files changed, 27 insertions(+), 24 deletions(-)

diff --git a/lightrag/kg/__init__.py b/lightrag/kg/__init__.py
index 18aa2796..4943fc1d 100644
--- a/lightrag/kg/__init__.py
+++ b/lightrag/kg/__init__.py
@@ -135,24 +135,23 @@ STORAGES = {
     "QdrantVectorDBStorage": ".kg.qdrant_impl",
 }
 
-def verify_storage_implementation(
-        storage_type: str, storage_name: str
-    ) -> None:
-        """Verify if storage implementation is compatible with specified storage type
 
-        Args:
-            storage_type: Storage type (KV_STORAGE, GRAPH_STORAGE etc.)
-            storage_name: Storage implementation name
+def verify_storage_implementation(storage_type: str, storage_name: str) -> None:
+    """Verify if storage implementation is compatible with specified storage type
 
-        Raises:
-            ValueError: If storage implementation is incompatible or missing required methods
-        """
-        if storage_type not in STORAGE_IMPLEMENTATIONS:
-            raise ValueError(f"Unknown storage type: {storage_type}")
+    Args:
+        storage_type: Storage type (KV_STORAGE, GRAPH_STORAGE etc.)
+        storage_name: Storage implementation name
 
-        storage_info = STORAGE_IMPLEMENTATIONS[storage_type]
-        if storage_name not in storage_info["implementations"]:
-            raise ValueError(
-                f"Storage implementation '{storage_name}' is not compatible with {storage_type}. "
-                f"Compatible implementations are: {', '.join(storage_info['implementations'])}"
-            )
\ No newline at end of file
+    Raises:
+        ValueError: If storage implementation is incompatible or missing required methods
+    """
+    if storage_type not in STORAGE_IMPLEMENTATIONS:
+        raise ValueError(f"Unknown storage type: {storage_type}")
+
+    storage_info = STORAGE_IMPLEMENTATIONS[storage_type]
+    if storage_name not in storage_info["implementations"]:
+        raise ValueError(
+            f"Storage implementation '{storage_name}' is not compatible with {storage_type}. "
+            f"Compatible implementations are: {', '.join(storage_info['implementations'])}"
+        )
diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py
index ca9733ce..5bb05764 100644
--- a/lightrag/lightrag.py
+++ b/lightrag/lightrag.py
@@ -8,7 +8,11 @@ from datetime import datetime
 from functools import partial
 from typing import Any, AsyncIterator, Callable, Iterator, cast, final
 
-from lightrag.kg import STORAGE_ENV_REQUIREMENTS, STORAGES, verify_storage_implementation
+from lightrag.kg import (
+    STORAGE_ENV_REQUIREMENTS,
+    STORAGES,
+    verify_storage_implementation,
+)
 
 from .base import (
     BaseGraphStorage,
@@ -251,6 +255,10 @@ class LightRAG:
     The default function is :func:`.utils.convert_response_to_json`.
     """
 
+    cosine_better_than_threshold: float = field(
+        default=float(os.getenv("COSINE_THRESHOLD", 0.2))
+    )
+
     _storages_status: StoragesStatus = field(default=StoragesStatus.NOT_CREATED)
 
     def __post_init__(self):
@@ -278,11 +286,8 @@ class LightRAG:
             # self.check_storage_env_vars(storage_name)
 
         # Ensure vector_db_storage_cls_kwargs has required fields
-        default_vector_db_kwargs = {
-            "cosine_better_than_threshold": float(os.getenv("COSINE_THRESHOLD", "0.2"))
-        }
         self.vector_db_storage_cls_kwargs = {
-            **default_vector_db_kwargs,
+            "cosine_better_than_threshold": self.cosine_better_than_threshold,
             **self.vector_db_storage_cls_kwargs,
         }
 
@@ -1463,7 +1468,6 @@ class LightRAG:
 
         return result
 
-
     def check_storage_env_vars(self, storage_name: str) -> None:
         """Check if all required environment variables for storage implementation exist