updated type

This commit is contained in:
Yannick Stephan
2025-02-09 21:12:39 +01:00
parent 3e49ced307
commit 068924b59f
2 changed files with 4 additions and 12 deletions

View File

@@ -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"""

View File

@@ -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)