llama-index

1
from abc import abstractmethod
2
from typing import List
3

4
from llama_index.legacy.indices.query.schema import QueryBundle, QueryType
5
from llama_index.legacy.prompts.mixin import PromptMixin
6
from llama_index.legacy.schema import NodeWithScore
7

8

9
class BaseImageRetriever(PromptMixin):
10
    """Base Image Retriever Abstraction."""
11

12
    def text_to_image_retrieve(
13
        self, str_or_query_bundle: QueryType
14
    ) -> List[NodeWithScore]:
15
        """Retrieve image nodes given query or single image input.
16

17
        Args:
18
            str_or_query_bundle (QueryType): a query text
19
            string or a QueryBundle object.
20
        """
21
        if isinstance(str_or_query_bundle, str):
22
            str_or_query_bundle = QueryBundle(query_str=str_or_query_bundle)
23
        return self._text_to_image_retrieve(str_or_query_bundle)
24

25
    @abstractmethod
26
    def _text_to_image_retrieve(
27
        self,
28
        query_bundle: QueryBundle,
29
    ) -> List[NodeWithScore]:
30
        """Retrieve image nodes or documents given query text.
31

32
        Implemented by the user.
33

34
        """
35

36
    def image_to_image_retrieve(
37
        self, str_or_query_bundle: QueryType
38
    ) -> List[NodeWithScore]:
39
        """Retrieve image nodes given single image input.
40

41
        Args:
42
            str_or_query_bundle (QueryType): a image path
43
            string or a QueryBundle object.
44
        """
45
        if isinstance(str_or_query_bundle, str):
46
            # leave query_str as empty since we are using image_path for image retrieval
47
            str_or_query_bundle = QueryBundle(
48
                query_str="", image_path=str_or_query_bundle
49
            )
50
        return self._image_to_image_retrieve(str_or_query_bundle)
51

52
    @abstractmethod
53
    def _image_to_image_retrieve(
54
        self,
55
        query_bundle: QueryBundle,
56
    ) -> List[NodeWithScore]:
57
        """Retrieve image nodes or documents given image.
58

59
        Implemented by the user.
60

61
        """
62

63
    # Async Methods
64
    async def atext_to_image_retrieve(
65
        self,
66
        str_or_query_bundle: QueryType,
67
    ) -> List[NodeWithScore]:
68
        if isinstance(str_or_query_bundle, str):
69
            str_or_query_bundle = QueryBundle(query_str=str_or_query_bundle)
70
        return await self._atext_to_image_retrieve(str_or_query_bundle)
71

72
    @abstractmethod
73
    async def _atext_to_image_retrieve(
74
        self,
75
        query_bundle: QueryBundle,
76
    ) -> List[NodeWithScore]:
77
        """Async retrieve image nodes or documents given query text.
78

79
        Implemented by the user.
80

81
        """
82

83
    async def aimage_to_image_retrieve(
84
        self,
85
        str_or_query_bundle: QueryType,
86
    ) -> List[NodeWithScore]:
87
        if isinstance(str_or_query_bundle, str):
88
            # leave query_str as empty since we are using image_path for image retrieval
89
            str_or_query_bundle = QueryBundle(
90
                query_str="", image_path=str_or_query_bundle
91
            )
92
        return await self._aimage_to_image_retrieve(str_or_query_bundle)
93

94
    @abstractmethod
95
    async def _aimage_to_image_retrieve(
96
        self,
97
        query_bundle: QueryBundle,
98
    ) -> List[NodeWithScore]:
99
        """Async retrieve image nodes or documents given image.
100

101
        Implemented by the user.
102

103
        """
104