llama-index

1
from typing import List, Optional
2

3
from llama_index.legacy.callbacks.base import CallbackManager
4
from llama_index.legacy.core.base_retriever import BaseRetriever
5
from llama_index.legacy.indices.query.query_transform.base import BaseQueryTransform
6
from llama_index.legacy.prompts.mixin import PromptMixinType
7
from llama_index.legacy.schema import NodeWithScore, QueryBundle
8

9

10
class TransformRetriever(BaseRetriever):
11
    """Transform Retriever.
12

13
    Takes in an existing retriever and a query transform and runs the query transform
14
    before running the retriever.
15

16
    """
17

18
    def __init__(
19
        self,
20
        retriever: BaseRetriever,
21
        query_transform: BaseQueryTransform,
22
        transform_metadata: Optional[dict] = None,
23
        callback_manager: Optional[CallbackManager] = None,
24
        object_map: Optional[dict] = None,
25
        verbose: bool = False,
26
    ) -> None:
27
        self._retriever = retriever
28
        self._query_transform = query_transform
29
        self._transform_metadata = transform_metadata
30
        super().__init__(
31
            callback_manager=callback_manager, object_map=object_map, verbose=verbose
32
        )
33

34
    def _get_prompt_modules(self) -> PromptMixinType:
35
        """Get prompt sub-modules."""
36
        # NOTE: don't include tools for now
37
        return {"query_transform": self._query_transform}
38

39
    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
40
        query_bundle = self._query_transform.run(
41
            query_bundle, metadata=self._transform_metadata
42
        )
43
        return self._retriever.retrieve(query_bundle)
44