Spaces:
Sleeping
Sleeping
File size: 4,459 Bytes
75ca235 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 |
"""Task definitions and data loaders for Bug Triage OpenEnv."""
from __future__ import annotations
import random
from copy import deepcopy
from datetime import datetime
from typing import Optional
from models import TicketGroundTruth, TicketModel
from . import task_easy, task_hard, task_medium
TASK_MODULES = {
task_easy.TASK_ID: task_easy,
task_medium.TASK_ID: task_medium,
task_hard.TASK_ID: task_hard,
}
class TaskDefinition:
"""Represents a task with tickets and ground truth."""
def __init__(
self,
task_id: str,
tickets: list[TicketModel],
ground_truths: list[TicketGroundTruth],
step_budget: int,
available_teams: list[str],
available_components: list[str],
):
self.task_id = task_id
self.tickets = tickets
self.ground_truths = ground_truths
self.step_budget = step_budget
self.available_teams = available_teams
self.available_components = available_components
self._validate_integrity()
self.truth_map = {gt.ticket_id: gt for gt in ground_truths}
def get_ground_truth(self, ticket_id: str) -> Optional[TicketGroundTruth]:
"""Get ground truth for a ticket."""
return self.truth_map.get(ticket_id)
def _validate_integrity(self) -> None:
"""Ensure ticket and ground-truth ids are unique and aligned."""
ticket_ids = [ticket.ticket_id for ticket in self.tickets]
ground_truth_ids = [truth.ticket_id for truth in self.ground_truths]
if len(ticket_ids) != len(set(ticket_ids)):
raise ValueError("Duplicate ticket_id values found in tickets")
if len(ground_truth_ids) != len(set(ground_truth_ids)):
raise ValueError("Duplicate ticket_id values found in ground truths")
ticket_id_set = set(ticket_ids)
ground_truth_id_set = set(ground_truth_ids)
if ticket_id_set != ground_truth_id_set:
missing_truth = sorted(ticket_id_set - ground_truth_id_set)
extra_truth = sorted(ground_truth_id_set - ticket_id_set)
raise ValueError(
"Tickets and ground truths must contain the same ticket ids. "
f"missing_ground_truth={missing_truth}, extra_ground_truth={extra_truth}"
)
def shuffle_tickets(self, seed: int):
"""Shuffle tickets deterministically."""
self._validate_integrity()
if not self.tickets:
self.truth_map = {}
return
rng = random.Random(seed)
shuffled_tickets = list(self.tickets)
rng.shuffle(shuffled_tickets)
ground_truth_by_id = {truth.ticket_id: truth for truth in self.ground_truths}
self.tickets = shuffled_tickets
self.ground_truths = [ground_truth_by_id[ticket.ticket_id] for ticket in self.tickets]
self.truth_map = dict(ground_truth_by_id)
def _task_data(task_id: str) -> dict:
module = TASK_MODULES.get(task_id)
if module is None:
raise FileNotFoundError(f"Unknown task id: {task_id}")
return deepcopy(module.TASK_DATA)
def load_task(task_id: str, seed: Optional[int] = None) -> TaskDefinition:
"""
Load a task from in-memory fixture.
Args:
task_id: One of 'bug_triage_easy', 'bug_triage_medium', 'bug_triage_hard'
seed: Random seed for deterministic shuffling
Returns:
TaskDefinition with loaded tickets and ground truth
"""
data = _task_data(task_id)
tickets = []
for ticket_data in data["tickets"]:
ticket_data["created_at"] = datetime.fromisoformat(ticket_data["created_at"])
tickets.append(TicketModel(**ticket_data))
ground_truths = [TicketGroundTruth(**gt_data) for gt_data in data["ground_truths"]]
task = TaskDefinition(
task_id=task_id,
tickets=tickets,
ground_truths=ground_truths,
step_budget=data["step_budget"],
available_teams=data["available_teams"],
available_components=data["available_components"],
)
if seed is not None:
task.shuffle_tickets(seed)
return task
def list_tasks() -> dict:
"""List all available tasks."""
return {
task_id: {
"difficulty": module.DIFFICULTY,
"description": module.DESCRIPTION,
}
for task_id, module in TASK_MODULES.items()
}
__all__ = [
"TaskDefinition",
"load_task",
"list_tasks",
]
|