Building a solver-judge workflow

import asyncio
import re

from openai import AsyncOpenAI
from rllm.types import Step, Trajectory


async def _generate_solutions(
    client: AsyncOpenAI, model: str, problem: str, n: int = 2
) -> list[Trajectory]:
    async def _solve() -> Trajectory:
        messages = [
            {
                "role": "user",
                "content": f"{problem}. Output the final answer within <answer>...</answer>",
            }
        ]
        response = await client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=1,
            max_tokens=1000,
        )
        content = response.choices[0].message.content or ""
        return Trajectory(
            name="solver",
            steps=[
                Step(
                    chat_completions=messages + [{"role": "assistant", "content": content}],
                    model_response=content,
                    action=_parse_answer(content),
                )
            ],
        )

    return await asyncio.gather(*(_solve() for _ in range(n)))


def _parse_answer(response: str) -> str:
    match = re.search(r"<answer>(.*?)</answer>", response, re.IGNORECASE | re.DOTALL)
    if match:
        return f"<answer>{match.group(1).strip()}</answer>"
    return "No solution found"

async def _judge_solutions(
    client: AsyncOpenAI, model: str, problem: str, solutions: list[str]
) -> Trajectory:
    prompt = _create_judge_prompt(problem, solutions)
    messages = [{"role": "user", "content": prompt}]
    response = await client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=1,
        max_tokens=1000,
    )
    content = response.choices[0].message.content or ""
    return Trajectory(
        name="judge",
        steps=[
            Step(
                chat_completions=messages + [{"role": "assistant", "content": content}],
                model_response=content,
                action=_parse_judge_response(content, solutions),
            )
        ],
    )


def _parse_judge_response(response: str, solutions: list[str]) -> str:
    match = re.search(r"<answer>(.*?)</answer>", response, re.IGNORECASE | re.DOTALL)
    if match:
        try:
            idx = int(match.group(1).strip())
            return solutions[idx - 1]
        except (ValueError, IndexError):
            return ""
    return ""


def _create_judge_prompt(problem: str, solutions: list[str]) -> str:
    prompt = (
        "You are an expert verifier. Given a countdown problem and multiple "
        "solution attempts, select a correct solution.\n"
        f"Problem:\n{problem}\nSolutions to evaluate:\n"
    )
    for i, solution in enumerate(solutions, 1):
        prompt += f"\nSolution {i}:\n{solution}\n"
    prompt += (
        "\nA correct solution must satisfy the following criteria:\n"
        "1. The solution uses only the given numbers.\n"
        "2. Each number is used exactly once.\n"
        "3. Only basic arithmetic operations (+, -, *, /) are used.\n"
        "4. The calculation results in the target number.\n"
        "5. The final answer is clearly marked within <answer>...</answer> tags.\n"
        "Output the index of your selected solution within <answer>...</answer> tags, "
        "e.g., <answer>1</answer> for the first solution. If multiple solutions are "
        "correct, output the index of the first correct one."
    )
    return prompt

import rllm
from rllm.types import AgentConfig, Episode, Task

N_SOLUTIONS = 2


@rllm.rollout(name="solver-judge")
async def solver_judge_flow(task: Task, config: AgentConfig) -> Episode:
    client = AsyncOpenAI(base_url=config.base_url, api_key="EMPTY")
    problem = task.instruction

    # 1. Solver generates N solutions in parallel.
    solver_trajectories = await _generate_solutions(
        client, config.model, problem, n=N_SOLUTIONS
    )

    # 2. Judge selects the best solution.
    solutions = [t.steps[0].action for t in solver_trajectories]
    judge_trajectory = await _judge_solutions(client, config.model, problem, solutions)

    # 3. Bundle everything into an Episode.
    selected = judge_trajectory.steps[0].action
    return Episode(
        trajectories=[*solver_trajectories, judge_trajectory],
        artifacts={"answer": selected},
    )