Upload inference.py

inference.py

@@ -195,7 +195,7 @@ Generate your next action (SQL or action name):"""
                 return text.upper()
             else:
                 # Try to extract SQL from response
-                sql_match = re.search(r'(SELECT.*?;?)', text, re.IGNORECASE | re.DOTALL)
+                sql_match = re.search(r'((?:SELECT|INSERT|UPDATE|DELETE|CREATE|ALTER|DROP)[\s\S]*?(?:;|$))', text, re.IGNORECASE | re.DOTALL)
                 if sql_match:
                     return sql_match.group(1).strip()
                 return 'QUERY'  # Default fallback
@@ -234,6 +234,10 @@ Generate your next action (SQL or action name):"""
                     action_type = 'submit'
                     sql = None
                     column_name = None
+                elif action_text.upper() == 'FIX' or action_text.upper().startswith('FIX'):
+                    action_type = 'fix'
+                    sql = action_text
+                    column_name = None
                 elif action_text.upper().startswith('PROFILE'):
                     action_type = 'profile'
                     sql = None
@@ -272,6 +276,8 @@ Generate your next action (SQL or action name):"""
                 except Exception as e:
                     error_msg = str(e)
                     current_action = f"{action_type}:{sql or column_name or 'N/A'}"
+                    steps_taken = step
+                    rewards.append(0.0)
                     log_step(step=step, action=current_action, reward=0.0, done=False, error=error_msg)
                     break
             
@@ -293,14 +299,6 @@ def main():
         print("Please set HF_TOKEN or ensure .env file exists with the token", flush=True)
         return
     
-    if not API_BASE_URL:
-        print("Error: API_BASE_URL environment variable is required", flush=True)
-        return
-    
-    if not MODEL_NAME:
-        print("Error: MODEL_NAME environment variable is required", flush=True)
-        return
-    
     # Initialize agent
     agent = SQLQualityAgent(API_BASE_URL, MODEL_NAME, HF_TOKEN, ENV_BASE_URL)