多轮：一道题要来回好几轮

import logging
import random
import re
from typing import Any, Dict, List, Optional, Union

import evals
import evals.metrics
from evals.api import CompletionFn
from evals.elsuite.twenty_questions.utils import PROMPTS, generate_task_state_for
from evals.eval import SolverEval
from evals.record import Recorder
from evals.registry import registry
from evals.solvers.human_cli_solver import HumanCliSolver
from evals.solvers.solver import Solver
from evals.solvers.utils import maybe_wrap_with_solver
from evals.task_state import Message

logger = logging.getLogger(__name__)
WORD_PATTERN = r"\[GUESS (.*?)\]"


class TwentyQuestions(SolverEval):
    def __init__(
        self,
        completion_fns: List[CompletionFn],
        samples_jsonl: str,
        gamemaster_spec: str,
        max_questions: int = 20,
        max_replies: int = 40,
        num_shortlist_items: int = 20,
        shortlist_variant: bool = False,
        seed: int = 222024,
        n_samples: Optional[int] = None,
        *args,
        **kwargs,
    ):
        super().__init__(completion_fns, seed=seed, *args, **kwargs)

        self.samples_jsonl = samples_jsonl
        self.gamemaster_solver = maybe_wrap_with_solver(
            registry.make_completion_fn(gamemaster_spec)
        )
        self.max_questions = max_questions

        if max_replies < max_questions:
            logger.warn(
                f"max_replies ({max_replies}) is less than max_questions ({max_questions}). Setting max_replies to {max_questions + 20}"
            )
        self.max_replies = max_replies if max_replies > max_questions else max_questions + 20
        self.num_shortlist_items = num_shortlist_items
        self.shortlist_variant = shortlist_variant

        self.n_samples = n_samples
        self.rng = random.Random(seed)

    def eval_sample(self, solver: Solver, sample: Dict, rng: random.Random) -> Dict[str, Any]:
        assert "word" in sample, "Sample must contain 'word' field"
        assert "difficulty" in sample, "Sample must contain 'difficulty' field"

        if not isinstance(solver, HumanCliSolver):
            logging.info(f"Running sample: {sample['word']}")

        # Generate the shortlist for the current sample if applicable.
        if self.shortlist_variant:
            assert self.num_shortlist_items <= len(
                self.shortlist
            ), "Number of shortlist items must be less than or equal to the total number of samples."
            shortlist_for_sample = rng.sample(self.shortlist, self.num_shortlist_items)
            if sample["word"] not in shortlist_for_sample:
                random_index = rng.randint(0, len(shortlist_for_sample) - 1)
                shortlist_for_sample[random_index] = sample["word"]
        else:
            shortlist_for_sample = None
        response = self._conversation_loop(solver, sample, shortlist_for_sample)

        return response

    def run(self, recorder: Recorder) -> Dict[str, Union[float, int]]:
        samples = self.get_samples()
        self.rng.shuffle(samples)
        samples = samples[: self.n_samples] if self.n_samples else samples

        if self.shortlist_variant:
            self.shortlist = [sample["word"] for sample in samples]

        self.eval_all_samples(recorder, samples)
        events = recorder.get_events("match")

        scores = [event.data["score"] for event in events]
        num_guesses = [event.data["num_guesses"] for event in events]
        num_questions = [event.data["num_questions"] for event in events]
        num_violations = [event.data["num_violations"] for event in events]
        num_gamemaster_refusals = [event.data["num_gamemaster_refusals"] for event in events]
        incorrect_guesses = [event.data["incorrect_guesses"] for event in events]
        word_difficulties = [event.data["word_difficulty"] for event in events]

        return {
            "score": sum(scores) / len(scores),
            "accuracy": evals.metrics.get_accuracy(events),
            "bootstrap_std": evals.metrics.get_bootstrap_accuracy_std(events),
            "average_num_guesses": sum(num_guesses) / len(num_guesses),
            "average_num_questions": sum(num_questions) / len(num_questions),
            "average_num_violations": sum(num_violations) / len(num_violations),
            "average_num_gamemaster_refusals": sum(num_gamemaster_refusals)
            / len(num_gamemaster_refusals),
            "average_num_incorrect_guesses": sum((len(ig) for ig in incorrect_guesses))
            / len(incorrect_guesses),
            "average_word_difficulty": sum(word_difficulties) / len(word_difficulties),
        }

    def _conversation_loop(
        self, solver: Solver, sample: Dict, shortlist: Optional[List[str]] = None
    ) -> Dict[str, Any]:
        """Maintains a conversation between the guesser and the gamemaster until the maximum number of questions is reached, or until a correct guess is made.

        Args:
            solver (Solver): any compatible solver, instantiated for the current sample.
            sample (Dict): current sample – one word to guess, and its associated difficulty.

        Returns:
            Dict[str, Any]: a dictionary containing the final result and metrics of the conversation.
        """

        metrics = {
            "num_guesses": 0,
            "num_questions": 0,
            "num_violations": 0,
            "num_guesser_replies": 0,  # num_guesses + num_questions + num_violations
            "num_gamemaster_refusals": 0,
            "incorrect_guesses": [],
        }
        conversation = []

        # Contains fall-back condition to avoid infinite loops for solvers which never output questions.
        while (
            metrics["num_questions"] < self.max_questions
            and metrics["num_guesser_replies"] < self.max_replies
        ):
            task_state = generate_task_state_for(
                "guesser", conversation, max_questions=self.max_questions, shortlist=shortlist
            )
            guesser_response = solver(task_state)
            conversation += [Message(content=guesser_response.output, role="guesser")]
            metrics["num_guesser_replies"] += 1

            # Check if guess made:
            match = re.search(WORD_PATTERN, guesser_response.output)
            if match is not None:
                metrics["num_guesses"] += 1
                guess = match.group(1)
                if guess.lower() == sample["word"].lower():
                    response = {
                        "correct": True,
                        "score": self.max_questions - metrics["num_questions"],
                        "expected": sample["word"],
                        "word_difficulty": sample["difficulty"],
                        "picked": guess,
                        "num_guesses": metrics["num_guesses"],
                        "num_questions": metrics["num_questions"],
                        "num_violations": metrics["num_violations"],
                        "num_gamemaster_refusals": metrics["num_gamemaster_refusals"],
                        "incorrect_guesses": metrics["incorrect_guesses"],
                    }
                    evals.record.record_match(**response)
                    return response
                else:
                    metrics["incorrect_guesses"] += [guess]
                    conversation += [
                        Message(
                            content=PROMPTS["incorrect_guess"].format(guess=guess), role="system"
                        )
                    ]
                    continue
            elif "?" in guesser_response.output.strip():
                metrics["num_questions"] += 1
            else:  # Neither guess nor question.
                # TODO: Maybe make the guesser retry here?
                logger.warn(
                    f"Rule violation, no guess or question in output: {guesser_response.output}"
                )
                metrics["num_violations"] += 1
                conversation += [Message(content=PROMPTS["rule_violation"], role="system")]
                continue

            task_state = generate_task_state_for("gamemaster", conversation, sample["word"])
            gamemaster_response = self.gamemaster_solver(task_state)
            conversation += [Message(content=gamemaster_response.output, role="gamemaster")]
            if gamemaster_response.output.lower() == "skip":
                metrics["num_gamemaster_refusals"] += 1

        logger.info(f"Ran out of questions for word: {sample['word']}")
        response = {
            "correct": False,
            "score": 0,
            "expected": sample["word"],
            "word_difficulty": sample["difficulty"],
            "num_guesses": metrics["num_guesses"],
            "num_questions": metrics["num_questions"],
            "num_violations": metrics["num_violations"],
            "num_gamemaster_refusals": metrics["num_gamemaster_refusals"],
            "incorrect_guesses": metrics["incorrect_guesses"],
        }
        evals.record.record_match(**response)
        return response

from dataclasses import dataclass
from typing import List

from evals.task_state import Message, TaskState


@dataclass
class Interaction:
    #   All messages we've seen (except for the task_description)
    messages: List[Message]

    #   IDs of the CoT private internal messages
    private_messages_ids: List[int]


class PersistentMemoryCache:
    def __init__(
        self,
        interaction_length: int,
    ):
        self.private_interaction_length = interaction_length
        self.last_interaction = None

    def save_private_interaction(self, task_state: TaskState):
        #   Save the interaction
        interaction_messages = task_state.messages
        num_interaction_messages = len(interaction_messages)
        private_messages_ids = (
            [] if self.last_interaction is None else self.last_interaction.private_messages_ids
        )
        private_messages_ids += list(
            range(
                num_interaction_messages - self.private_interaction_length - 1,
                num_interaction_messages - 1,
            )
        )
        self.last_interaction = Interaction(interaction_messages, private_messages_ids)

    def load_private_interaction(self, task_state: TaskState) -> List[Message]:
        if self.last_interaction is None:
            return task_state.messages

        #   Check if task_state matches our last interaction
        interaction = self.last_interaction
        task_state_message_ix = 0
        for our_message_ix in range(0, len(interaction.messages)):
            if our_message_ix in interaction.private_messages_ids:
                continue
            else:
                if (
                    task_state.messages[task_state_message_ix]
                    != interaction.messages[our_message_ix]
                ):
                    raise ValueError(
                        (
                            f"task_state message {task_state_message_ix} different than the corresponding message "
                            "in the interaction history.\n"
                            f"task_state.messages:\n{task_state.messages}\n"
                            f"interaction.messages:\n{interaction.messages}\n"
                        )
                    )
                task_state_message_ix += 1

        return interaction.messages + task_state.messages[task_state_message_ix:]

import copy
import logging
import re
from dataclasses import dataclass
from typing import Any, Optional

from evals.elsuite.bugged_tools.tools import Tool, ToolTaskState
from evals.solvers.solver import Solver, SolverResult
from evals.task_state import Message, TaskState

logger = logging.getLogger(__name__)


@dataclass
class ToolCall:
    tool_name: str
    input: str
    output: Any


@dataclass
class ParsedSolverResult:
    tool_calls: list[ToolCall]
    final_answer: Optional[str]


@dataclass
class RunnerResult:
    final_task_state: ToolTaskState
    final_solver_result: SolverResult
    metrics: dict


class Runner:
    def __init__(
        self,
        solver: Solver,
        sample: Any,
        name_to_tool: dict,
        max_turns: int,
        default_task_description: str,
        default_reminder_message: str,
    ):
        self.solver = solver
        self.sample = sample
        self.name_to_tool = name_to_tool
        self.max_turns = max_turns
        self.default_task_description = default_task_description
        self.default_reminder_message = default_reminder_message

    def run(self) -> RunnerResult:
        # Prepare initial task state
        tools = self.name_to_tool.values()
        tool_names_and_descriptions = self._get_tool_names_and_descriptions(tools)
        task_description = self.default_task_description.format(
            tool_names_and_descriptions=tool_names_and_descriptions
        )
        task_message = self.sample["task"]
        messages = [
            Message(role="user", content=task_message),
        ]
        task_state = TaskState(
            task_description=task_description,
            messages=messages,
            current_state=None,
        )

        # Loops until solver completes task or hits turn limit
        turn = 0
        final_answer = None
        while turn < self.max_turns:
            # Get result from solver
            solver_result = self.solver(task_state)
            parsed_solver_result = self._parse_solver_result(solver_result)
            final_answer = parsed_solver_result.final_answer

            # If solver failed to call tool or give final answer, prompt them to try again
            if parsed_solver_result.tool_calls == [] and final_answer is None:
                content = self.default_reminder_message
                task_state = self._add_eval_message(task_state, solver_result, content=content)
                turn += 1
                continue

            if final_answer is not None:
                return self._finish_run(task_state, solver_result, final_answer, turn)

            # Run tools. If solver gave tool incorrect input, prompt them to try again.
            assert parsed_solver_result.tool_calls != []
            tool_outputs = [self._run_tool_call(i) for i in parsed_solver_result.tool_calls]
            if any([i is None for i in tool_outputs]):
                content = self.default_reminder_message
                task_state = self._add_eval_message(task_state, solver_result, content=content)
                turn += 1
                continue

            # Add user message containing tool outputs
            task_state = self._add_tool_outputs(task_state, solver_result, tool_outputs)
            turn += 1

        return self._finish_run(task_state, solver_result, None, turn)

    def _get_tool_names_and_descriptions(self, tools: list[Tool]):
        """
        Given sequence of tools, creates a string of each tools name
        and description, each tool's info separated by a newline
        """
        s = ""
        for tool in tools:
            s += f"{tool._name}: {tool._desc}\n"
        return s

    def _parse_solver_result(self, solver_result: SolverResult) -> ParsedSolverResult:
        output = solver_result.output
        tool_calls = self._parse_tool_calls(output)
        final_answer = self._parse_final_answer(output)
        return ParsedSolverResult(tool_calls=tool_calls, final_answer=final_answer)

    def _parse_tool_calls(self, output: str) -> Optional[list[ToolCall]]:
        tool_message_matches = self._find_tool_messages(output)
        if tool_message_matches == []:
            return []

        tool_calls = []
        for tool_name, tool_message in tool_message_matches:
            # Log warning if solver calls a tool that doesn't exist
            try:
                self.name_to_tool[tool_name]
            except KeyError:
                logger.warn(f"Solver tried to call '{tool_name}' tool which doesn't exist!")
                continue

            tool_call = ToolCall(tool_name=tool_name, input=tool_message, output=None)
            tool_calls.append(tool_call)
        return tool_calls

    def _find_tool_messages(self, text: str) -> list[tuple[str, str]]:
        """
        Finds all tool calls, which are formatted [NAME: INPUT],
        where NAME != "Answer" and NAME != "Bugged"
        """
        pattern = r"\(@(?!Answer|Bugged)(\w+): (.+?)\)"
        matches = re.findall(pattern, text, re.DOTALL)
        return matches

    def _parse_final_answer(self, output: str) -> Optional[str]:
        """
        If a final answer exists of form [Answer: OUTPUT], returns the output,
        otherwise returns None
        """
        match = re.search(r"\(@Answer: (.*?)\)", output, re.DOTALL)
        return match.group(1) if match else None

    def _run_tool_call(self, tool_call: ToolCall) -> ToolCall:
        # Prepare task state
        tool_name = tool_call.tool_name
        tool = self.name_to_tool[tool_name]
        tool_input = tool_call.input
        tool_desc = self.name_to_tool[tool_name]._desc

        # Remove quotes if solver wrapped input
        if tool_input.startswith(("'", '"')) and tool_input.endswith(("'", '"')):
            tool_input = tool_input[1:-1]

        task_description = (
            f"Your name is {tool_name}. A description of your purpose is shown below:\n{tool_desc}"
        )
        messages = [Message(role="user", content=tool_input)]
        task_state = ToolTaskState(
            task_description=task_description, messages=messages, current_state=None
        )
        try:
            out = tool(task_state)
        except (TypeError, ValueError, IndexError):
            out = None

        if out is None:
            return None

        tool_call.output = out.output
        return tool_call

    def _add_eval_message(
        self,
        task_state: TaskState,
        solver_output: SolverResult,
        content: str,
    ) -> TaskState:
        messages = copy.deepcopy(task_state.messages)
        messages.append(Message(role="assistant", content=solver_output.output))
        # NOTE: we assume that the order of tool_outputs is the same as the order of tool_calls

        messages.append(Message(role="user", content=content))
        new_task_state = TaskState(
            task_description=task_state.task_description,
            messages=messages,
            current_state=None,
        )
        return new_task_state

    def _add_tool_outputs(
        self,
        task_state: TaskState,
        solver_output: SolverResult,
        tool_outputs: list[ToolCall],
    ) -> TaskState:
        content = ""
        for tool_output in tool_outputs:
            name = tool_output.tool_name
            input = tool_output.input
            output = tool_output.output
            content += f"{name} output on input {input}: {output}\n"

        return self._add_eval_message(task_state, solver_output, content)

    def _finish_run(
        self,
        final_task_state: TaskState,
        solver_result: SolverResult,
        final_answer: Optional[str],
        turn: int,
    ) -> RunnerResult:
        expected_answer = self.sample["answer"]
        is_correct = False
        if final_answer is not None:
            final_answer = final_answer.lower().strip()
            # Remove quotes if solver wrapped input
            if final_answer.startswith(("'", '"')) and final_answer.endswith(("'", '"')):
                final_answer = final_answer[1:-1]
            is_correct = final_answer == expected_answer.lower().strip()

        metrics = {
            "is_correct": is_correct,
            "num_turns": turn + 1,  # zero-indexed,
        }

        return RunnerResult(
            final_task_state,
            solver_result,
            metrics,
        )