server/grader.py

"""
Deterministic grader for the SQL Query Environment.

Compares the agent's query result against the expected result.
Produces a score between 0.0 and 1.0 with partial credit:

  0.0  — SQL error or no result returned
  0.1  — Query runs but returns wrong column count
  0.2  — Correct column count but wrong column names
  0.3  — Correct columns but wrong number of rows
  0.5  — Correct structure, some rows match
  0.7  — Most rows match (>= 70% of expected rows found)
  0.9  — All rows match but in wrong order
  1.0  — Exact match (columns, rows, order)
"""

from tasks import TaskDefinition


def _normalize_value(val) -> str:
    """Normalize a value for comparison (handle float rounding, casing)."""
    if val is None:
        return "NULL"
    if isinstance(val, float):
        return f"{val:.2f}"
    return str(val).strip().lower()


def _normalize_row(row: tuple) -> tuple:
    return tuple(_normalize_value(v) for v in row)


def grade_query(
    task: TaskDefinition,
    agent_rows: list[tuple] | None,
    agent_columns: list[str] | None,
    query_error: str | None,
) -> tuple[float, str]:
    """
    Grade the agent's SQL query result against the expected result.

    Returns:
        (score, feedback_string)
    """
    # SQL error → 0.0
    if query_error is not None:
        return 0.0, f"SQL error: {query_error}"

    # No result (e.g. INSERT/UPDATE instead of SELECT) → 0.0
    if agent_rows is None or agent_columns is None:
        return 0.0, "Query did not return any result rows. Expected a SELECT query."

    expected_cols = [c.lower() for c in task.expected_columns]
    actual_cols = [c.lower() for c in agent_columns]

    # Wrong column count → 0.1
    if len(actual_cols) != len(expected_cols):
        return 0.1, (
            f"Wrong number of columns. Expected {len(expected_cols)} "
            f"({', '.join(expected_cols)}), got {len(actual_cols)} "
            f"({', '.join(actual_cols)})."
        )

    # Wrong column names → 0.2
    if actual_cols != expected_cols:
        return 0.2, (
            f"Column names don't match. Expected {expected_cols}, got {actual_cols}. "
            f"Note: column names must match exactly."
        )

    # Columns match — now compare rows
    expected_normalized = [_normalize_row(r) for r in task.expected_rows]
    actual_normalized = [_normalize_row(r) for r in agent_rows]

    # Wrong row count → partial credit based on matching rows
    if len(actual_normalized) != len(expected_normalized):
        # Count how many expected rows appear in actual
        expected_set = set(expected_normalized)
        actual_set = set(actual_normalized)
        matching = len(expected_set & actual_set)

        if matching == 0:
            return 0.3, (
                f"Correct columns but wrong number of rows. "
                f"Expected {len(expected_normalized)} rows, got {len(actual_normalized)}. "
                f"No matching rows found."
            )
        ratio = matching / len(expected_normalized)
        score = 0.3 + ratio * 0.4  # Scale from 0.3 to 0.7
        return round(score, 2), (
            f"Correct columns but wrong row count. "
            f"Expected {len(expected_normalized)} rows, got {len(actual_normalized)}. "
            f"{matching}/{len(expected_normalized)} expected rows found."
        )

    # Same row count — check content
    expected_set = set(expected_normalized)
    actual_set = set(actual_normalized)

    if expected_set != actual_set:
        # Some rows match
        matching = len(expected_set & actual_set)
        if matching == 0:
            return 0.5, (
                "Correct structure (columns and row count) but row values "
                "don't match any expected rows."
            )
        ratio = matching / len(expected_normalized)
        score = 0.5 + ratio * 0.2
        return round(score, 2), (
            f"Partially correct. {matching}/{len(expected_normalized)} rows match."
        )

    # All rows present — check order
    if actual_normalized != expected_normalized:
        return 0.9, (
            "All correct rows found but in wrong order. "
            "Check your ORDER BY clause."
        )

    # Perfect match!
    return 1.0, "Correct! Query returned the exact expected result."