def generate_task_category()

in project/paperbench/paperbench/rubric/tasks.py [0:0]


def generate_task_category(node: TaskNode, model: str = "gpt-4o") -> str:
    """Uses an LLM to generate a task category for the given `node` based on its `requirements`."""

    client = get_openai_client()

    messages = [
        {
            "role": "system",
            "content": (
                "You are an assistant that classifies tasks into predefined categories. "
                "The categories are:\n"
            )
            + "\n".join(f"- {category}" for category in VALID_TASK_CATEGORIES)
            + "\n\n"
            + (
                "Please read the following task requirement and determine the most appropriate "
                "category from the list above. Respond with only the category name, exactly as "
                "written."
            ),
        },
        {
            "role": "user",
            "content": f"Task requirement: {node.requirements}",
        },
    ]
    completion = oai_completion_with_retry(
        client.chat.completions.create,
        messages=messages,
        model=model,
    )
    response = completion.choices[0].message.content.strip()

    if response not in VALID_TASK_CATEGORIES:
        raise ValueError(f"Invalid task category generated: {response}")

    return response