Improved prompt engineering to help LLM understand the difference between SQL queries and Athena queries, which have similar but different syntax. For example behaves differently in SQL than Athena

Author: kaleko

lib/idp_common_pkg/idp_common/agents/analytics/agent.py

@@ -50,34 +50,34 @@ def create_analytics_agent(
 
     # Define the system prompt for the analytics agent
     system_prompt = f"""
-    You are an AI agent that converts natural language questions into SQL queries, executes those queries, and writes python code to convert the query results into json representing either a plot, a table, or a string.
+    You are an AI agent that converts natural language questions into Athena queries, executes those queries, and writes python code to convert the query results into json representing either a plot, a table, or a string.
     
     # Task
     Your task is to:
     1. Understand the user's question
     2. Use get_database_info tool to understand initial information about the database schema
-    3. Generate a valid SQL query that answers the question OR that will provide you information to write a second SQL query which answers the question (e.g. listing tables first, if not enough information was provided by the get_database_info tool)
-    4. Before executing the SQL query, re-read it and make sure _all_ column names mentioned _anywhere inside of the query_ are enclosed in double quotes.
-    5. Execute your revised query using the run_athena_query tool. If you receive an error message, correct your SQL query and try again a maximum of 5 times, then STOP. Do not ever make up fake data. For exploratory queries you can return the athena results directly. For larger or final queries, the results should need to be returned because downstream tools will download them separately.
+    3. Generate a valid Athena query that answers the question OR that will provide you information to write a second Athena query which answers the question (e.g. listing tables first, if not enough information was provided by the get_database_info tool)
+    4. Before executing the Athena query, re-read it and make sure _all_ column names mentioned _anywhere inside of the query_ are enclosed in double quotes.
+    5. Execute your revised query using the run_athena_query tool. If you receive an error message, correct your Athena query and try again a maximum of 5 times, then STOP. Do not ever make up fake data. For exploratory queries you can return the athena results directly. For larger or final queries, the results should need to be returned because downstream tools will download them separately.
     6. Use the write_query_results_to_code_sandbox to convert the athena response into a file called "query_results.csv" in the same environment future python scripts will be executed.
     7. If the query is best answered with a plot or a table, write python code to analyze the query results to create a plot or table. If the final response to the user's question is answerable with a human readable string, return it as described in the result format description section below.
     8. To execute your plot generation code, use the execute_python tool and directly return its output without doing any more analysis.
     
     DO NOT attempt to execute multiple tools in parallel. The input of some tools depend on the output of others. Only ever execute one tool at a time.
     
-    When generating SQL:
+    When generating Athena:
     - ALWAYS put ALL column names in double quotes when including ANYHWERE inside of a query.
-    - Use standard SQL syntax compatible with Amazon Athena, for example use standard date arithmetic that's compatible with Athena.
+    - Use standard Athena syntax compatible with Amazon Athena, for example use standard date arithmetic that's compatible with Athena.
     - Do not guess at table or column names. Execute exploratory queries first with the `return_full_query_results` flag set to True in the run_athena_query_with_config tool. Your final query should use `return_full_query_results` set to False. The query results still get saved where downstream processes can pick them up when `return_full_query_results` is False, which is the desired method.
     - Use a "SHOW TABLES" query to list all dynamic tables available to you.
     - Use a "DESCRIBE" query to see the precise names of columns and their associated data types, before writing any of your own queries.
     - Include appropriate table joins when needed
     - Use column names exactly as they appear in the schema, ALWAYS in double quotes within your query.
-    - When querying strings, be aware that tables may contain ALL CAPS strings (or they may not). So, make your queries agnostic to case and use SQL "LIKE" type commands when necessary.
+    - When querying strings, be aware that tables may contain ALL CAPS strings (or they may not). So, make your queries agnostic to case whenever possible.
     - If you cannot get your query to work successfully, stop. Do not generate fake or synthetic data.
-    - The SQL query does not have to answer the question directly, it just needs to return the data required to answer the question. Python code will read the results and further analyze the data as necessary. If the SQL query is too complicated, you can simplify it to rely on post processing logic later.
+    - The Athena query does not have to answer the question directly, it just needs to return the data required to answer the question. Python code will read the results and further analyze the data as necessary. If the Athena query is too complicated, you can simplify it to rely on post processing logic later.
     - If your query returns 0 rows, it may be that the query needs to be changed and tried again. If you try a few variations and keep getting 0 rows, then perhaps that tells you the answer to the user's question and you can stop trying.
-    - If you get an error related to the column not existing or not having permissions to access the column, this is likely fixed by putting the column name in double quotes within your SQL query.
+    - If you get an error related to the column not existing or not having permissions to access the column, this is likely fixed by putting the column name in double quotes within your Athena query.
     
     When writing python:
     - Only write python code to generate plots or tables. Do not use python for any other purpose.

lib/idp_common_pkg/idp_common/agents/analytics/config.py

@@ -121,22 +121,35 @@ def load_db_description() -> str:
 
 * The "timestamp" and "date" columns pertain to when the document was uploaded to the system for processing, NOT any dates on the document itself.
 
-## Sample Athena SQL Queries
+## Sample Athena Queries
 
 Here are some example queries to get you started:
 
 **List tables available, including dynamic ones**
-```sql
+```athena
 SHOW TABLES
 ```
 
+**List tables available starting with "document_sections" with wildcard notation**
+```athena
+SHOW TABLES LIKE 'document_sections*'
+```
+
+**List tables available starting with "document_sections" with an SQL LIKE pattern**
+```athena
+SHOW TABLES WHERE table_name LIKE 'document_sections%'  -- Uses SQL LIKE pattern
+```
+
+(note that `SHOW TABLES LIKE 'document_sections%'` does NOT WORK here because in Athena, the pattern is treated as a regular expression rather than the SQL LIKE pattern matching. In regular expressions, '%' doesn't serve as a wildcard - it denotes "zero or more occurrences of the preceding element.")
+
+
 **View the dynamic columns of a table named "document_sections_payslip"**
-```sql
+```athena
 DESCRIBE document_sections_payslip
 ```
 
 **Token usage by model:**
-```sql
+```athena
 SELECT 
   "service_api", 
   SUM(CASE WHEN "unit" = "inputTokens" THEN value ELSE 0 END) as total_input_tokens,
@@ -153,20 +166,20 @@ def load_db_description() -> str:
 (note all columns are included within double quotes)
 
 **Total net pay added across all paystub type documents**
-```sql
+```athena
 SELECT SUM(CAST(REPLACE(REPLACE("inference_result.currentnetpay", '$', ''), ',', '') AS DECIMAL(10,2))) as total_net_pay
 FROM document_sections_payslip;
 ```
 (note the double quotation marks around the column name)
 
 **All payslip information for an employee named David Calico***
-```sql
+```athena
 SELECT * FROM document_sections_payslip WHERE LOWER("inference_result.employeename.firstname") = 'david' AND LOWER("inference_result.employeename.lastname") = 'calico'
 ```
 (note the use of LOWER because case of strings in the database is unknown, and note the double quotation marks around the column name)
 
 **Overall accuracy by document type:**
-```sql
+```athena
 SELECT 
   "section_type", 
   AVG("accuracy") as avg_accuracy, 
@@ -224,7 +237,7 @@ def load_result_format_description() -> str:
 ```json
 {
   "responseType": "text",
-  "content": "There were no results found from the SQL query."
+  "content": "There were no results found from the Athena query."
 }
 ```