[Python] Remapping SQL agent's query response to the original table/column names

Introduction


I'm currently building an AI assistant that has SQL agent where it's idea is simply generate sql quereies based on the user's question. using Python and it's libraries.

The problem was at start was that the database has above 150 tables, where each the table/column names isn't clear fot the sql llm as it has some arabic words written in english, french words also written in english, unusuall names like companyname rather than company_name, or some typos.

I've solved this problem by

1. embedding phase
   1. creating a json file that contain mapping of the old names with better, clearer ones.
   2. replaced hard words with easier ones to make the llm understand, as the scheme is now simpler
   3. using another LLM: generate descriptions on each table/column
   4. create victors using an embedder to simplify search
2. retrival phase
   1. RAG is used on the table description, with keyword search on the column's.
   2. the results of the search is passed to 2 llm calls, one to choose the table, another is used to choose the columns related, based also on descriptions
   3. after the selection, the related tables is passed to the sql agent which will generate the sql query
   4. after the sql generation the table names should return to the original names, also for the column names

The main problem


the implementation of the remapper, which it is supposed to revert table/column names to it's original values of the sql agent query response. The remap function should handle any generated sql code simple or complex ones.

More info


I'm using sqlcoder:7b to generate the queries, it's okay if i changed the remap function but i'm not sure what is the best way to do it, i tried sqlparse once but with no luck as it couldn't handle easy cases where there isn't an alias

Output


The output should be a query that has tables/column names reverted to the original ones

I've implemented everything but the remapper as the sql agent could generate complex queries or it might not use alias or a single way to write queries.

I tried something but it didn't work:

Remap function


def remap_sql_using_selected_tables(self, sql_query: str, selected_tables: dict) -> str:
"""
Remap friendly table/column names to original database names using selected tables.
Handles aliases, qualified columns, unqualified columns (if unambiguous), and CTEs.

Args:
sql_query (str): SQL using friendly names.
selected_tables (dict): Selected tables from previous steps.

Returns:
str: SQL with friendly names replaced by original names.
"""
#return it to the form that is accepted within the function, as it accepts columns in list form without any information
for tbl in selected_tables.keys():
list_of_columns = list(selected_tables[tbl]['table_columns'].keys())
selected_tables[tbl]['table_columns'] = list_of_columns

#Build reverse mapping only for selected tables
reverse_mapping = {
tbl: mapping for tbl, mapping in self.original_table_column_mapping.items()
if tbl in selected_tables
}

#Normalize SQL and aliases (remove quotes for consistency)
sql_query_clean = sql_query.replace('"', '').replace('`', '')

#Detect aliases in FROM, JOIN, and CTEs
alias_pattern = re.compile(r'\b(FROM|JOIN)\s+([^\s,()]+)(?:\s+(?:AS\s+)?(\w+))?', re.IGNORECASE)
aliases = {} # alias -> friendly table
for match in alias_pattern.finditer(sql_query_clean):
_, table_or_subquery, alias = match.groups()
table_clean = table_or_subquery.strip()
if alias:
aliases[alias.strip()] = table_clean
else:
aliases[table_clean] = table_clean

#Detect CTEs
cte_pattern = re.compile(r'\bWITH\s+(\w+)\s+AS\s*\(', re.IGNORECASE)
for match in cte_pattern.finditer(sql_query_clean):
cte_name = match.group(1)
aliases[cte_name] = cte_name

#Replace table names first (FROM, JOIN)
for friendly_table, mapping in reverse_mapping.items():
sql_query_clean = re.sub(rf'\b{re.escape(friendly_table)}\b',
mapping["original"],
sql_query_clean,
flags=re.IGNORECASE)

# Replace qualified columns (alias.column or table.column)
for alias, friendly_table in aliases.items():
if friendly_table not in reverse_mapping:
continue
col_map = reverse_mapping[friendly_table]["columns"]
for friendly_col, orig_col in col_map.items():
# Replace alias.column
sql_query_clean = re.sub(rf'\b{re.escape(alias)}\.{re.escape(friendly_col)}\b',
f'{alias}.{orig_col}',
sql_query_clean,
flags=re.IGNORECASE)
# Replace table.column
sql_query_clean = re.sub(rf'\b{re.escape(reverse_mapping[friendly_table]["original"])}\.{re.escape(friendly_col)}\b',
f'{reverse_mapping[friendly_table]["original"]}.{orig_col}',
sql_query_clean,
flags=re.IGNORECASE)

#Replace unqualified columns if unambiguous
all_columns = {}
for table, mapping in reverse_mapping.items():
for friendly_col, orig_col in mapping["columns"].items():
all_columns.setdefault(friendly_col, set()).add(orig_col)

for friendly_col, orig_set in all_columns.items():
if len(orig_set) == 1:
sql_query_clean = re.sub(rf'(?<!\.)\b{re.escape(friendly_col)}\b',
list(orig_set)[0],
sql_query_clean,
flags=re.IGNORECASE)

return sql_query_clean

Example of selected table variable


selected_tables = {'employees':
{'table_description': 'table description where it talks about the content and what it is used for',
'table_columns': {'company_id': '(bigint), nullable',
'supervisor_id': '(bigint), nullable'}}
}

Example of self.original_table_column_mapping


# the "original" key contains the table original name
# the keys of "columns" contain the columns original names
self.original_table_column_mapping = {"employees": {
"original": "users_table",
"columns": {
"company_id": "companies_id",
"supervisor_id": "supervisorid",
}
}
}

Output example


Generated SQL before remapping:

SELECT COUNT(*) AS total_employees FROM employees e JOIN user_company uc ON e.company_id = uc.company_id WHERE uc.company_id = 1234567;


Generated SQL after remapping:

SELECT COUNT(*) AS total_employees FROM users_table e JOIN company uc ON e.companies_id = uc.id WHERE uc.id = 1234567;


also it should work in cases of no alias or more complex ones, as it does not work in these cases.

Continue reading...