llama-index

1
"""Answer inserter."""
2

3
from abc import abstractmethod
4
from typing import Any, Dict, List, Optional
5

6
from llama_index.legacy.prompts.base import BasePromptTemplate, PromptTemplate
7
from llama_index.legacy.prompts.mixin import (
8
    PromptDictType,
9
    PromptMixin,
10
    PromptMixinType,
11
)
12
from llama_index.legacy.query_engine.flare.schema import QueryTask
13
from llama_index.legacy.service_context import ServiceContext
14

15

16
class BaseLookaheadAnswerInserter(PromptMixin):
17
    """Lookahead answer inserter.
18

19
    These are responsible for insert answers into a lookahead answer template.
20

21
    E.g.
22
    lookahead answer: Red is for [Search(What is the meaning of Ghana's
23
        flag being red?)], green for forests, and gold for mineral wealth.
24
    query: What is the meaning of Ghana's flag being red?
25
    query answer: "the blood of those who died in the country's struggle
26
        for independence"
27
    final answer: Red is for the blood of those who died in the country's
28
        struggle for independence, green for forests, and gold for mineral wealth.
29

30
    """
31

32
    def _get_prompt_modules(self) -> PromptMixinType:
33
        """Get prompt sub-modules."""
34
        return {}
35

36
    @abstractmethod
37
    def insert(
38
        self,
39
        response: str,
40
        query_tasks: List[QueryTask],
41
        answers: List[str],
42
        prev_response: Optional[str] = None,
43
    ) -> str:
44
        """Insert answers into response."""
45

46

47
DEFAULT_ANSWER_INSERT_PROMPT_TMPL = """
48
An existing 'lookahead response' is given below. The lookahead response
49
contains `[Search(query)]` tags. Some queries have been executed and the
50
response retrieved. The queries and answers are also given below.
51
Also the previous response (the response before the lookahead response)
52
is given below.
53
Given the lookahead template, previous response, and also queries and answers,
54
please 'fill in' the lookahead template with the appropriate answers.
55

56
NOTE: Please make sure that the final response grammatically follows
57
the previous response + lookahead template. For example, if the previous
58
response is "New York City has a population of " and the lookahead
59
template is "[Search(What is the population of New York City?)]", then
60
the final response should be "8.4 million".
61

62
NOTE: the lookahead template may not be a complete sentence and may
63
contain trailing/leading commas, etc. Please preserve the original
64
formatting of the lookahead template if possible.
65

66
NOTE:
67

68
NOTE: the exception to the above rule is if the answer to a query
69
is equivalent to "I don't know" or "I don't have an answer". In this case,
70
modify the lookahead template to indicate that the answer is not known.
71

72
NOTE: the lookahead template may contain multiple `[Search(query)]` tags
73
    and only a subset of these queries have been executed.
74
    Do not replace the `[Search(query)]` tags that have not been executed.
75

76
Previous Response:
77

78

79
Lookahead Template:
80
Red is for [Search(What is the meaning of Ghana's \
81
    flag being red?)], green for forests, and gold for mineral wealth.
82

83
Query-Answer Pairs:
84
Query: What is the meaning of Ghana's flag being red?
85
Answer: The red represents the blood of those who died in the country's struggle \
86
    for independence
87

88
Filled in Answers:
89
Red is for the blood of those who died in the country's struggle for independence, \
90
    green for forests, and gold for mineral wealth.
91

92
Previous Response:
93
One of the largest cities in the world
94

95
Lookahead Template:
96
, the city contains a population of [Search(What is the population \
97
    of New York City?)]
98

99
Query-Answer Pairs:
100
Query: What is the population of New York City?
101
Answer: The population of New York City is 8.4 million
102

103
Synthesized Response:
104
, the city contains a population of 8.4 million
105

106
Previous Response:
107
the city contains a population of
108

109
Lookahead Template:
110
[Search(What is the population of New York City?)]
111

112
Query-Answer Pairs:
113
Query: What is the population of New York City?
114
Answer: The population of New York City is 8.4 million
115

116
Synthesized Response:
117
8.4 million
118

119
Previous Response:
120
{prev_response}
121

122
Lookahead Template:
123
{lookahead_response}
124

125
Query-Answer Pairs:
126
{query_answer_pairs}
127

128
Synthesized Response:
129
"""
130
DEFAULT_ANSWER_INSERT_PROMPT = PromptTemplate(DEFAULT_ANSWER_INSERT_PROMPT_TMPL)
131

132

133
class LLMLookaheadAnswerInserter(BaseLookaheadAnswerInserter):
134
    """LLM Lookahead answer inserter.
135

136
    Takes in a lookahead response and a list of query tasks, and the
137
        lookahead answers, and inserts the answers into the lookahead response.
138

139
    Args:
140
        service_context (ServiceContext): Service context.
141

142
    """
143

144
    def __init__(
145
        self,
146
        service_context: Optional[ServiceContext] = None,
147
        answer_insert_prompt: Optional[BasePromptTemplate] = None,
148
    ) -> None:
149
        """Init params."""
150
        self._service_context = service_context or ServiceContext.from_defaults()
151
        self._answer_insert_prompt = (
152
            answer_insert_prompt or DEFAULT_ANSWER_INSERT_PROMPT
153
        )
154

155
    def _get_prompts(self) -> Dict[str, Any]:
156
        """Get prompts."""
157
        return {
158
            "answer_insert_prompt": self._answer_insert_prompt,
159
        }
160

161
    def _update_prompts(self, prompts: PromptDictType) -> None:
162
        """Update prompts."""
163
        if "answer_insert_prompt" in prompts:
164
            self._answer_insert_prompt = prompts["answer_insert_prompt"]
165

166
    def insert(
167
        self,
168
        response: str,
169
        query_tasks: List[QueryTask],
170
        answers: List[str],
171
        prev_response: Optional[str] = None,
172
    ) -> str:
173
        """Insert answers into response."""
174
        prev_response = prev_response or ""
175

176
        query_answer_pairs = ""
177
        for query_task, answer in zip(query_tasks, answers):
178
            query_answer_pairs += f"Query: {query_task.query_str}\nAnswer: {answer}\n"
179

180
        return self._service_context.llm.predict(
181
            self._answer_insert_prompt,
182
            lookahead_response=response,
183
            query_answer_pairs=query_answer_pairs,
184
            prev_response=prev_response,
185
        )
186

187

188
class DirectLookaheadAnswerInserter(BaseLookaheadAnswerInserter):
189
    """Direct lookahead answer inserter.
190

191
    Simple inserter module that directly inserts answers into
192
        the [Search(query)] tags in the lookahead response.
193

194
    Args:
195
        service_context (ServiceContext): Service context.
196

197
    """
198

199
    def _get_prompts(self) -> Dict[str, Any]:
200
        """Get prompts."""
201
        return {}
202

203
    def _update_prompts(self, prompts: PromptDictType) -> None:
204
        """Update prompts."""
205

206
    def insert(
207
        self,
208
        response: str,
209
        query_tasks: List[QueryTask],
210
        answers: List[str],
211
        prev_response: Optional[str] = None,
212
    ) -> str:
213
        """Insert answers into response."""
214
        for query_task, answer in zip(query_tasks, answers):
215
            response = (
216
                response[: query_task.start_idx]
217
                + answer
218
                + response[query_task.end_idx + 1 :]
219
            )
220
        return response
221