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

Fix cosine threshold parameter setting error for chroma

Browse source
This commit is contained in:
yangdx 2025-01-29 22:41:18 +08:00
parent 90c765c724
commit 20d6355a4a
2 changed files with 10 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
lightrag/kg/chroma_impl.py
View file

@ -1,3 +1,4 @@
import os
import asyncio
from dataclasses import dataclass
from typing import Union
@ -12,16 +13,16 @@ from lightrag.utils import logger
class ChromaVectorDBStorage(BaseVectorStorage):
"""ChromaDB vector storage implementation."""
cosine_better_than_threshold: float = 0.2
cosine_better_than_threshold: float = float(os.getenv("COSINE_THRESHOLD", "0.2"))
def __post_init__(self):
try:
# Use global config value if specified, otherwise use default
self.cosine_better_than_threshold = self.global_config.get(
config = self.global_config.get("vector_db_storage_cls_kwargs", {})
self.cosine_better_than_threshold = config.get(
"cosine_better_than_threshold", self.cosine_better_than_threshold
)
config = self.global_config.get("vector_db_storage_cls_kwargs", {})
user_collection_settings = config.get("collection_settings", {})
# Default HNSW index settings for ChromaDB
default_collection_settings = {

14
lightrag/kg/nano_vector_db_impl.py
View file

@ -76,6 +76,12 @@ class NanoVectorDBStorage(BaseVectorStorage):
cosine_better_than_threshold: float = float(os.getenv("COSINE_THRESHOLD", "0.2"))
def __post_init__(self):
# Use global config value if specified, otherwise use default
config = self.global_config.get("vector_db_storage_cls_kwargs", {})
self.cosine_better_than_threshold = config.get(
"cosine_better_than_threshold", self.cosine_better_than_threshold
)
self._client_file_name = os.path.join(
self.global_config["working_dir"], f"vdb_{self.namespace}.json"
)
@ -83,14 +89,6 @@ class NanoVectorDBStorage(BaseVectorStorage):
self._client = NanoVectorDB(
self.embedding_func.embedding_dim, storage_file=self._client_file_name
)
# get cosine_better_than_threshold from LightRAG
vector_db_kwargs = self.global_config.get("vector_db_storage_cls_kwargs", {})
self.cosine_better_than_threshold = vector_db_kwargs.get(
"cosine_better_than_threshold",
self.global_config.get(
"cosine_better_than_threshold", self.cosine_better_than_threshold
),
)
async def upsert(self, data: dict[str, dict]):
logger.info(f"Inserting {len(data)} vectors to {self.namespace}")