agent_system.py

import asyncio
import re
import time
import traceback
from copy import deepcopy

from vime.rollout.rm_hub import batched_async_rm
from vime.rollout.vllm_rollout import (
    _build_inference_sampling_params,
    _inference_generate_meta_info,
    _inference_generate_tokens_and_logprobs,
)
from vime.utils.http_utils import post
from vime.utils.types import Sample

from .prompts import SOLVER_PROMPT_TEMPLATE, generate_rewriter_template, generate_select_template


async def generate_response(args, prompt, key):
    try:
        sampling_params = args.sampling_params
        tokenizer = args.tokenizer
        max_context_length = args.rollout_max_context_len
        sample = deepcopy(args.sample)

        url = f"http://{args.vllm_router_ip}:{args.vllm_router_port}/inference/v1/generate"

        sample.prompt = prompt
        prompt_token_ids = tokenizer(sample.prompt, add_special_tokens=False)["input_ids"]
        sample.tokens = prompt_token_ids
        prompt_length = len(prompt_token_ids)
        current_sampling_params = deepcopy(sampling_params)
        current_sampling_params["max_new_tokens"] = min(
            sampling_params["max_new_tokens"], max_context_length - prompt_length
        )

        if current_sampling_params["max_new_tokens"] <= 0:
            return None

        payload = {
            "model": args.hf_checkpoint,
            "token_ids": prompt_token_ids,
            "sampling_params": _build_inference_sampling_params(current_sampling_params),
        }

        output = await post(url, payload)
        choice = output["choices"][0]
        finish_reason = choice.get("finish_reason") or "stop"
        new_response_tokens, new_response_log_probs = _inference_generate_tokens_and_logprobs(choice)
        response_text = tokenizer.decode(new_response_tokens, skip_special_tokens=False) if new_response_tokens else ""

        sample.append_response_tokens(
            args,
            tokens=new_response_tokens,
            log_probs=new_response_log_probs,
            trainable=True,
            meta_info=_inference_generate_meta_info(output),
        )
        assert len(sample.rollout_log_probs) == sample.response_length, (
            f"rollout logprob length mismatch: {len(sample.rollout_log_probs)} logprobs "
            f"vs {sample.response_length} response tokens"
        )
        sample.response = response_text

        match finish_reason:
            case "length":
                sample.status = Sample.Status.TRUNCATED
            case "abort" | "cancelled":
                sample.status = Sample.Status.ABORTED
            case _:
                sample.status = Sample.Status.COMPLETED

        args.results_dict[key].append(sample)

        final = response_text.replace("<|user|>", "")
        if "</think>" in final:
            contents = final.split("</think>")
            if len(contents) == 2 and contents[1] != "":
                reason_content = contents[0].strip()
                response_content = contents[1].strip()
                sample.reason_content = reason_content
                sample.response_content = response_content
                return response_content
        sample.reason_content = None
        sample.response_content = None
        return None
    except Exception as e:
        print(f"Error generating response: {e}")
        return None


class Agent:
    """A base class for our AI agents."""

    def __init__(self):
        pass

    async def run(self, args, prompt, max_retries: int = 1, key: str = None) -> str:
        """Runs the agent by sending a prompt to the LLM."""
        for _i in range(max_retries):
            try:
                response = await generate_response(args, prompt, key=key)
                return response
            except Exception as e:
                print(f"Error querying LLM: {e}")
                time.sleep(1)
        print(f"Failed to query LLM after {max_retries} retries")
        return None


class SolverAgent(Agent):
    """The agent responsible for generating and improving solutions."""

    def __init__(self):
        super().__init__()

    async def generate_initial_solution(self, args, problem_statement) -> str:
        """Generates the first solution attempt."""
        prompt = SOLVER_PROMPT_TEMPLATE.format(problem_statement=problem_statement)
        return await self.run(args, prompt, max_retries=3, key="solver")


class RewriterAgent(Agent):
    """The agent responsible for rewriting solutions."""

    def __init__(self):
        super().__init__()

    async def rewrite(self, args, problem_statement, previous_solutions: list[str]) -> str:
        """Generates the rewrited solution."""

        # Build the prompt template dynamically.
        template = generate_rewriter_template(len(previous_solutions))

        # Populate the template arguments.
        format_params = {"problem_statement": problem_statement}
        for i, solution in enumerate(previous_solutions):
            format_params[f"solution{i+1}"] = solution

        prompt = template.format(**format_params)
        return await self.run(args, prompt, max_retries=1, key="rewriter")


class SelectorAgent(Agent):
    """The agent responsible for selecting solutions."""

    def __init__(self):
        super().__init__()

    async def select(self, args, problem_statement, candidate_solutions: list[str]) -> str:
        """Generates the rewrited solution."""

        # Build the prompt template dynamically.
        template = generate_select_template(len(candidate_solutions))

        # Populate the template arguments.
        format_params = {"problem_statement": problem_statement}
        for i, solution in enumerate(candidate_solutions):
            format_params[f"solution{i+1}"] = solution

        prompt = template.format(**format_params)
        return await self.run(args, prompt, max_retries=10, key="selector")

    def extract_selected_solution_idx(self, response: str, candidate_solutions: list[str]) -> int:
        """Extracts the selected solution ID from the response."""
        PATTERN = re.compile(r"Judgment:\s*(\d+)")
        matched = PATTERN.findall(response)
        try:
            selected_id = int(matched[0]) - 1
            if selected_id < len(candidate_solutions) and selected_id >= 0:
                return selected_id
            else:
                return None
        except Exception as e:
            print(f"extract_selected_solution_idx error: {e}")
            return None


async def rewrite_worker(args, previous_solutions, problem_statement, worker_id):
    rewriter = RewriterAgent()
    new_solution = await rewriter.rewrite(args, problem_statement, previous_solutions)
    return new_solution


async def solver_worker(args, problem_statement, worker_id):
    """
    Run a single solver pipeline.
    """

    try:
        solver = SolverAgent()
        current_solution = await solver.generate_initial_solution(args, problem_statement)
        return current_solution

    except Exception as e:
        print(f"[Worker-{worker_id}] exception: {e}")
        print(f"[Worker-{worker_id}] traceback: {traceback.format_exc()}")
        return None


async def run_agent_system(args, sample):
    """
    Run `num_parallel` pipelines concurrently.
    """

    args = deepcopy(args)  # Deep copy args because rollout_with_multi_agents mutates them.
    args.sample = sample
    args.results_dict = {"solver": [], "rewriter": [], "selector": []}
    # Every sample emitted below is a training sample split out of this one
    # rollout execution (the input ``sample``). Stamp the shared rollout id on
    # every collected sample at each return point so the per-rollout loss
    # reducer aggregates the solver / rewriter / selector siblings as one
    # rollout instead of N, and the by-rollout step splitter keeps them in
    # the same step. Captured here because ``sample`` gets shadowed by zip-
    # loop variables further down.
    input_rollout_id = sample.index

    def _emit(samples_list):
        for emitted_sample in samples_list:
            emitted_sample.rollout_id = input_rollout_id
        return samples_list

    problem_statement = sample.prompt
    tasks = [solver_worker(args, problem_statement, worker_id) for worker_id in range(args.num_parallel)]
    results = await asyncio.gather(*tasks, return_exceptions=True)

    rewards = await batched_async_rm(args, args.results_dict["solver"])
    for sample, reward in zip(args.results_dict["solver"], rewards, strict=False):
        sample.reward = reward

    previous_solutions = [item for item in results if isinstance(item, str)]

    def reward_adjustment(samples, reward_weight):
        for sample in samples:
            sample.reward = sample.reward * reward_weight
        return samples

    if len(previous_solutions) == 0:
        reward_adjustment(args.results_dict["solver"], args.incorrect_reward_weight)
        return _emit(args.results_dict["solver"])

    # Rewriting
    tasks = [
        rewrite_worker(args, previous_solutions, problem_statement, worker_id)
        for worker_id in range(args.num_parallel)
    ]
    rewrited_solutions_raw = await asyncio.gather(*tasks, return_exceptions=True)

    # Filter out failed tasks and keep only valid rewritten solutions.
    rewrited_solutions = []
    for _i, result in enumerate(rewrited_solutions_raw):
        if isinstance(result, str):
            rewrited_solutions.append(result)

    rewards = await batched_async_rm(args, args.results_dict["rewriter"])
    for sample, reward in zip(args.results_dict["rewriter"], rewards, strict=False):
        sample.reward = reward

    if len(rewrited_solutions) == 0:
        reward_adjustment(args.results_dict["solver"], args.incorrect_reward_weight)
        reward_adjustment(args.results_dict["rewriter"], args.incorrect_reward_weight)
        return _emit(args.results_dict["solver"] + args.results_dict["rewriter"])

    # Selection
    selector = SelectorAgent()
    response = await selector.select(args, problem_statement, rewrited_solutions)
    if len(args.results_dict["selector"]) == 0:
        reward_adjustment(args.results_dict["solver"], args.incorrect_reward_weight)
        reward_adjustment(args.results_dict["rewriter"], args.incorrect_reward_weight)
        return _emit(args.results_dict["solver"] + args.results_dict["rewriter"])

    assert (
        len(args.results_dict["selector"]) == 1
    ), f"selector should only return one solution, but got {len(args.results_dict['selector'])}"
    if response is None:
        args.results_dict["selector"][0].reward = 0
    else:
        selected_solution_idx = selector.extract_selected_solution_idx(response, rewrited_solutions)
        if selected_solution_idx is None:
            args.results_dict["selector"][0].reward = 0
        else:
            selected_solution = rewrited_solutions[selected_solution_idx]
            for sample in args.results_dict["rewriter"]:
                if sample.response_content is not None and selected_solution in sample.response_content:
                    args.results_dict["selector"][0].reward = sample.reward
                    break

    ## If the final answer is correct, add a positive bonus to all rewards.
    ## Otherwise, add a negative penalty to all rewards.
    if args.results_dict["selector"][0].reward == 1:
        reward_adjustment(args.results_dict["solver"], args.correct_reward_weight)
        reward_adjustment(args.results_dict["rewriter"], args.correct_reward_weight)
        reward_adjustment(args.results_dict["selector"], args.correct_reward_weight)
    else:
        reward_adjustment(args.results_dict["solver"], args.incorrect_reward_weight)
        reward_adjustment(args.results_dict["rewriter"], args.incorrect_reward_weight)
        reward_adjustment(args.results_dict["selector"], args.incorrect_reward_weight)

    return _emit(args.results_dict["solver"] + args.results_dict["rewriter"] + args.results_dict["selector"])