fix humaneval first-line indent recovery

src/benchkit/benchmarks/humaneval.py

@@ -1,6 +1,7 @@
 """HumanEval benchmark - 164 code generation tasks, scored by pass@1."""
 
 import json
+import textwrap
 from pathlib import Path
 
 from benchkit.benchmarks.base import Task
@@ -18,26 +19,52 @@
 
 
 def _extract_code(response: str) -> str:
-    """Strip markdown fences if the model wrapped its output."""
+    """Strip reasoning and a single Markdown code fence."""
     text = strip_think_tags(response).rstrip()
 
     if "```python" in text:
         text = text.split("```python", 1)[1].split("```", 1)[0]
     elif "```" in text:
         text = text.split("```", 1)[1].split("```", 1)[0]
 
-    text = text.strip("\n")
-    lines = text.split("\n")
-
-    if any(ln.startswith(("def ", "class ", "import ", "from ", "@")) for ln in lines):
-        return text
-
-    # Otherwise it's a bare body -> indent only if not already indented.
-    first = next((ln for ln in lines if ln.strip()), "")
-    if first and not first.startswith((" ", "\t")):
-        return "\n".join(("    " + ln) if ln.strip() else ln for ln in lines)
-
-    return text
+    return text.strip("\n")
+
+
+def _assemble_solution(prompt: str, entry: str, code: str) -> str:
+    """Join a completion onto the HumanEval prompt.
+
+    Bare bodies are indented as a function body. If the first line is flush
+    and the rest is already indented, indenting every line over-indents the
+    body, so recover by indenting only that first line.
+    """
+    if f"def {entry}" in code:
+        imports = [
+            line for line in prompt.split("\n") if line.startswith(("import ", "from "))
+        ]
+        return "\n".join(imports) + "\n\n" + code if imports else code
+
+    normalized = textwrap.dedent(code)
+    solution = prompt + textwrap.indent(normalized, "    ")
+    try:
+        compile(solution, "<benchkit-humaneval>", "exec")
+        return solution
+    except (IndentationError, TabError):
+        pass
+
+    lines = code.splitlines()
+    first_index = next(
+        (index for index, line in enumerate(lines) if line.strip()), None
+    )
+    if first_index is not None and not lines[first_index].startswith((" ", "\t")):
+        lines[first_index] = "    " + lines[first_index]
+        normalized = textwrap.dedent("\n".join(lines))
+        repaired = prompt + textwrap.indent(normalized, "    ")
+        try:
+            compile(repaired, "<benchkit-humaneval>", "exec")
+        except (IndentationError, TabError):
+            return solution
+        return repaired
+    return solution
 
 
 class HumanEval:
@@ -70,18 +97,7 @@ def evaluate(self, task: Task, response: str) -> bool:
     def evaluate_with_feedback(self, task: Task, response: str) -> EvaluationResult:
         code = _extract_code(response)
         entry = task.metadata["entry_point"]
-
-        if f"def {entry}" in code:
-            # Model gave full function - prepend any imports from the prompt
-            imports = [
-                line
-                for line in task.prompt.split("\n")
-                if line.startswith(("import ", "from "))
-            ]
-            fn_code = "\n".join(imports) + "\n\n" + code if imports else code
-        else:
-            fn_code = task.prompt + code
-
+        fn_code = _assemble_solution(task.prompt, entry, code)
         full = fn_code + "\n\n" + task.metadata["test"] + f"\ncheck({entry})\n"
         result = execute_with_feedback(full)
         return EvaluationResult(float(result.passed), result.feedback)

tests/test_humaneval.py

@@ -0,0 +1,67 @@
+"""Regression tests for HumanEval response assembly."""
+
+from benchkit.benchmarks.humaneval import HumanEval, _assemble_solution, _extract_code
+
+CLOSE_ELEMENTS_FLUSH = """\
+sorted_numbers = sorted(numbers)
+for i in range(len(sorted_numbers) - 1):
+    if sorted_numbers[i + 1] - sorted_numbers[i] < threshold:
+        return True
+return False"""
+
+CLOSE_ELEMENTS_INDENTED = """\
+    sorted_numbers = sorted(numbers)
+    for i in range(len(sorted_numbers) - 1):
+        if sorted_numbers[i + 1] - sorted_numbers[i] < threshold:
+            return True
+    return False"""
+
+CLOSE_ELEMENTS_FIRST_LINE_FLUSH = """\
+sorted_numbers = sorted(numbers)
+    for i in range(len(sorted_numbers) - 1):
+        if sorted_numbers[i + 1] - sorted_numbers[i] < threshold:
+            return True
+    return False"""
+
+
+def _task(task_id: str = "HumanEval/0"):
+    for task in HumanEval().load_tasks():
+        if task.id == task_id:
+            return task
+    raise AssertionError(f"missing {task_id}")
+
+
+def test_body_recovers_first_line_only_indent_loss() -> None:
+    task = _task()
+    solution = _assemble_solution(
+        task.prompt, task.metadata["entry_point"], CLOSE_ELEMENTS_FIRST_LINE_FLUSH
+    )
+    compile(solution, "<test-humaneval-body>", "exec")
+    assert HumanEval().evaluate(task, CLOSE_ELEMENTS_FIRST_LINE_FLUSH)
+
+
+def test_flush_and_indented_bodies_still_pass() -> None:
+    task = _task()
+    bench = HumanEval()
+    assert bench.evaluate(task, CLOSE_ELEMENTS_FLUSH)
+    assert bench.evaluate(task, CLOSE_ELEMENTS_INDENTED)
+    assert bench.evaluate(task, f"```python\n{CLOSE_ELEMENTS_FIRST_LINE_FLUSH}\n```")
+
+
+def test_full_function_completion_still_passes() -> None:
+    task = _task()
+    response = (
+        "def has_close_elements(numbers: list[float], threshold: float) -> bool:\n"
+        + CLOSE_ELEMENTS_INDENTED
+    )
+    assert HumanEval().evaluate(task, response)
+
+
+def test_wrong_body_still_fails() -> None:
+    task = _task()
+    assert not HumanEval().evaluate(task, "    return False")
+
+
+def test_extract_code_strips_fences_without_reindenting() -> None:
+    extracted = _extract_code(f"```python\n{CLOSE_ELEMENTS_FIRST_LINE_FLUSH}\n```")
+    assert extracted == CLOSE_ELEMENTS_FIRST_LINE_FLUSH