MetaGPT

1
"""Base retriever."""
2

3
from abc import abstractmethod
4

5
from llama_index.core.retrievers import BaseRetriever
6
from llama_index.core.schema import BaseNode, NodeWithScore, QueryType
7

8
from metagpt.utils.reflection import check_methods
9

10

11
class RAGRetriever(BaseRetriever):
12
    """Inherit from llama_index"""
13

14
    @abstractmethod
15
    async def _aretrieve(self, query: QueryType) -> list[NodeWithScore]:
16
        """Retrieve nodes"""
17

18
    def _retrieve(self, query: QueryType) -> list[NodeWithScore]:
19
        """Retrieve nodes"""
20

21

22
class ModifiableRAGRetriever(RAGRetriever):
23
    """Support modification."""
24

25
    @classmethod
26
    def __subclasshook__(cls, C):
27
        if cls is ModifiableRAGRetriever:
28
            return check_methods(C, "add_nodes")
29
        return NotImplemented
30

31
    @abstractmethod
32
    def add_nodes(self, nodes: list[BaseNode], **kwargs) -> None:
33
        """To support add docs, must inplement this func"""
34

35

36
class PersistableRAGRetriever(RAGRetriever):
37
    """Support persistent."""
38

39
    @classmethod
40
    def __subclasshook__(cls, C):
41
        if cls is PersistableRAGRetriever:
42
            return check_methods(C, "persist")
43
        return NotImplemented
44

45
    @abstractmethod
46
    def persist(self, persist_dir: str, **kwargs) -> None:
47
        """To support persist, must inplement this func"""
48