sql: do not hard code the LIMIT clause in the table_info section (#1563)

Seeing a lot of issues in Discord in which the LLM is not using the
correct LIMIT clause for different SQL dialects. ie, it's using `LIMIT`
for mssql instead of `TOP`, or instead of `ROWNUM` for Oracle, etc.
I think this could be due to us specifying the LIMIT statement in the
example rows portion of `table_info`. So the LLM is seeing the `LIMIT`
statement used in the prompt.
Since we can't specify each dialect's method here, I think it's fine to
just replace the `SELECT... LIMIT 3;` statement with `3 rows from
table_name table:`, and wrap everything in a block comment directly
following the `CREATE` statement. The Rajkumar et al paper wrapped the
example rows and `SELECT` statement in a block comment as well anyway.
Thoughts @fpingham?
This commit is contained in:
Jon Luo
2023-03-14 02:08:27 -04:00
committed by GitHub
parent 9ee2713272
commit 0a1b1806e9
4 changed files with 41 additions and 37 deletions

View File

@@ -126,12 +126,6 @@ class SQLDatabase:
# build the select command
command = select(table).limit(self._sample_rows_in_table_info)
# save the command in string format
select_star = (
f"SELECT * FROM '{table.name}' LIMIT "
f"{self._sample_rows_in_table_info}"
)
# save the columns in string format
columns_str = "\t".join([col.name for col in table.columns])
@@ -152,16 +146,18 @@ class SQLDatabase:
except ProgrammingError:
sample_rows_str = ""
# build final info for table
tables.append(
create_table
+ select_star
+ ";\n"
+ columns_str
+ "\n"
+ sample_rows_str
table_info = (
f"{create_table.rstrip()}\n"
f"/*\n"
f"{self._sample_rows_in_table_info} rows from {table.name} table:\n"
f"{columns_str}\n"
f"{sample_rows_str}\n"
f"*/"
)
# build final info for table
tables.append(table_info)
else:
tables.append(create_table)