from inspect_evals.infinite_bench.constants import (
    DEFAULT_SYSTEM_MSG,
    MATH_CALC_SYSTEM_MSG,
    MODEL_CONTEXT_SIZES,
)


# Based on source code from https://github.com/OpenBMB/InfiniteBench
def create_system_msg(task_name: str) -> str:
    """
    Create a system message based on the task name.

    - For "math_calc", use MATH_CALC_SYSTEM_MSG.
    - Otherwise, use DEFAULT_SYSTEM_MSG.

    Args:
        task_name (str): The name of the task to create the system message for.

    Returns:
        str: The system message for the task.
    """
    if task_name == "math_calc":
        return MATH_CALC_SYSTEM_MSG
    else:
        return DEFAULT_SYSTEM_MSG


def get_model_context_size(model_name: str) -> int:
    """
    Get the model context size.

    Args:
        model_name (str): The name of the model to get the context size for.

    Returns:
        int: The context size for the model.
    """
    max_tokens = MODEL_CONTEXT_SIZES.get(model_name, 128 * 1000)

    return max_tokens
