NexQuant/rdagent/components/task_implementation/factor_implementation/evolving/scheduler.py

import json
from pathlib import Path

from jinja2 import Environment, StrictUndefined

from rdagent.components.task_implementation.factor_implementation.evolving.evolvable_subjects import (
    FactorEvolvingItem,
)
from rdagent.components.task_implementation.factor_implementation.utils import (
    get_data_folder_intro,
)
from rdagent.core.conf import RD_AGENT_SETTINGS
from rdagent.core.log import RDAgentLog
from rdagent.core.prompts import Prompts
from rdagent.oai.llm_utils import APIBackend

scheduler_prompts = Prompts(file_path=Path(__file__).parent.parent / "prompts.yaml")


def RandomSelect(to_be_finished_task_index, implementation_factors_per_round):
    import random

    to_be_finished_task_index = random.sample(
        to_be_finished_task_index,
        implementation_factors_per_round,
    )

    RDAgentLog().info(f"The random selection is: {to_be_finished_task_index}")
    return to_be_finished_task_index


def LLMSelect(to_be_finished_task_index, implementation_factors_per_round, evo: FactorEvolvingItem, former_trace):
    tasks = []
    for i in to_be_finished_task_index:
        # find corresponding former trace for each task
        target_factor_task_information = evo.sub_tasks[i].get_factor_information()
        if target_factor_task_information in former_trace:
            tasks.append((i, evo.sub_tasks[i], former_trace[target_factor_task_information]))

    system_prompt = (
        Environment(undefined=StrictUndefined)
        .from_string(
            scheduler_prompts["select_implementable_factor_system"],
        )
        .render(
            data_info=get_data_folder_intro(),
        )
    )

    session = APIBackend(use_chat_cache=False).build_chat_session(
        session_system_prompt=system_prompt,
    )

    while True:
        user_prompt = (
            Environment(undefined=StrictUndefined)
            .from_string(
                scheduler_prompts["select_implementable_factor_user"],
            )
            .render(
                factor_num=implementation_factors_per_round,
                sub_tasks=tasks,
            )
        )
        if (
            session.build_chat_completion_message_and_calculate_token(
                user_prompt,
            )
            < RD_AGENT_SETTINGS.chat_token_limit
        ):
            break

    response = session.build_chat_completion(
        user_prompt=user_prompt,
        json_mode=True,
    )
    try:
        selection = json.loads(response)["selected_factor"]
        if not isinstance(selection, list):
            return to_be_finished_task_index
        selection_index = [x for x in selection if isinstance(x, int)]
    except:
        return to_be_finished_task_index

    return selection_index