Format table and types metadata into string to be used in prompt
(table_names: List[str] = [])
| 64 | |
| 65 | # TODO: refac this to access JSON fields instead of tables |
| 66 | def get_table_schemas_str(table_names: List[str] = []) -> str: |
| 67 | """ |
| 68 | Format table and types metadata into string to be used in prompt |
| 69 | """ |
| 70 | global ENUMS_METADATA_DICT |
| 71 | global TABLES_METADATA_DICT |
| 72 | |
| 73 | tables_to_use = [] |
| 74 | if table_names: |
| 75 | tables_to_use = [TABLES_METADATA_DICT[t_name] for t_name in table_names] |
| 76 | else: |
| 77 | tables_to_use = [t for t in TABLES_METADATA_DICT.values()] |
| 78 | |
| 79 | enums_to_use = set() |
| 80 | tables_str_list = [] |
| 81 | for table in tables_to_use: |
| 82 | tables_str = f"table name: {table['name']}\n" |
| 83 | if table.get("description"): |
| 84 | tables_str += f"table description: {table.get('description')}\n" |
| 85 | columns_str_list = [] |
| 86 | for column in table.get("columns", []): |
| 87 | columns_str_list.append(f"{column['name']} [{column['type']}]") |
| 88 | if column.get("type") in ENUMS_METADATA_DICT.keys(): |
| 89 | enums_to_use.add(column.get("type")) |
| 90 | tables_str += f"table columns: {', '.join(columns_str_list)}\n" |
| 91 | tables_str_list.append(tables_str) |
| 92 | tables_details = "\n\n".join(tables_str_list) |
| 93 | |
| 94 | enums_str_list = [] |
| 95 | for custom_type_str in enums_to_use: |
| 96 | custom_type = ENUMS_METADATA_DICT.get(custom_type_str) |
| 97 | if custom_type: |
| 98 | enums_str = f"enum: {custom_type['type']}\n" |
| 99 | enums_str += f"valid values: {', '.join(custom_type.get('valid_values'))}\n" |
| 100 | enums_str_list.append(enums_str) |
| 101 | enums_details = "\n\n".join(enums_str_list) |
| 102 | |
| 103 | return enums_details + "\n\n" + tables_details |
| 104 | |
| 105 | |
| 106 | def get_relevant_tables_from_pinecone(natural_language_query, index_name="text_to_sql") -> List[str]: |
no outgoing calls
no test coverage detected