gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Inject TiDB同LightRAG storage when needed

Browse Source
This commit is contained in:
yangdx
2025-02-11 04:27:45 +08:00
parent 5408e7ea02
commit c5c606f491
2 changed files with 58 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
lightrag/kg/tidb_impl.py
View File

@@ -102,6 +102,8 @@ class TiDB:
@dataclass
class TiDBKVStorage(BaseKVStorage):
# should pass db object to self.db
db: TiDB = None
def __post_init__(self):
self._data = {}
self._max_batch_size = self.global_config["embedding_batch_num"]
@@ -208,6 +210,8 @@ class TiDBKVStorage(BaseKVStorage):
@dataclass
class TiDBVectorDBStorage(BaseVectorStorage):
# should pass db object to self.db
db: TiDB = None
cosine_better_than_threshold: float = float(os.getenv("COSINE_THRESHOLD", "0.2"))
def __post_init__(self):
@@ -329,6 +333,9 @@ class TiDBVectorDBStorage(BaseVectorStorage):
@dataclass
class TiDBGraphStorage(BaseGraphStorage):
# should pass db object to self.db
db: TiDB = None
def __post_init__(self):
self._max_batch_size = self.global_config["embedding_batch_num"]