Upload engine.py

engine.py

@@ -2,7 +2,7 @@ import json
 import os
 from functools import lru_cache
 from openai import OpenAI
-from datetime import datetime
+from datetime import datetime, date, timedelta
 import re
 
 # =========================
@@ -37,10 +37,19 @@ def load_metadata():
     }
 
 
-def resolve_operator(op, value):
-    # 🔴 Normalize operator input
-    op = op.lower().strip()
+# =========================
+# OPERATOR RESOLUTION (COMPLETE FIXED VERSION)
+# =========================
+
+def resolve_operator(op, value, field_type=None):
+    """
+    Resolve operator and format value based on data type
+    FIXED: Properly handles numeric types without quotes
+    """
+    # Normalize operator input
+    op = op.lower().strip().replace(" ", "_")
 
+    # Extended operator aliases for all your operators
     OPERATOR_ALIASES = {
         "=": "equals",
         "==": "equals",
@@ -52,11 +61,24 @@ def resolve_operator(op, value):
         ">=": "greater_or_equal",
         "<=": "less_or_equal",
         "greater than": "greater_than",
-        "less than": "less_than"
+        "less than": "less_than",
+        "greaterthan": "greater_than",
+        "lessthan": "less_than",
+        "greaterthanorequal": "greater_or_equal",
+        "lessthanorequal": "less_or_equal",
+        "does_not_contain": "not_contains",
+        "is_blank": "is_empty",
+        "is_not_blank": "is_not_empty",
+        "on": "equals",
+        "date_equals": "equals",
+        "date_between": "between",
+        "startswith": "starts_with",
+        "endswith": "ends_with"
     }
 
     op = OPERATOR_ALIASES.get(op, op)
 
+    # SQL operator mapping
     mapping = {
         "equals": "=",
         "not_equals": "!=",
@@ -65,10 +87,32 @@ def resolve_operator(op, value):
         "greater_or_equal": ">=",
         "less_or_equal": "<=",
         "contains": "LIKE",
+        "not_contains": "NOT LIKE",
         "starts_with": "LIKE",
         "ends_with": "LIKE",
         "in": "IN",
-        "not_in": "NOT IN"
+        "not_in": "NOT IN",
+        "is_empty": "IS NULL",
+        "is_not_empty": "IS NOT NULL",
+        "between": "BETWEEN",
+        "not_between": "NOT BETWEEN",
+        "before": "<",
+        "after": ">",
+        # Date relative operators
+        "today": "=",
+        "yesterday": "=",
+        "tomorrow": "=",
+        "this_week": "BETWEEN",
+        "last_week": "BETWEEN",
+        "next_week": "BETWEEN",
+        "this_month": "BETWEEN",
+        "last_month": "BETWEEN",
+        "next_month": "BETWEEN",
+        "this_quarter": "BETWEEN",
+        "last_quarter": "BETWEEN",
+        "next_quarter": "BETWEEN",
+        "this_year": "BETWEEN",
+        "last_year": "BETWEEN"
     }
 
     if op not in mapping:
@@ -76,50 +120,179 @@ def resolve_operator(op, value):
 
     sql_op = mapping[op]
 
-    # 🔴 Escape string values safely
+    # ✅ Determine if field is numeric
+    is_numeric = field_type in ['integer', 'decimal', 'float', 'number', 'int', 'bigint']
+    is_date = field_type in ['date', 'datetime', 'timestamp']
+    is_boolean = field_type in ['boolean', 'bool']
+    
+    # Escape string values safely
     def sql_escape(val):
+        if val is None:
+            return 'NULL'
         return str(val).replace("'", "''")
 
+    # Handle NULL operators
+    if op in ("is_empty", "is_not_empty"):
+        return sql_op, ""
+
+    # Handle date relative operators
+    if op in ("today", "yesterday", "tomorrow", "this_week", "last_week", "next_week",
+              "this_month", "last_month", "next_month", "this_quarter", "last_quarter",
+              "next_quarter", "this_year", "last_year"):
+        today = date.today()
+        
+        if op == "today":
+            return "=", f"'{today}'"
+        elif op == "yesterday":
+            return "=", f"'{today - timedelta(days=1)}'"
+        elif op == "tomorrow":
+            return "=", f"'{today + timedelta(days=1)}'"
+        elif op == "this_week":
+            start = today - timedelta(days=today.weekday())
+            end = start + timedelta(days=6)
+            return "BETWEEN", f"'{start}' AND '{end}'"
+        elif op == "this_month":
+            start = today.replace(day=1)
+            if today.month == 12:
+                end = today.replace(day=31)
+            else:
+                end = (today.replace(month=today.month+1, day=1) - timedelta(days=1))
+            return "BETWEEN", f"'{start}' AND '{end}'"
+        elif op == "this_year":
+            start = today.replace(month=1, day=1)
+            end = today.replace(month=12, day=31)
+            return "BETWEEN", f"'{start}' AND '{end}'"
+        # Add more as needed
+
+    # Handle LIKE operators
     if op == "contains":
         return sql_op, f"'%{sql_escape(value)}%'"
 
+    if op == "not_contains":
+        return sql_op, f"'%{sql_escape(value)}%'"
+
     if op == "starts_with":
         return sql_op, f"'{sql_escape(value)}%'"
 
     if op == "ends_with":
         return sql_op, f"'%{sql_escape(value)}'"
 
+    # Handle BETWEEN operator
+    if op in ("between", "not_between"):
+        if not isinstance(value, (list, tuple)) or len(value) != 2:
+            raise ValueError("BETWEEN operator requires array of 2 values")
+        
+        if is_numeric:
+            return sql_op, f"{value[0]} AND {value[1]}"
+        else:
+            return sql_op, f"'{sql_escape(value[0])}' AND '{sql_escape(value[1])}'"
+
+    # ✅ Handle IN operators with type checking
     if op in ("in", "not_in"):
         if not isinstance(value, list):
-            raise ValueError("IN operator requires list")
-        escaped = [f"'{sql_escape(v)}'" for v in value]
-        return sql_op, f"({','.join(escaped)})"
-
-    return sql_op, f"'{sql_escape(value)}'"
-
+            value = [value]
+        
+        if is_numeric:
+            escaped = [str(v) for v in value]  # ✅ No quotes for numbers
+        else:
+            escaped = [f"'{sql_escape(v)}'" for v in value]
+        
+        return sql_op, f"({', '.join(escaped)})"
+
+    # ✅ Handle regular comparison operators with type awareness
+    if is_numeric:
+        return sql_op, str(value)  # ✅ No quotes for numbers
+    elif is_boolean:
+        if isinstance(value, bool):
+            return sql_op, "1" if value else "0"
+        return sql_op, str(value)
+    elif is_date:
+        return sql_op, f"'{sql_escape(value)}'"
+    else:
+        return sql_op, f"'{sql_escape(value)}'"
 
 
 # =========================
-# JOIN RESOLUTION
+# JOIN RESOLUTION (FIXED)
 # =========================
 
 def resolve_join_path(start_table, end_table):
+    """
+    Find join path between two tables
+    FIXED: Handles your join_graph.json structure
+    """
     joins = load_metadata()["joins"]
-
-    for path in joins.values():
+    
+    # Try direct lookup with double underscore
+    key = f"{start_table}__{end_table}"
+    if key in joins:
+        return joins[key]
+    
+    # Try searching by start and end table
+    for path_key, path in joins.items():
         if path["start_table"] == start_table and path["end_table"] == end_table:
-            return path["steps"]
+            return path
 
     raise ValueError(
         f"No join path found from {start_table} to {end_table}"
     )
 
 
+def build_join_sql(base_table, join_path):
+    """
+    Build JOIN SQL from join path
+    FIXED: Properly handles multi-step joins with from_previous_step flag
+    """
+    steps = join_path["steps"]
+    sql = []
+    
+    # Sort steps by step number
+    sorted_steps = sorted(steps, key=lambda x: x.get("step", 0))
+    
+    for i, step in enumerate(sorted_steps):
+        alias = step["alias"]
+        table = step["table"]
+        join_type = step["join_type"].upper()
+        
+        # ✅ Determine the left side of the join
+        if i == 0:
+            # First join always references base table
+            left_ref = base_table
+        else:
+            # Subsequent joins: check from_previous_step flag
+            if step.get("from_previous_step", False):
+                left_ref = sorted_steps[i-1]["alias"]  # ✅ Use previous alias
+            else:
+                left_ref = base_table
+        
+        # Build basic join condition
+        join_condition = f"{left_ref}.{step['base_column']} = {alias}.{step['foreign_column']}"
+        
+        # ✅ Add extra conditions if present
+        if "extra_conditions" in step and step["extra_conditions"]:
+            for extra in step["extra_conditions"]:
+                condition = f"{alias}.{extra['column']} {extra['operator']} {extra['value']}"
+                join_condition += f" AND {condition}"
+        
+        sql.append(
+            f"{join_type} JOIN {table} {alias} ON {join_condition}"
+        )
+    
+    return "\n".join(sql)
+
+
+# =========================
+# FIELD RESOLUTION
+# =========================
+
 FIELD_ALIASES = {
-    "join_date": "start_date",
-    "joining_date": "start_date",
-    "joined": "start_date",
-    "hire_date": "start_date"
+    "join_date": "date_of_joining",
+    "joining_date": "date_of_joining",
+    "joined": "date_of_joining",
+    "hire_date": "date_of_joining",
+    "emp_code": "employee_code",
+    "emp_name": "full_name",
+    "dept": "department"
 }
 
 def resolve_field(field_name, module):
@@ -127,7 +300,7 @@ def resolve_field(field_name, module):
     fields = meta["fields"]
 
     # 🔹 Normalize field name
-    field_name = field_name.lower().strip()
+    field_name = field_name.lower().strip().replace(" ", "_")
     field_name = FIELD_ALIASES.get(field_name, field_name)
 
     # 🔹 Validate existence
@@ -151,21 +324,6 @@ def resolve_field(field_name, module):
     return field
 
 
-
-def build_join_sql(base_table, steps):
-    sql = []
-    prev_alias = base_table  # alias == table name
-
-    for step in steps:
-        alias = step["alias"]
-        sql.append(
-            f"{step['join_type'].upper()} JOIN {step['table']} {alias} "
-            f"ON {prev_alias}.{step['base_column']} = {alias}.{step['foreign_column']}"
-        )
-        prev_alias = alias
-
-    return "\n".join(sql)
-
 # =========================
 # JSON SAFETY
 # =========================
@@ -174,7 +332,12 @@ def safe_json_loads(text):
     try:
         return json.loads(text)
     except json.JSONDecodeError:
-        match = re.search(r"\{.*\}", text, re.S)
+        # Try to extract JSON from markdown
+        match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', text, re.DOTALL)
+        if match:
+            return json.loads(match.group(1))
+        
+        match = re.search(r"\{.*\}", text, re.DOTALL)
         if match:
             return json.loads(match.group())
         raise ValueError("LLM returned invalid JSON")
@@ -194,7 +357,7 @@ def parse_intent(question, retries=2):
         if (fields := [
             f for f in meta["fields"]
             if meta["fields"][f]["module"] == module
-        ])
+        ][:20])  # Limit to 20 fields per module for token efficiency
     ])
 
     prompt = f"""
@@ -230,7 +393,7 @@ User question:
     for attempt in range(retries):
         try:
             res = client.chat.completions.create(
-                model="gpt-4.1-mini",
+                model="gpt-4o-mini",
                 messages=[
                     {
                         "role": "system",
@@ -241,7 +404,7 @@ User question:
                 temperature=0
             )
 
-            content = res.choices[0].message.content
+            content = res.choices[0].message.content.strip()
             plan = safe_json_loads(content)
 
             # ✅ NORMALIZE + STABILIZE INTENT SHAPE
@@ -253,12 +416,12 @@ User question:
 
             return plan
 
-        except Exception:
+        except Exception as e:
             if attempt == retries - 1:
-                raise ValueError("LLM failed to return valid JSON")
+                raise ValueError(f"LLM failed to return valid JSON: {str(e)}")
 
 # =========================
-# SQL GENERATOR
+# SQL GENERATOR (FIXED)
 # =========================
 
 def build_sql(plan):
@@ -273,7 +436,7 @@ def build_sql(plan):
     base_table = meta["modules"][module]["base_table"]
 
     joins = []
-    joined_tables = set()
+    joined_tables = {base_table}  # ✅ Track all joined tables
     where_clauses = []
 
     # ---------- SELECT ----------
@@ -284,7 +447,7 @@ def build_sql(plan):
         for f in select_fields:
             field = resolve_field(f, module)
             select_columns.append(
-                f"{field['table']}.{field['column']}"
+                f"{field['table']}.{field['column']} AS {f}"
             )
         select_sql = ", ".join(select_columns)
     else:
@@ -296,47 +459,61 @@ def build_sql(plan):
 
         table = field["table"]
         column = field["column"]
+        field_type = field.get("type")  # ✅ Get field type
 
+        # Add join if needed
         if table != base_table and table not in joined_tables:
-            join_steps = resolve_join_path(base_table, table)
-            joins.append(build_join_sql(base_table, join_steps))
-            joined_tables.add(table)
-
-        sql_op, sql_value = resolve_operator(f["operator"], f["value"])
-        where_clauses.append(
-            f"{table}.{column} {sql_op} {sql_value}"
-        )
+            join_path = resolve_join_path(base_table, table)
+            joins.append(build_join_sql(base_table, join_path))
+            
+            # ✅ Track all tables in join path
+            for step in join_path["steps"]:
+                joined_tables.add(step["table"])
+
+        # ✅ Pass field_type to resolve_operator
+        sql_op, sql_value = resolve_operator(f["operator"], f["value"], field_type)
+        
+        if sql_value:  # Has value
+            where_clauses.append(f"{table}.{column} {sql_op} {sql_value}")
+        else:  # IS NULL / IS NOT NULL
+            where_clauses.append(f"{table}.{column} {sql_op}")
 
     # 🔴 FIX: safe WHERE clause
     where_sql = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""
 
     # ---------- FINAL SQL ----------
-    sql = f"""
-    SELECT {select_sql}
-    FROM {base_table}
-    {' '.join(joins)}
-    {where_sql}
-    LIMIT 100
-    """
+    sql_parts = [
+        f"SELECT {select_sql}",
+        f"FROM {base_table}"
+    ]
+    
+    if joins:
+        sql_parts.extend(joins)
+    
+    if where_sql:
+        sql_parts.append(where_sql)
+    
+    sql_parts.append("LIMIT 100")
+    
+    sql = "\n".join(sql_parts)
 
     return sql.strip()
 
 
-
-
 # =========================
 # VALIDATION
 # =========================
 
 def validate_sql(sql):
-    sql = sql.lower()
+    sql_lower = sql.lower()
 
-    if not sql.startswith("select"):
+    if not sql_lower.strip().startswith("select"):
         raise ValueError("Only SELECT allowed")
 
-    forbidden = ["drop", "delete", "update", "insert", "truncate"]
-    if any(x in sql for x in forbidden):
-        raise ValueError("Unsafe SQL")
+    forbidden = ["drop", "delete", "update", "insert", "truncate", "alter", "create"]
+    for keyword in forbidden:
+        if re.search(rf'\b{keyword}\b', sql_lower):
+            raise ValueError(f"Unsafe SQL: '{keyword}' not allowed")
 
     return sql
 
@@ -365,3 +542,26 @@ def run(question):
         "query_plan": plan,
         "sql": sql
     }
+
+
+# =========================
+# TEST
+# =========================
+
+if __name__ == "__main__":
+    test_queries = [
+        "Show all employees",
+        "Find departments with more than 50 employees",
+        "Show employees in departments 1, 2, 3",
+        "List employees who joined this month"
+    ]
+    
+    for q in test_queries:
+        print(f"\n{'='*80}")
+        print(f"Q: {q}")
+        print('='*80)
+        try:
+            result = run(q)
+            print("SQL:", result["sql"])
+        except Exception as e:
+            print("ERROR:", e)