cleaned code
This commit is contained in:
@@ -50,7 +50,7 @@ Usage:
|
|||||||
|
|
||||||
import os
|
import os
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import Any, Union, Dict
|
from typing import Any, Union
|
||||||
|
|
||||||
from lightrag.utils import (
|
from lightrag.utils import (
|
||||||
logger,
|
logger,
|
||||||
@@ -85,18 +85,18 @@ class JsonDocStatusStorage(DocStatusStorage):
|
|||||||
]
|
]
|
||||||
)
|
)
|
||||||
|
|
||||||
async def get_status_counts(self) -> Dict[str, int]:
|
async def get_status_counts(self) -> dict[str, int]:
|
||||||
"""Get counts of documents in each status"""
|
"""Get counts of documents in each status"""
|
||||||
counts = {status: 0 for status in DocStatus}
|
counts = {status: 0 for status in DocStatus}
|
||||||
for doc in self._data.values():
|
for doc in self._data.values():
|
||||||
counts[doc["status"]] += 1
|
counts[doc["status"]] += 1
|
||||||
return counts
|
return counts
|
||||||
|
|
||||||
async def get_failed_docs(self) -> Dict[str, DocProcessingStatus]:
|
async def get_failed_docs(self) -> dict[str, DocProcessingStatus]:
|
||||||
"""Get all failed documents"""
|
"""Get all failed documents"""
|
||||||
return {k: v for k, v in self._data.items() if v["status"] == DocStatus.FAILED}
|
return {k: v for k, v in self._data.items() if v["status"] == DocStatus.FAILED}
|
||||||
|
|
||||||
async def get_pending_docs(self) -> Dict[str, DocProcessingStatus]:
|
async def get_pending_docs(self) -> dict[str, DocProcessingStatus]:
|
||||||
"""Get all pending documents"""
|
"""Get all pending documents"""
|
||||||
return {k: v for k, v in self._data.items() if v["status"] == DocStatus.PENDING}
|
return {k: v for k, v in self._data.items() if v["status"] == DocStatus.PENDING}
|
||||||
|
|
||||||
|
@@ -74,6 +74,7 @@ class MongoKVStorage(BaseKVStorage):
|
|||||||
"""Drop the collection"""
|
"""Drop the collection"""
|
||||||
await self._data.drop()
|
await self._data.drop()
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class MongoGraphStorage(BaseGraphStorage):
|
class MongoGraphStorage(BaseGraphStorage):
|
||||||
"""
|
"""
|
||||||
|
Reference in New Issue
Block a user