code/local_eval/live_protocol.py

"""Local copies of the live Aug 24 pre-eval / scoring protocol.

Retry unusable turns the way ``remote.worker._generate_retrying_bad_turns`` and
``sanity_service.dispatcher`` do. Detect submits the way the dispatcher does:
the marker in the first bash command, not a byte-exact clause.
"""

from __future__ import annotations

from dataclasses import replace

from albedo_eval_service.remote.dataset import EvalSample, format_messages
from albedo_eval_service.remote.generation import GenerationResult
from albedo_eval_service.shared.observation_format import (
    MAX_CONSECUTIVE_BAD_TURNS,
    retry_feedback,
    unusable_turn,
)
from albedo_eval_service.shared.submit_protocol import first_bash_command, is_exact_submission
from albedo_eval_service.simulator.prompt_simulator import COMPLETE_MARKER

from .constants import TOKENIZER_DIR

__all__ = [
    "MAX_CONSECUTIVE_BAD_TURNS",
    "generate_retrying_bad_turns",
    "is_live_submit",
]


def is_live_submit(text: str, *, command: str = "", marker: str = "") -> bool:
    """True when the first bash block is a submit, matching live pre-eval."""
    cmd = first_bash_command(text)
    if marker and marker in cmd:
        return True
    if command and is_exact_submission(text, command):
        return True
    if not command and not marker:
        return COMPLETE_MARKER in text
    return False


def generate_retrying_bad_turns(generator, samples: list[EvalSample]) -> list[GenerationResult]:
    """Generate one turn; re-ask empty / truncated / no-command up to 3 times."""
    results = generator.generate(samples)
    by_id = {sample.sample_id: sample for sample in samples}
    for _ in range(MAX_CONSECUTIVE_BAD_TURNS - 1):
        redo = [
            (by_id[r.sample_id], unusable_turn(r.text, truncated=r.truncated))
            for r in results
            if not r.error and r.sample_id in by_id and unusable_turn(r.text, truncated=r.truncated)
        ]
        if not redo:
            break
        retry_samples = [_with_retry_feedback(sample, reason) for sample, reason in redo]
        fresh = {r.sample_id: r for r in generator.generate(retry_samples)}
        results = [fresh.get(r.sample_id, r) for r in results]
    return results


def _with_retry_feedback(sample: EvalSample, reason: str) -> EvalSample:
    messages = _base_messages(sample) + [{"role": "user", "content": retry_feedback(reason)}]
    return replace(
        sample,
        prompt=format_messages(
            messages, tokenizer_path=str(TOKENIZER_DIR), enable_thinking=True
        ),
        messages=messages,
    )


def _base_messages(sample: EvalSample) -> list[dict[str, str]]:
    return list(sample.messages or [{"role": "user", "content": sample.prompt}])