File size: 7,224 Bytes
8a58cf3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
"""Subclasses from base prompt."""
from typing import List

from gpt_index.prompts.base import Prompt
from gpt_index.prompts.prompt_type import PromptType


class SummaryPrompt(Prompt):
    """Summary prompt.

    Prompt to summarize the provided `context_str`.

    Required template variables: `context_str`

    Args:
        template (str): Template for the prompt.
        **prompt_kwargs: Keyword arguments for the prompt.

    """

    prompt_type: PromptType = PromptType.SUMMARY
    input_variables: List[str] = ["context_str"]


class TreeInsertPrompt(Prompt):
    """Tree Insert prompt.

    Prompt to insert a new chunk of text `new_chunk_text` into the tree index.
    More specifically, this prompt has the LLM select the relevant candidate
    child node to continue tree traversal.

    Required template variables: `num_chunks`, `context_list`, `new_chunk_text`

    Args:
        template (str): Template for the prompt.
        **prompt_kwargs: Keyword arguments for the prompt.

    """

    prompt_type: PromptType = PromptType.TREE_INSERT
    input_variables: List[str] = ["num_chunks", "context_list", "new_chunk_text"]


class TreeSelectPrompt(Prompt):
    """Tree select prompt.

    Prompt to select a candidate child node out of all child nodes
    provided in `context_list`, given a query `query_str`. `num_chunks` is
    the number of child nodes in `context_list`.

    Required template variables: `num_chunks`, `context_list`, `query_str`

    Args:
        template (str): Template for the prompt.
        **prompt_kwargs: Keyword arguments for the prompt.

    """

    prompt_type: PromptType = PromptType.TREE_SELECT
    input_variables: List[str] = ["num_chunks", "context_list", "query_str"]


class TreeSelectMultiplePrompt(Prompt):
    """Tree select multiple prompt.

    Prompt to select multiple candidate child nodes out of all
    child nodes provided in `context_list`, given a query `query_str`.
    `branching_factor` refers to the number of child nodes to select, and
    `num_chunks` is the number of child nodes in `context_list`.

    Required template variables: `num_chunks`, `context_list`, `query_str`,
        `branching_factor`

    Args:
        template (str): Template for the prompt.
        **prompt_kwargs: Keyword arguments for the prompt.

    """

    prompt_type = PromptType.TREE_SELECT_MULTIPLE
    input_variables: List[str] = [
        "num_chunks",
        "context_list",
        "query_str",
        "branching_factor",
    ]


class RefinePrompt(Prompt):
    """Refine prompt.

    Prompt to refine an existing answer `existing_answer` given a context `context_msg`,
    and a query `query_str`.

    Required template variables: `query_str`, `existing_answer`, `context_msg`

    Args:
        template (str): Template for the prompt.
        **prompt_kwargs: Keyword arguments for the prompt.

    """

    # TODO: rename context_msg to context_str

    prompt_type: PromptType = PromptType.REFINE
    input_variables: List[str] = ["query_str", "existing_answer", "context_msg"]


class QuestionAnswerPrompt(Prompt):
    """Question Answer prompt.

    Prompt to answer a question `query_str` given a context `context_str`.

    Required template variables: `context_str`, `query_str`

    Args:
        template (str): Template for the prompt.
        **prompt_kwargs: Keyword arguments for the prompt.

    """

    prompt_type: PromptType = PromptType.QUESTION_ANSWER
    input_variables: List[str] = ["context_str", "query_str"]


class KeywordExtractPrompt(Prompt):
    """Keyword extract prompt.

    Prompt to extract keywords from a text `text` with a maximum of
    `max_keywords` keywords.

    Required template variables: `text`, `max_keywords`

    Args:
        template (str): Template for the prompt.
        **prompt_kwargs: Keyword arguments for the prompt.

    """

    prompt_type: PromptType = PromptType.KEYWORD_EXTRACT
    input_variables: List[str] = ["text", "max_keywords"]


class QueryKeywordExtractPrompt(Prompt):
    """Query keyword extract prompt.

    Prompt to extract keywords from a query `query_str` with a maximum
    of `max_keywords` keywords.

    Required template variables: `query_str`, `max_keywords`

    Args:
        template (str): Template for the prompt.
        **prompt_kwargs: Keyword arguments for the prompt.

    """

    prompt_type: PromptType = PromptType.QUERY_KEYWORD_EXTRACT
    input_variables: List[str] = ["question", "max_keywords"]


class SchemaExtractPrompt(Prompt):
    """Schema extract prompt.

    Prompt to extract schema from unstructured text `text`.

    Required template variables: `text`, `schema`

    Args:
        template (str): Template for the prompt.
        **prompt_kwargs: Keyword arguments for the prompt.

    """

    prompt_type: PromptType = PromptType.SCHEMA_EXTRACT
    input_variables: List[str] = ["text", "schema"]


class TextToSQLPrompt(Prompt):
    """Text to SQL prompt.

    Prompt to translate a natural language query into SQL,
    given a schema `schema`.

    Required template variables: `query_str`, `schema`

    Args:
        template (str): Template for the prompt.
        **prompt_kwargs: Keyword arguments for the prompt.

    """

    prompt_type: PromptType = PromptType.TEXT_TO_SQL
    input_variables: List[str] = ["query_str", "schema"]


class TableContextPrompt(Prompt):
    """Table context prompt.

    Prompt to generate a table context given a table schema `schema`,
    as well as unstructured text context `context_str`, and
    a task `query_str`.
    This includes both a high-level description of the table
    as well as a description of each column in the table.

    Args:
        template (str): Template for the prompt.
        **prompt_kwargs: Keyword arguments for the prompt.

    """

    prompt_type: PromptType = PromptType.TABLE_CONTEXT
    input_variables: List[str] = ["schema", "context_str", "query_str"]


class RefineTableContextPrompt(Prompt):
    """Refine Table context prompt.

    Prompt to refine a table context given a table schema `schema`,
    as well as unstructured text context `context_msg`, and
    a task `query_str`.
    This includes both a high-level description of the table
    as well as a description of each column in the table.

    Args:
        template (str): Template for the prompt.
        **prompt_kwargs: Keyword arguments for the prompt.

    """

    # TODO: rename context_msg to context_str

    prompt_type: PromptType = PromptType.TABLE_CONTEXT
    input_variables: List[str] = [
        "schema",
        "context_msg",
        "query_str",
        "existing_answer",
    ]


class KnowledgeGraphPrompt(Prompt):
    """Define the knowledge graph triplet extraction prompt."""

    prompt_type: PromptType = PromptType.KNOWLEDGE_TRIPLET_EXTRACT
    input_variables: List[str] = ["max_knowledge_triplets", "text"]


class SimpleInputPrompt(Prompt):
    """Simple Input prompt.

    Required template variables: `query_str`.

    Args:
        template (str): Template for the prompt.
        **prompt_kwargs: Keyword arguments for the prompt.

    """

    prompt_type: PromptType = PromptType.SIMPLE_INPUT
    input_variables: List[str] = ["query_str"]