ZennyKenny commited on
Commit
61d9b40
·
verified ·
1 Parent(s): 35cddc5

try to fix output

Browse files
Files changed (1) hide show
  1. app.py +22 -2
app.py CHANGED
@@ -43,16 +43,36 @@ def query_sql(user_query: str) -> str:
43
  Returns:
44
  The query result from the database as a formatted string.
45
  """
46
- # Provide the AI with the correct schema
47
  schema_info = (
48
  "The database has a table named 'receipts' with the following schema:\n"
49
  "- receipt_id (INTEGER, primary key)\n"
50
  "- customer_name (VARCHAR(16))\n"
51
  "- price (FLOAT)\n"
52
  "- tip (FLOAT)\n"
53
- "Generate an SQL query using ONLY these column names.\n"
 
54
  )
55
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
56
  # Generate SQL query using the provided schema
57
  generated_sql = agent.run(f"{schema_info} Convert this request into SQL: {user_query}")
58
 
 
43
  Returns:
44
  The query result from the database as a formatted string.
45
  """
46
+ # Provide the AI with the correct schema and strict instructions
47
  schema_info = (
48
  "The database has a table named 'receipts' with the following schema:\n"
49
  "- receipt_id (INTEGER, primary key)\n"
50
  "- customer_name (VARCHAR(16))\n"
51
  "- price (FLOAT)\n"
52
  "- tip (FLOAT)\n"
53
+ "Generate a valid SQL SELECT query using ONLY these column names.\n"
54
+ "DO NOT explain your reasoning, and DO NOT return anything other than the SQL query itself."
55
  )
56
 
57
+ # Generate SQL query using the provided schema
58
+ generated_sql = agent.run(f"{schema_info} Convert this request into SQL: {user_query}")
59
+
60
+ # Log the generated SQL for debugging
61
+ print(f"Generated SQL: {generated_sql}")
62
+
63
+ # Ensure we only execute valid SELECT queries
64
+ if not generated_sql.strip().lower().startswith(("select", "show", "pragma")):
65
+ return "Error: Only SELECT queries are allowed."
66
+
67
+ # Execute the SQL query and return the result
68
+ result = sql_engine(generated_sql)
69
+
70
+ # Log the SQL query result
71
+ print(f"SQL Query Result: {result}")
72
+
73
+ return result # Return the final result, NOT the generated SQL
74
+
75
+
76
  # Generate SQL query using the provided schema
77
  generated_sql = agent.run(f"{schema_info} Convert this request into SQL: {user_query}")
78