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, scope="USA", session_id=None)
| 35 | ) |
| 36 | |
| 37 | def text_to_sql_with_retry(natural_language_query, table_names, k=3, messages=None, scope="USA", session_id=None): |
| 38 | """ |
| 39 | Tries to take a natural language query and generate valid SQL to answer it K times |
| 40 | """ |
| 41 | if scope == "SF": |
| 42 | model = "gpt-3.5-turbo" |
| 43 | else: |
| 44 | model = "gpt-3.5-turbo" |
| 45 | |
| 46 | example_messages = [] |
| 47 | enums_message = [{'role': 'user', 'content': ''}] |
| 48 | schema_message = [{'role': 'user', 'content': ''}] |
| 49 | message_history = [] |
| 50 | |
| 51 | if not messages: |
| 52 | table_text, enum_text = get_table_and_enums(table_names, scope) |
| 53 | |
| 54 | schema_message[0]['content'] = table_text |
| 55 | enums_message[0]['content'] = enum_text |
| 56 | |
| 57 | enc = tiktoken.encoding_for_model(model) |
| 58 | |
| 59 | instruction_length = len(enc.encode(table_text + '\n\n' + enum_text)) |
| 60 | |
| 61 | content = get_retry_prompt(DIALECT, natural_language_query, scope) |
| 62 | |
| 63 | max_messages = -1 |
| 64 | if instruction_length > 1000: |
| 65 | max_messages = 3 |
| 66 | elif instruction_length > 1500: |
| 67 | max_messages = 2 |
| 68 | elif instruction_length > 2000: |
| 69 | max_messages = 1 |
| 70 | |
| 71 | example_messages = make_default_messages('', scope, n=max_messages) |
| 72 | message_history.append({ |
| 73 | "role": "user", |
| 74 | "content": content |
| 75 | }) |
| 76 | |
| 77 | assistant_message = None |
| 78 | sql_query = "" |
| 79 | for attempt_number in range(k): |
| 80 | sql_query_data = {} |
| 81 | try: |
| 82 | purpose = "text_to_sql" if attempt_number == 0 else "text_to_sql_retry" |
| 83 | try: |
| 84 | payload = schema_message + message_history |
| 85 | if (attempt_number == 0): |
| 86 | payload = example_messages + enums_message + payload |
| 87 | assistant_message = get_assistant_message_from_openai(payload, model=model, scope=scope, purpose=purpose, session_id=session_id) |
| 88 | except: |
| 89 | continue |
| 90 | |
| 91 | sql_query_data = extract_sql_query_from_message(assistant_message["message"]["content"]) |
| 92 | |
| 93 | if sql_query_data.get('MissingData'): |
| 94 | return {"MissingData": sql_query_data['MissingData']}, "" |
no test coverage detected