llama-index

1
from typing import List, Optional, Sequence
2

3
from llama_index.legacy.callbacks.base import CallbackManager
4
from llama_index.legacy.core.base_query_engine import BaseQueryEngine
5
from llama_index.legacy.core.response.schema import RESPONSE_TYPE
6
from llama_index.legacy.indices.query.query_transform.base import BaseQueryTransform
7
from llama_index.legacy.prompts.mixin import PromptMixinType
8
from llama_index.legacy.schema import NodeWithScore, QueryBundle
9

10

11
class TransformQueryEngine(BaseQueryEngine):
12
    """Transform query engine.
13

14
    Applies a query transform to a query bundle before passing
15
        it to a query engine.
16

17
    Args:
18
        query_engine (BaseQueryEngine): A query engine object.
19
        query_transform (BaseQueryTransform): A query transform object.
20
        transform_metadata (Optional[dict]): metadata to pass to the
21
            query transform.
22
        callback_manager (Optional[CallbackManager]): A callback manager.
23

24
    """
25

26
    def __init__(
27
        self,
28
        query_engine: BaseQueryEngine,
29
        query_transform: BaseQueryTransform,
30
        transform_metadata: Optional[dict] = None,
31
        callback_manager: Optional[CallbackManager] = None,
32
    ) -> None:
33
        self._query_engine = query_engine
34
        self._query_transform = query_transform
35
        self._transform_metadata = transform_metadata
36
        super().__init__(callback_manager)
37

38
    def _get_prompt_modules(self) -> PromptMixinType:
39
        """Get prompt sub-modules."""
40
        return {
41
            "query_transform": self._query_transform,
42
            "query_engine": self._query_engine,
43
        }
44

45
    def retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
46
        query_bundle = self._query_transform.run(
47
            query_bundle, metadata=self._transform_metadata
48
        )
49
        return self._query_engine.retrieve(query_bundle)
50

51
    def synthesize(
52
        self,
53
        query_bundle: QueryBundle,
54
        nodes: List[NodeWithScore],
55
        additional_source_nodes: Optional[Sequence[NodeWithScore]] = None,
56
    ) -> RESPONSE_TYPE:
57
        query_bundle = self._query_transform.run(
58
            query_bundle, metadata=self._transform_metadata
59
        )
60
        return self._query_engine.synthesize(
61
            query_bundle=query_bundle,
62
            nodes=nodes,
63
            additional_source_nodes=additional_source_nodes,
64
        )
65

66
    async def asynthesize(
67
        self,
68
        query_bundle: QueryBundle,
69
        nodes: List[NodeWithScore],
70
        additional_source_nodes: Optional[Sequence[NodeWithScore]] = None,
71
    ) -> RESPONSE_TYPE:
72
        query_bundle = self._query_transform.run(
73
            query_bundle, metadata=self._transform_metadata
74
        )
75
        return await self._query_engine.asynthesize(
76
            query_bundle=query_bundle,
77
            nodes=nodes,
78
            additional_source_nodes=additional_source_nodes,
79
        )
80

81
    def _query(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
82
        """Answer a query."""
83
        query_bundle = self._query_transform.run(
84
            query_bundle, metadata=self._transform_metadata
85
        )
86
        return self._query_engine.query(query_bundle)
87

88
    async def _aquery(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
89
        """Answer a query."""
90
        query_bundle = self._query_transform.run(
91
            query_bundle, metadata=self._transform_metadata
92
        )
93
        return await self._query_engine.aquery(query_bundle)
94