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 import BaseEvaluator
8
from llama_index.legacy.indices.list.base import SummaryIndex
9
from llama_index.legacy.prompts.mixin import PromptMixinType
10
from llama_index.legacy.query_engine.retriever_query_engine import RetrieverQueryEngine
11
from llama_index.legacy.schema import Document, QueryBundle
12
from llama_index.legacy.service_context import ServiceContext
13

14
logger = logging.getLogger(__name__)
15

16

17
class RetrySourceQueryEngine(BaseQueryEngine):
18
    """Retry with different source nodes."""
19

20
    def __init__(
21
        self,
22
        query_engine: RetrieverQueryEngine,
23
        evaluator: BaseEvaluator,
24
        service_context: Optional[ServiceContext] = None,
25
        max_retries: int = 3,
26
        callback_manager: Optional[CallbackManager] = None,
27
    ) -> None:
28
        """Run a BaseQueryEngine with retries."""
29
        self._query_engine = query_engine
30
        self._evaluator = evaluator
31
        self._service_context = service_context
32
        self.max_retries = max_retries
33
        super().__init__(callback_manager)
34

35
    def _get_prompt_modules(self) -> PromptMixinType:
36
        """Get prompt sub-modules."""
37
        return {"query_engine": self._query_engine, "evaluator": self._evaluator}
38

39
    def _query(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
40
        response = self._query_engine._query(query_bundle)
41
        if self.max_retries <= 0:
42
            return response
43
        typed_response = (
44
            response if isinstance(response, Response) else response.get_response()
45
        )
46
        query_str = query_bundle.query_str
47
        eval = self._evaluator.evaluate_response(query_str, typed_response)
48
        if eval.passing:
49
            logger.debug("Evaluation returned True.")
50
            return response
51
        else:
52
            logger.debug("Evaluation returned False.")
53
            # Test source nodes
54
            source_evals = [
55
                self._evaluator.evaluate(
56
                    query=query_str,
57
                    response=typed_response.response,
58
                    contexts=[source_node.get_content()],
59
                )
60
                for source_node in typed_response.source_nodes
61
            ]
62
            orig_nodes = typed_response.source_nodes
63
            assert len(source_evals) == len(orig_nodes)
64
            new_docs = []
65
            for node, eval_result in zip(orig_nodes, source_evals):
66
                if eval_result:
67
                    new_docs.append(Document(text=node.node.get_content()))
68
            if len(new_docs) == 0:
69
                raise ValueError("No source nodes passed evaluation.")
70
            new_index = SummaryIndex.from_documents(
71
                new_docs,
72
                service_context=self._service_context,
73
            )
74
            new_retriever_engine = RetrieverQueryEngine(new_index.as_retriever())
75
            new_query_engine = RetrySourceQueryEngine(
76
                new_retriever_engine,
77
                self._evaluator,
78
                self._service_context,
79
                self.max_retries - 1,
80
            )
81
            return new_query_engine.query(query_bundle)
82

83
    async def _aquery(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
84
        """Not supported."""
85
        return self._query(query_bundle)
86