cleaned code

This commit is contained in:
Yannick Stephan
2025-02-09 15:24:30 +01:00
parent 28b53144d9
commit 7116bd1872
4 changed files with 9 additions and 8 deletions

View File

@@ -50,7 +50,7 @@ Usage:
import os
from dataclasses import dataclass
from typing import Any, Union, Dict
from typing import Any, Union
from lightrag.utils import (
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"""
counts = {status: 0 for status in DocStatus}
for doc in self._data.values():
counts[doc["status"]] += 1
return counts
async def get_failed_docs(self) -> Dict[str, DocProcessingStatus]:
async def get_failed_docs(self) -> dict[str, DocProcessingStatus]:
"""Get all failed documents"""
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"""
return {k: v for k, v in self._data.items() if v["status"] == DocStatus.PENDING}

View File

@@ -74,6 +74,7 @@ class MongoKVStorage(BaseKVStorage):
"""Drop the collection"""
await self._data.drop()
@dataclass
class MongoGraphStorage(BaseGraphStorage):
"""