Update app.py

app.py

@@ -86,9 +86,34 @@ def generate_sql(question: str, schema: str) -> str:
     table_match = re.search(r'CREATE TABLE\s+"?(\w+)"?', schema, re.IGNORECASE)
     table_name = table_match.group(1) if table_match else "data"
     quoted = f'"{table_name}"'
-
-    # Extract column names to inject into prompt — helps T5-small stay grounded
     col_match = re.findall(r'"(\w+)"', schema)
+
+    # ── Rule-based shortcuts (fast + accurate) ────────────────────────────────
+    q = question.lower().strip()
+    if re.search(r'show.*(first|top).*\d+|first.*\d+.*row|top.*\d+', q):
+        n = re.search(r'\d+', q)
+        return f'SELECT * FROM {quoted} LIMIT {n.group() if n else 10}'
+    if re.search(r'(show|display|get|give).*(first|all).*row|first.*row|show.*row', q):
+        return f'SELECT * FROM {quoted} LIMIT 10'
+    if re.search(r'count.*(total|all|record|row)|total.*(record|row|count)|how many', q):
+        return f'SELECT COUNT(*) FROM {quoted}'
+    if re.search(r'show.*(all|every).*row|all.*row|select all', q):
+        return f'SELECT * FROM {quoted} LIMIT 50'
+    if re.search(r'average|avg', q) and col_match:
+        num_col = next((c for c in col_match if re.search(r'num|price|val|amt|count|qty|sal|rev|cost|pm|aqi|no|co|so|o3', c, re.I)), col_match[1] if len(col_match) > 1 else col_match[0])
+        return f'SELECT AVG("{num_col}") FROM {quoted}'
+    if re.search(r'unique|distinct', q) and col_match:
+        return f'SELECT COUNT(DISTINCT "{col_match[0]}") FROM {quoted}'
+    if re.search(r'group by', q) and col_match:
+        return f'SELECT "{col_match[0]}", COUNT(*) FROM {quoted} GROUP BY "{col_match[0]}"'
+    if re.search(r'max|maximum|highest', q) and col_match:
+        num_col = col_match[1] if len(col_match) > 1 else col_match[0]
+        return f'SELECT MAX("{num_col}") FROM {quoted}'
+    if re.search(r'min|minimum|lowest', q) and col_match:
+        num_col = col_match[1] if len(col_match) > 1 else col_match[0]
+        return f'SELECT MIN("{num_col}") FROM {quoted}'
+
+    # ── T5 model fallback ─────────────────────────────────────────────────────
     col_hint = ", ".join(col_match) if col_match else ""
     prompt = f"tables:\n{schema}\ncolumns: {col_hint}\nquery for: {question}"
     inputs = tokenizer(
@@ -197,4 +222,3 @@ async def query(req: QueryRequest):
 @app.get("/health")
 def health():
     return {"status": "ok", "model": MODEL_NAME, "device": DEVICE}
-