update response extraction

duckdb-nsql/eval/prompt_formatters.py

@@ -65,77 +65,16 @@ class RajkumarFormatter:
 
     @classmethod
     def format_model_output(cls, output_sql: str, prompt: str) -> str:
-        def clean_code_block(block: str) -> str:
-            """Clean a code block by removing markdown syntax and extra whitespace."""
-            # Remove markdown indicators and common SQL prefixes
-            cleaned = (block
-                       .replace('```sql\n', '')
-                       .replace('```duckdb\n', '')
-                       .replace('```\n', '')
-                       .replace('```', '')
-                       .strip())
-
-            return cleaned
-
-        def ensure_semicolon(sql: str) -> str:
-            """Ensure the SQL query ends with exactly one semicolon."""
-            sql = sql.strip()
-            # Remove any existing trailing semicolons
-            while sql.endswith(';'):
-                sql = sql[:-1].strip()
-            # Add back exactly one semicolon
-            return sql + ";"
-
-        # First, try to find SQL-specific code blocks
-        sql_blocks = []
-        start_pos = 0
-        while True:
-            start = output_sql.find('```sql', start_pos)
-            if start == -1:
-                start = output_sql.find('```duckdb', start_pos)
-            if start == -1:
-                break
-
-            end = output_sql.find('```', start + 4)
-            if end == -1:
-                break
-
-            sql_blocks.append(output_sql[start:end+3])
-            start_pos = end + 3
-
-        # If SQL blocks found, use the last one
-        if sql_blocks:
-            return ensure_semicolon(clean_code_block(sql_blocks[-1])).replace('```sql','').replace('```','').strip()
-
-        # If no SQL blocks, look for generic code blocks
-        generic_blocks = []
-        start_pos = 0
-        while True:
-            start = output_sql.find('```', start_pos)
-            if start == -1:
-                break
-
-            end = output_sql.find('```', start + 3)
-            if end == -1:
-                break
-
-            block = output_sql[start:end+3]
-            # Skip if this is actually an SQL block (we already handled those)
-            if not block.startswith('```sql') and not block.startswith('```duckdb'):
-                generic_blocks.append(block)
-            start_pos = end + 3
-
-        # If generic blocks found, use the last one
-        if generic_blocks:
-            return ensure_semicolon(clean_code_block(generic_blocks[-1])).replace('```sql','').replace('```','').strip()
-
-        # If no code blocks found at all, take everything up to first semicolon
-        semicolon_pos = output_sql.find(';')
-        if semicolon_pos != -1:
-            return ensure_semicolon(output_sql[:semicolon_pos].strip()).replace('```sql','').replace('```','').strip()
-
-        # If no semicolon found, use the entire text
-        return ensure_semicolon(output_sql.strip()).replace('```sql','').replace('```','').strip()
+        pattern = r"```(?:sql|mysql|duckdb)?\n?(.*?)```"
+        match = re.search(pattern, output_sql, re.DOTALL)
+        sql = match.group(1).strip() if match else output_sql.strip()
+        
+        # Handle edge case where regex captured empty content
+        if not sql:
+            sql = output_sql.strip()
+    
+        # Ensure single trailing semicolon 
+        return sql.rstrip(';') + ';'
 
     @classmethod
     def format_gold_output(cls, output_sql: str) -> str: