agentsociety.taskloader.taskloader

TaskLoader module for loading and managing tasks from JSON/JSONL files.

This module provides a PyTorch DataLoader-like interface for task management, supporting task loading, extraction, and state tracking.

Module Contents

Classes

TaskStatus	Task execution status enumeration.
Task	Represents a single task with context, status, and results.
TaskLoader	A PyTorch DataLoader-like class for loading and managing tasks from JSON/JSONL files.

API

class agentsociety.taskloader.taskloader.TaskStatus(*args, **kwds)

Bases: enum.Enum

Task execution status enumeration.

Initialization

PENDING

‘pending’

RUNNING

‘running’

COMPLETED

‘completed’

class agentsociety.taskloader.taskloader.Task

Represents a single task with context, status, and results.

• Description: A task object that contains execution context, status tracking, ground truth data, and result storage capabilities.

• Args:

  • ground_truth (Any): Ground truth data for the task

  • task_id (int): Unique identifier for the task

  • status (TaskStatus): Current execution status of the task

  • result (Optional[Any]): Execution result of the task

  • assigned_agent_id (Optional[int]): ID of the agent assigned to this task

ground_truth: Any

None

task_id: int

None

status: agentsociety.taskloader.taskloader.TaskStatus

None

result: Optional[Any]

None

assigned_agent_id: Optional[int]

None

get_task_context() → dict[str, Any]

Returns the task information excluding basic fields.

• Description: Returns all task information except for ground_truth, task_id, status, and result. This method dynamically includes all fields from the task object and its subclasses.

• Returns:

  • dict[str, Any]: Task information dictionary

set_result(result: Any) → None

Sets the task result and marks it as completed.

• Description: Updates the task result and automatically changes status to completed.

• Args:

  • result (Any): The execution result to store

set_running() → None

Marks the task as running.

• Description: Changes the task status to running when execution begins.

reset() → None

Resets the task to pending status and clears results.

• Description: Resets the task to its initial state for re-execution.

assign_to_agent(agent_id: int) → None

Assigns this task to a specific agent.

• Description: Sets the assigned_agent_id field to track which agent is responsible for executing this task.

• Args:

  • agent_id (int): The ID of the agent assigned to this task

class agentsociety.taskloader.taskloader.TaskLoader(task_type: type[agentsociety.taskloader.taskloader.Task], file_path: str, shuffle: bool = False, max_tasks: Optional[int] = None)

A PyTorch DataLoader-like class for loading and managing tasks from JSON/JSONL files.

• Description: Loads tasks from JSON or JSONL files and provides methods to extract tasks for execution. Supports task state management and result tracking.

• Args:

  • file_path (str): Path to the JSON/JSONL file containing tasks

  • shuffle (bool): Whether to shuffle tasks when loading

  • max_tasks (Optional[int]): Maximum number of tasks to load

Initialization

_load_tasks() → None

Loads tasks from the specified file.

• Description: Reads the JSON/JSONL file and converts each entry into a Task object. Supports both JSON (list of tasks) and JSONL (one task per line) formats.

next(n: int = 1) → Optional[Union[agentsociety.taskloader.taskloader.Task, List[agentsociety.taskloader.taskloader.Task]]]

Extracts the next N tasks for execution.

• Description: Returns the next N pending tasks and marks them as running. If only one task is requested, returns a single Task object. If multiple tasks are requested, returns a list of Task objects. Returns None if no pending tasks are available.

• Args:

  • n (int): Number of tasks to extract (default: 1)

• Returns:

  • Optional[Union[Task, List[Task]]]: Single task, list of tasks, or None

get_pending_count() → int

Returns the number of pending tasks.

• Description: Counts tasks that are still in pending status.

• Returns:

  • int: Number of pending tasks

get_completed_count() → int

Returns the number of completed tasks.

• Description: Counts tasks that have been completed.

• Returns:

  • int: Number of completed tasks

get_running_count() → int

Returns the number of running tasks.

• Description: Counts tasks that are currently running.

• Returns:

  • int: Number of running tasks

reset_all() → None

Resets all tasks to pending status and clears collection tracking.

• Description: Resets all tasks to their initial state for re-execution and clears the collection tracking set.

get_task_by_id(task_id: str) → Optional[agentsociety.taskloader.taskloader.Task]

Retrieves a task by its ID.

• Description: Finds and returns a specific task using its unique identifier.

• Args:

  • task_id (str): The unique identifier of the task

• Returns:

  • Optional[Task]: The task if found, None otherwise

get_tasks_by_status(status: agentsociety.taskloader.taskloader.TaskStatus) → List[agentsociety.taskloader.taskloader.Task]

Returns all tasks with the specified status.

• Description: Filters tasks by their execution status.

• Args:

  • status (TaskStatus): The status to filter by

• Returns:

  • List[Task]: List of tasks with the specified status

__len__() → int

Returns the total number of tasks.

• Description: Returns the total count of all tasks in the loader.

• Returns:

  • int: Total number of tasks

__iter__()

Makes TaskLoader iterable.

• Description: Allows TaskLoader to be used in for loops, yielding pending tasks.

__next__() → agentsociety.taskloader.taskloader.Task

Returns the next pending task in iteration.

• Description: Returns the next pending task and marks it as running. Raises StopIteration when no more pending tasks are available.

• Returns:

  • Task: The next pending task

• Raises:

  • StopIteration: When no more pending tasks are available

get_task_results()

Extracts newly completed tasks and converts them to StorageTaskResult format.

• Description: Retrieves only the newly completed tasks (not previously collected) and converts them to StorageTaskResult format for storage. Uses the assigned_agent_id from each task to determine the agent.

• Args:

  • agent_id (int): The ID of the agent that executed the tasks

• Returns:

  • List[StorageTaskResult]: List of new task results in storage format

get_uncollected_completed_count() → int

Returns the number of completed tasks that haven’t been collected yet.

• Description: Counts completed tasks that haven’t been collected by getTaskResults yet.

• Returns:

  • int: Number of uncollected completed tasks