Tries to take a natural language query and generate valid SQL to answer it K times
(natural_language_query, table_names, k=3, messages=None)
| 131 | |
| 132 | |
| 133 | def text_to_sql_with_retry(natural_language_query, table_names, k=3, messages=None): |
| 134 | """ |
| 135 | Tries to take a natural language query and generate valid SQL to answer it K times |
| 136 | """ |
| 137 | if not messages: |
| 138 | # ask the assistant to rephrase before generating the query |
| 139 | schemas_str = get_table_schemas_str(table_names) |
| 140 | # rephrase = [{ |
| 141 | # "role": "user", |
| 142 | # "content": make_rephrase_msg_with_schema_and_warnings().format( |
| 143 | # natural_language_query=natural_language_query, |
| 144 | # schemas_str=schemas_str |
| 145 | # ) |
| 146 | # }] |
| 147 | # rephrased_query = get_assistant_message(rephrase)["message"]["content"] |
| 148 | # natural_language_query = rephrased_query |
| 149 | |
| 150 | content = make_msg_with_schema_and_warnings().format( |
| 151 | natural_language_query=natural_language_query, |
| 152 | schemas_str=schemas_str |
| 153 | ) |
| 154 | |
| 155 | messages = make_default_messages(schemas_str) |
| 156 | messages.append({ |
| 157 | "role": "user", |
| 158 | "content": content |
| 159 | }) |
| 160 | |
| 161 | assistant_message = None |
| 162 | |
| 163 | for _ in range(k): |
| 164 | try: |
| 165 | # model = "gpt-4" |
| 166 | # model = "gpt-3.5-turbo" |
| 167 | model = "gpt-3.5-turbo-0301" |
| 168 | assistant_message = get_assistant_message(messages, model=model) |
| 169 | sql_query = extract_sql_query_from_message(assistant_message["message"]["content"]) |
| 170 | |
| 171 | response = execute_sql(sql_query) |
| 172 | # Generated SQL query did not produce exception. Return result |
| 173 | return response, sql_query |
| 174 | |
| 175 | except Exception as e: |
| 176 | messages.append({ |
| 177 | "role": "assistant", |
| 178 | "content": assistant_message["message"]["content"] |
| 179 | }) |
| 180 | messages.append({ |
| 181 | "role": "user", |
| 182 | "content": MSG_WITH_ERROR_TRY_AGAIN.format(error_message=str(e)) |
| 183 | }) |
| 184 | |
| 185 | print("Could not generate SQL query after {k} tries.".format(k=k)) |
| 186 | return None, None |
no test coverage detected