From 7848a38a45ad116e6ae53dd784b55a1cb5c3bbbd Mon Sep 17 00:00:00 2001 From: Yannick Stephan Date: Sun, 16 Feb 2025 14:11:19 +0100 Subject: [PATCH] added all abstractmethod --- lightrag/base.py | 72 +++++++++++++++++++++++++----------------------- 1 file changed, 38 insertions(+), 34 deletions(-) diff --git a/lightrag/base.py b/lightrag/base.py index 8e3a7ecf..fc4702d4 100644 --- a/lightrag/base.py +++ b/lightrag/base.py @@ -1,5 +1,6 @@ from __future__ import annotations +from abc import ABC, abstractmethod import os from dataclasses import dataclass, field from enum import Enum @@ -79,126 +80,126 @@ class QueryParam: @dataclass -class StorageNameSpace: +class StorageNameSpace(ABC): namespace: str global_config: dict[str, Any] + @abstractmethod async def index_done_callback(self) -> None: """Commit the storage operations after indexing""" - pass @dataclass -class BaseVectorStorage(StorageNameSpace): +class BaseVectorStorage(StorageNameSpace, ABC): embedding_func: EmbeddingFunc meta_fields: set[str] = field(default_factory=set) + @abstractmethod async def query(self, query: str, top_k: int) -> list[dict[str, Any]]: """Query the vector storage and retrieve top_k results.""" - raise NotImplementedError + @abstractmethod async def upsert(self, data: dict[str, dict[str, Any]]) -> None: """Insert or update vectors in the storage.""" - raise NotImplementedError + @abstractmethod async def delete_entity(self, entity_name: str) -> None: """Delete a single entity by its name.""" - raise NotImplementedError + @abstractmethod async def delete_entity_relation(self, entity_name: str) -> None: """Delete relations for a given entity.""" - raise NotImplementedError @dataclass -class BaseKVStorage(StorageNameSpace): - embedding_func: EmbeddingFunc | None = None +class BaseKVStorage(StorageNameSpace, ABC): + embedding_func: EmbeddingFunc + @abstractmethod async def get_by_id(self, id: str) -> dict[str, Any] | None: """Get value by id""" - raise NotImplementedError + @abstractmethod async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]: """Get values by ids""" - raise NotImplementedError + @abstractmethod async def filter_keys(self, keys: set[str]) -> set[str]: """Return un-exist keys""" - raise NotImplementedError + @abstractmethod async def upsert(self, data: dict[str, dict[str, Any]]) -> None: """Upsert data""" - raise NotImplementedError + @abstractmethod async def drop(self) -> None: """Drop the storage""" - raise NotImplementedError @dataclass -class BaseGraphStorage(StorageNameSpace): - embedding_func: EmbeddingFunc | None = None +class BaseGraphStorage(StorageNameSpace, ABC): + embedding_func: EmbeddingFunc """Check if a node exists in the graph.""" + @abstractmethod async def has_node(self, node_id: str) -> bool: """Check if an edge exists in the graph.""" - raise NotImplementedError + @abstractmethod async def has_edge(self, source_node_id: str, target_node_id: str) -> bool: """Get the degree of a node.""" - raise NotImplementedError + @abstractmethod async def node_degree(self, node_id: str) -> int: """Get the degree of an edge.""" - raise NotImplementedError + @abstractmethod async def edge_degree(self, src_id: str, tgt_id: str) -> int: """Get a node by its id.""" - raise NotImplementedError + @abstractmethod async def get_node(self, node_id: str) -> dict[str, str] | None: """Get an edge by its source and target node ids.""" - raise NotImplementedError + @abstractmethod async def get_edge( self, source_node_id: str, target_node_id: str ) -> dict[str, str] | None: """Get all edges connected to a node.""" - raise NotImplementedError + @abstractmethod async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None: """Upsert a node into the graph.""" - raise NotImplementedError + @abstractmethod async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None: """Upsert an edge into the graph.""" - raise NotImplementedError + @abstractmethod async def upsert_edge( self, source_node_id: str, target_node_id: str, edge_data: dict[str, str] ) -> None: """Delete a node from the graph.""" - raise NotImplementedError + @abstractmethod async def delete_node(self, node_id: str) -> None: """Embed nodes using an algorithm.""" - raise NotImplementedError + @abstractmethod async def embed_nodes( self, algorithm: str ) -> tuple[np.ndarray[Any, Any], list[str]]: """Get all labels in the graph.""" - raise NotImplementedError("Node embedding is not used in lightrag.") + @abstractmethod async def get_all_labels(self) -> list[str]: """Get a knowledge graph of a node.""" - raise NotImplementedError + @abstractmethod async def get_knowledge_graph( self, node_label: str, max_depth: int = 5 ) -> KnowledgeGraph: """Retrieve a subgraph of the knowledge graph starting from a given node.""" - raise NotImplementedError class DocStatus(str, Enum): @@ -234,29 +235,32 @@ class DocProcessingStatus: """Additional metadata""" -class DocStatusStorage(BaseKVStorage): +@dataclass +class DocStatusStorage(BaseKVStorage, ABC): """Base class for document status storage""" + @abstractmethod async def get_status_counts(self) -> dict[str, int]: """Get counts of documents in each status""" - raise NotImplementedError + @abstractmethod async def get_failed_docs(self) -> dict[str, DocProcessingStatus]: """Get all failed documents""" - raise NotImplementedError + @abstractmethod async def get_pending_docs(self) -> dict[str, DocProcessingStatus]: """Get all pending documents""" raise NotImplementedError + @abstractmethod async def get_processing_docs(self) -> dict[str, DocProcessingStatus]: """Get all processing documents""" - raise NotImplementedError + @abstractmethod async def get_processed_docs(self) -> dict[str, DocProcessingStatus]: """Get all procesed documents""" - raise NotImplementedError + @abstractmethod async def update_doc_status(self, data: dict[str, Any]) -> None: """Updates the status of a document. By default, it calls upsert.""" await self.upsert(data)