llama-index

1
import logging
2
from typing import Optional
3

4
from llama_index.legacy.callbacks.base import CallbackManager
5
from llama_index.legacy.core.base_query_engine import BaseQueryEngine
6
from llama_index.legacy.core.response.schema import RESPONSE_TYPE, Response
7
from llama_index.legacy.evaluation.base import BaseEvaluator
8
from llama_index.legacy.evaluation.guideline import GuidelineEvaluator
9
from llama_index.legacy.indices.query.query_transform.feedback_transform import (
10
    FeedbackQueryTransformation,
11
)
12
from llama_index.legacy.prompts.mixin import PromptMixinType
13
from llama_index.legacy.schema import QueryBundle
14

15
logger = logging.getLogger(__name__)
16

17

18
class RetryQueryEngine(BaseQueryEngine):
19
    """Does retry on query engine if it fails evaluation.
20

21
    Args:
22
        query_engine (BaseQueryEngine): A query engine object
23
        evaluator (BaseEvaluator): An evaluator object
24
        max_retries (int): Maximum number of retries
25
        callback_manager (Optional[CallbackManager]): A callback manager object
26
    """
27

28
    def __init__(
29
        self,
30
        query_engine: BaseQueryEngine,
31
        evaluator: BaseEvaluator,
32
        max_retries: int = 3,
33
        callback_manager: Optional[CallbackManager] = None,
34
    ) -> None:
35
        self._query_engine = query_engine
36
        self._evaluator = evaluator
37
        self.max_retries = max_retries
38
        super().__init__(callback_manager)
39

40
    def _get_prompt_modules(self) -> PromptMixinType:
41
        """Get prompt sub-modules."""
42
        return {"query_engine": self._query_engine, "evaluator": self._evaluator}
43

44
    def _query(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
45
        """Answer a query."""
46
        response = self._query_engine._query(query_bundle)
47
        if self.max_retries <= 0:
48
            return response
49
        typed_response = (
50
            response if isinstance(response, Response) else response.get_response()
51
        )
52
        query_str = query_bundle.query_str
53
        eval = self._evaluator.evaluate_response(query_str, typed_response)
54
        if eval.passing:
55
            logger.debug("Evaluation returned True.")
56
            return response
57
        else:
58
            logger.debug("Evaluation returned False.")
59
            new_query_engine = RetryQueryEngine(
60
                self._query_engine, self._evaluator, self.max_retries - 1
61
            )
62
            query_transformer = FeedbackQueryTransformation()
63
            new_query = query_transformer.run(query_bundle, {"evaluation": eval})
64
            return new_query_engine.query(new_query)
65

66
    async def _aquery(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
67
        """Not supported."""
68
        return self._query(query_bundle)
69

70

71
class RetryGuidelineQueryEngine(BaseQueryEngine):
72
    """Does retry with evaluator feedback
73
    if query engine fails evaluation.
74

75
    Args:
76
        query_engine (BaseQueryEngine): A query engine object
77
        guideline_evaluator (GuidelineEvaluator): A guideline evaluator object
78
        resynthesize_query (bool): Whether to resynthesize query
79
        max_retries (int): Maximum number of retries
80
        callback_manager (Optional[CallbackManager]): A callback manager object
81
    """
82

83
    def __init__(
84
        self,
85
        query_engine: BaseQueryEngine,
86
        guideline_evaluator: GuidelineEvaluator,
87
        resynthesize_query: bool = False,
88
        max_retries: int = 3,
89
        callback_manager: Optional[CallbackManager] = None,
90
        query_transformer: Optional[FeedbackQueryTransformation] = None,
91
    ) -> None:
92
        self._query_engine = query_engine
93
        self._guideline_evaluator = guideline_evaluator
94
        self.max_retries = max_retries
95
        self.resynthesize_query = resynthesize_query
96
        self.query_transformer = query_transformer or FeedbackQueryTransformation(
97
            resynthesize_query=self.resynthesize_query
98
        )
99
        super().__init__(callback_manager)
100

101
    def _get_prompt_modules(self) -> PromptMixinType:
102
        """Get prompt sub-modules."""
103
        return {
104
            "query_engine": self._query_engine,
105
            "guideline_evalator": self._guideline_evaluator,
106
        }
107

108
    def _query(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
109
        """Answer a query."""
110
        response = self._query_engine._query(query_bundle)
111
        if self.max_retries <= 0:
112
            return response
113
        typed_response = (
114
            response if isinstance(response, Response) else response.get_response()
115
        )
116
        query_str = query_bundle.query_str
117
        eval = self._guideline_evaluator.evaluate_response(query_str, typed_response)
118
        if eval.passing:
119
            logger.debug("Evaluation returned True.")
120
            return response
121
        else:
122
            logger.debug("Evaluation returned False.")
123
            new_query_engine = RetryGuidelineQueryEngine(
124
                self._query_engine,
125
                self._guideline_evaluator,
126
                self.resynthesize_query,
127
                self.max_retries - 1,
128
                self.callback_manager,
129
            )
130
            new_query = self.query_transformer.run(query_bundle, {"evaluation": eval})
131
            logger.debug("New query: %s", new_query.query_str)
132
            return new_query_engine.query(new_query)
133

134
    async def _aquery(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
135
        """Not supported."""
136
        return self._query(query_bundle)
137