diff --git a/lightrag/kg/jsondocstatus_impl.py b/lightrag/kg/jsondocstatus_impl.py index 35d31fba..1750b202 100644 --- a/lightrag/kg/jsondocstatus_impl.py +++ b/lightrag/kg/jsondocstatus_impl.py @@ -122,16 +122,8 @@ class JsonDocStatusStorage(DocStatusStorage): self._data.update(data) await self.index_done_callback() - async def get_by_id(self, id: str) -> dict[str, Any]: - return self._data.get(id, {}) - - async def get(self, doc_id: str) -> Union[DocProcessingStatus, None]: - """Get document status by ID""" - data = self._data.get(doc_id) - if data: - return DocProcessingStatus(**data) - else: - return None + async def get_by_id(self, id: str) -> Union[dict[str, Any], None]: + return self._data.get(id) async def delete(self, doc_ids: list[str]): """Delete document status by IDs""" diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index 440dbf78..40cc8978 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -531,7 +531,7 @@ class LightRAG: if not doc: await self.full_docs.upsert({doc_id: {"content": content}}) - async def insert_doc_status(doc_id: str, chunks: dict[str, Any]): + async def insert_text_chunks(doc_id: str, chunks: dict[str, Any]): # Check if chunks are already processed doc = await self.text_chunks.get_by_id(doc_id) if not doc: @@ -598,7 +598,7 @@ class LightRAG: tasks.append(self.chunks_vdb.upsert(chunks)) tasks.append(self._process_entity_relation_graph(chunks)) tasks.append(insert_full_doc(doc_id, status_doc.content)) - tasks.append(insert_doc_status(doc_id, chunks)) + tasks.append(insert_text_chunks(doc_id, chunks)) try: await asyncio.gather(*tasks)