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

cleanup code

Browse Source
This commit is contained in:
Yannick Stephan
2025-02-08 23:58:15 +01:00
parent 2929d1fc39
commit 50c7f26262
8 changed files with 63 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
lightrag/base.py
View File

@@ -104,9 +104,12 @@ class BaseKVStorage(StorageNameSpace):
async def drop(self) -> None:
raise NotImplementedError
async def get_by_status_and_ids(self, status: str) -> Union[list[dict[str, Any]], None]:
async def get_by_status_and_ids(
self, status: str
) -> Union[list[dict[str, Any]], None]:
raise NotImplementedError
@dataclass
class BaseGraphStorage(StorageNameSpace):
embedding_func: EmbeddingFunc = None

5
lightrag/kg/json_kv_impl.py
View File

@@ -12,6 +12,7 @@ from lightrag.base import (
BaseKVStorage,
)
@dataclass
class JsonKVStorage(BaseKVStorage):
def __post_init__(self):
@@ -50,6 +51,8 @@ class JsonKVStorage(BaseKVStorage):
async def drop(self) -> None:
self._data = {}
async def get_by_status_and_ids(self, status: str) -> Union[list[dict[str, Any]], None]:
async def get_by_status_and_ids(
self, status: str
) -> Union[list[dict[str, Any]], None]:
result = [v for _, v in self._data.items() if v["status"] == status]
return result if result else None

4
lightrag/kg/mongo_impl.py
View File

@@ -77,7 +77,9 @@ class MongoKVStorage(BaseKVStorage):
"""Drop the collection"""
await self._data.drop()
async def get_by_status_and_ids(self, status: str) -> Union[list[dict[str, Any]], None]:
async def get_by_status_and_ids(
self, status: str
) -> Union[list[dict[str, Any]], None]:
"""Get documents by status and ids"""
return self._data.find({"status": status})

1
lightrag/kg/oracle_impl.py
View File

@@ -327,6 +327,7 @@ class OracleKVStorage(BaseKVStorage):
):
logger.info("full doc and chunk data had been saved into oracle db!")
@dataclass
class OracleVectorDBStorage(BaseVectorStorage):
# should pass db object to self.db

4
lightrag/kg/postgres_impl.py
View File

@@ -238,7 +238,9 @@ class PGKVStorage(BaseKVStorage):
else:
return None
async def get_by_status_and_ids(self, status: str) -> Union[list[dict[str, Any]], None]:
async def get_by_status_and_ids(
self, status: str
) -> Union[list[dict[str, Any]], None]:
"""Specifically for llm_response_cache."""
SQL = SQL_TEMPLATES["get_by_status_" + self.namespace]
params = {"workspace": self.db.workspace, "status": status}

5
lightrag/kg/redis_impl.py
View File

@@ -59,10 +59,11 @@ class RedisKVStorage(BaseKVStorage):
if keys:
await self._redis.delete(*keys)
async def get_by_status_and_ids(self, status: str) -> Union[list[dict[str, Any]], None]:
async def get_by_status_and_ids(
self, status: str
) -> Union[list[dict[str, Any]], None]:
pipe = self._redis.pipeline()
for key in await self._redis.keys(f"{self.namespace}:*"):
pipe.hgetall(key)
results = await pipe.execute()
return [data for data in results if data.get("status") == status] or None

5
lightrag/kg/tidb_impl.py
View File

@@ -333,11 +333,14 @@ class TiDBVectorDBStorage(BaseVectorStorage):
merge_sql = SQL_TEMPLATES["insert_relationship"]
await self.db.execute(merge_sql, data)
async def get_by_status_and_ids(self, status: str) -> Union[list[dict[str, Any]], None]:
async def get_by_status_and_ids(
self, status: str
) -> Union[list[dict[str, Any]], None]:
SQL = SQL_TEMPLATES["get_by_status_" + self.namespace]
params = {"workspace": self.db.workspace, "status": status}
return await self.db.query(SQL, params, multirows=True)
@dataclass
class TiDBGraphStorage(BaseGraphStorage):
def __post_init__(self):

39
lightrag/lightrag.py
View File

@@ -629,12 +629,7 @@ class LightRAG:
# 4. Store original document
for doc_id, doc in new_docs.items():
await self.full_docs.upsert(
{
doc_id: {
"content": doc["content"],
"status": DocStatus.PENDING
}
}
{doc_id: {"content": doc["content"], "status": DocStatus.PENDING}}
)
logger.info(f"Stored {len(new_docs)} new unique documents")
@@ -643,8 +638,12 @@ class LightRAG:
# 1. get all pending and failed documents
_todo_doc_keys = []
_failed_doc = await self.full_docs.get_by_status_and_ids(status=DocStatus.FAILED)
_pendding_doc = await self.full_docs.get_by_status_and_ids(status=DocStatus.PENDING)
_failed_doc = await self.full_docs.get_by_status_and_ids(
status=DocStatus.FAILED
)
_pendding_doc = await self.full_docs.get_by_status_and_ids(
status=DocStatus.PENDING
)
if _failed_doc:
_todo_doc_keys.extend([doc["id"] for doc in _failed_doc])
@@ -690,10 +689,14 @@ class LightRAG:
# Store chunks in vector database
await self.chunks_vdb.upsert(chunks)
# Update doc status
await self.text_chunks.upsert({**chunks, "status": DocStatus.PENDING})
await self.text_chunks.upsert(
{**chunks, "status": DocStatus.PENDING}
)
except Exception as e:
# Mark as failed if any step fails
await self.text_chunks.upsert({**chunks, "status": DocStatus.FAILED})
await self.text_chunks.upsert(
{**chunks, "status": DocStatus.FAILED}
)
raise e
except Exception as e:
import traceback
@@ -707,8 +710,12 @@ class LightRAG:
"""Get pendding or failed chunks, extract entities and relationships from each chunk"""
# 1. get all pending and failed chunks
_todo_chunk_keys = []
_failed_chunks = await self.text_chunks.get_by_status_and_ids(status=DocStatus.FAILED)
_pendding_chunks = await self.text_chunks.get_by_status_and_ids(status=DocStatus.PENDING)
_failed_chunks = await self.text_chunks.get_by_status_and_ids(
status=DocStatus.FAILED
)
_pendding_chunks = await self.text_chunks.get_by_status_and_ids(
status=DocStatus.PENDING
)
if _failed_chunks:
_todo_chunk_keys.extend([doc["id"] for doc in _failed_chunks])
if _pendding_chunks:
@@ -742,11 +749,15 @@ class LightRAG:
if maybe_new_kg is None:
logger.info("No entities or relationships extracted!")
# Update status to processed
await self.text_chunks.upsert({chunk_id: {"status": DocStatus.PROCESSED}})
await self.text_chunks.upsert(
{chunk_id: {"status": DocStatus.PROCESSED}}
)
except Exception as e:
logger.error("Failed to extract entities and relationships")
# Mark as failed if any step fails
await self.text_chunks.upsert({chunk_id: {"status": DocStatus.FAILED}})
await self.text_chunks.upsert(
{chunk_id: {"status": DocStatus.FAILED}}
)
raise e
with tqdm_async(