added all abstractmethod
This commit is contained in:
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from abc import ABC, abstractmethod
|
||||||
import os
|
import os
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
@@ -79,126 +80,126 @@ class QueryParam:
|
|||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class StorageNameSpace:
|
class StorageNameSpace(ABC):
|
||||||
namespace: str
|
namespace: str
|
||||||
global_config: dict[str, Any]
|
global_config: dict[str, Any]
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def index_done_callback(self) -> None:
|
async def index_done_callback(self) -> None:
|
||||||
"""Commit the storage operations after indexing"""
|
"""Commit the storage operations after indexing"""
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class BaseVectorStorage(StorageNameSpace):
|
class BaseVectorStorage(StorageNameSpace, ABC):
|
||||||
embedding_func: EmbeddingFunc
|
embedding_func: EmbeddingFunc
|
||||||
meta_fields: set[str] = field(default_factory=set)
|
meta_fields: set[str] = field(default_factory=set)
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def query(self, query: str, top_k: int) -> list[dict[str, Any]]:
|
async def query(self, query: str, top_k: int) -> list[dict[str, Any]]:
|
||||||
"""Query the vector storage and retrieve top_k results."""
|
"""Query the vector storage and retrieve top_k results."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
||||||
"""Insert or update vectors in the storage."""
|
"""Insert or update vectors in the storage."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def delete_entity(self, entity_name: str) -> None:
|
async def delete_entity(self, entity_name: str) -> None:
|
||||||
"""Delete a single entity by its name."""
|
"""Delete a single entity by its name."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def delete_entity_relation(self, entity_name: str) -> None:
|
async def delete_entity_relation(self, entity_name: str) -> None:
|
||||||
"""Delete relations for a given entity."""
|
"""Delete relations for a given entity."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class BaseKVStorage(StorageNameSpace):
|
class BaseKVStorage(StorageNameSpace, ABC):
|
||||||
embedding_func: EmbeddingFunc | None = None
|
embedding_func: EmbeddingFunc
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def get_by_id(self, id: str) -> dict[str, Any] | None:
|
async def get_by_id(self, id: str) -> dict[str, Any] | None:
|
||||||
"""Get value by id"""
|
"""Get value by id"""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
|
async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
|
||||||
"""Get values by ids"""
|
"""Get values by ids"""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def filter_keys(self, keys: set[str]) -> set[str]:
|
async def filter_keys(self, keys: set[str]) -> set[str]:
|
||||||
"""Return un-exist keys"""
|
"""Return un-exist keys"""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
||||||
"""Upsert data"""
|
"""Upsert data"""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def drop(self) -> None:
|
async def drop(self) -> None:
|
||||||
"""Drop the storage"""
|
"""Drop the storage"""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class BaseGraphStorage(StorageNameSpace):
|
class BaseGraphStorage(StorageNameSpace, ABC):
|
||||||
embedding_func: EmbeddingFunc | None = None
|
embedding_func: EmbeddingFunc
|
||||||
"""Check if a node exists in the graph."""
|
"""Check if a node exists in the graph."""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def has_node(self, node_id: str) -> bool:
|
async def has_node(self, node_id: str) -> bool:
|
||||||
"""Check if an edge exists in the graph."""
|
"""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:
|
async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
|
||||||
"""Get the degree of a node."""
|
"""Get the degree of a node."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def node_degree(self, node_id: str) -> int:
|
async def node_degree(self, node_id: str) -> int:
|
||||||
"""Get the degree of an edge."""
|
"""Get the degree of an edge."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def edge_degree(self, src_id: str, tgt_id: str) -> int:
|
async def edge_degree(self, src_id: str, tgt_id: str) -> int:
|
||||||
"""Get a node by its id."""
|
"""Get a node by its id."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def get_node(self, node_id: str) -> dict[str, str] | None:
|
async def get_node(self, node_id: str) -> dict[str, str] | None:
|
||||||
"""Get an edge by its source and target node ids."""
|
"""Get an edge by its source and target node ids."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def get_edge(
|
async def get_edge(
|
||||||
self, source_node_id: str, target_node_id: str
|
self, source_node_id: str, target_node_id: str
|
||||||
) -> dict[str, str] | None:
|
) -> dict[str, str] | None:
|
||||||
"""Get all edges connected to a node."""
|
"""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:
|
async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None:
|
||||||
"""Upsert a node into the graph."""
|
"""Upsert a node into the graph."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
|
async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
|
||||||
"""Upsert an edge into the graph."""
|
"""Upsert an edge into the graph."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def upsert_edge(
|
async def upsert_edge(
|
||||||
self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
|
self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Delete a node from the graph."""
|
"""Delete a node from the graph."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def delete_node(self, node_id: str) -> None:
|
async def delete_node(self, node_id: str) -> None:
|
||||||
"""Embed nodes using an algorithm."""
|
"""Embed nodes using an algorithm."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def embed_nodes(
|
async def embed_nodes(
|
||||||
self, algorithm: str
|
self, algorithm: str
|
||||||
) -> tuple[np.ndarray[Any, Any], list[str]]:
|
) -> tuple[np.ndarray[Any, Any], list[str]]:
|
||||||
"""Get all labels in the graph."""
|
"""Get all labels in the graph."""
|
||||||
raise NotImplementedError("Node embedding is not used in lightrag.")
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def get_all_labels(self) -> list[str]:
|
async def get_all_labels(self) -> list[str]:
|
||||||
"""Get a knowledge graph of a node."""
|
"""Get a knowledge graph of a node."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def get_knowledge_graph(
|
async def get_knowledge_graph(
|
||||||
self, node_label: str, max_depth: int = 5
|
self, node_label: str, max_depth: int = 5
|
||||||
) -> KnowledgeGraph:
|
) -> KnowledgeGraph:
|
||||||
"""Retrieve a subgraph of the knowledge graph starting from a given node."""
|
"""Retrieve a subgraph of the knowledge graph starting from a given node."""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
|
||||||
class DocStatus(str, Enum):
|
class DocStatus(str, Enum):
|
||||||
@@ -234,29 +235,32 @@ class DocProcessingStatus:
|
|||||||
"""Additional metadata"""
|
"""Additional metadata"""
|
||||||
|
|
||||||
|
|
||||||
class DocStatusStorage(BaseKVStorage):
|
@dataclass
|
||||||
|
class DocStatusStorage(BaseKVStorage, ABC):
|
||||||
"""Base class for document status storage"""
|
"""Base class for document status storage"""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
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"""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
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"""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
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"""
|
||||||
raise NotImplementedError
|
raise NotImplementedError
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def get_processing_docs(self) -> dict[str, DocProcessingStatus]:
|
async def get_processing_docs(self) -> dict[str, DocProcessingStatus]:
|
||||||
"""Get all processing documents"""
|
"""Get all processing documents"""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def get_processed_docs(self) -> dict[str, DocProcessingStatus]:
|
async def get_processed_docs(self) -> dict[str, DocProcessingStatus]:
|
||||||
"""Get all procesed documents"""
|
"""Get all procesed documents"""
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
async def update_doc_status(self, data: dict[str, Any]) -> None:
|
async def update_doc_status(self, data: dict[str, Any]) -> None:
|
||||||
"""Updates the status of a document. By default, it calls upsert."""
|
"""Updates the status of a document. By default, it calls upsert."""
|
||||||
await self.upsert(data)
|
await self.upsert(data)
|
||||||
|
Reference in New Issue
Block a user