Implementation:Princeton nlp Tree of thought llm Get Task Factory

Knowledge Sources	tree-of-thought-llm
Domains	Software_Design, CLI_Design
Last Updated	2026-02-14 03:30 GMT

Overview

Concrete tool for registering new tasks in the factory function and CLI argument parser provided by the Tree of Thoughts framework.

Description

Task registration requires modifying two files:

• src/tot/tasks/__init__.py — Add an elif branch in get_task() that imports and instantiates the new task class.
• run.py — Add the new task name to the choices list in the --task argparse argument definition (line 48).

The factory uses lazy imports (importing the task class inside the elif block) to avoid loading all task dependencies when only one task is needed.

Usage

Follow this implementation when adding a new task to the framework. After implementing the task subclass and prompt templates, these two registration edits make the task runnable via the CLI.

Code Reference

Source Location

• Repository: tree-of-thought-llm
• File: src/tot/tasks/__init__.py (Lines 1-12), run.py (Line 48)

Signature

# src/tot/tasks/__init__.py — Current implementation
def get_task(name):
    if name == 'game24':
        from tot.tasks.game24 import Game24Task
        return Game24Task()
    elif name == 'text':
        from tot.tasks.text import TextTask
        return TextTask()
    elif name == 'crosswords':
        from tot.tasks.crosswords import MiniCrosswordsTask
        return MiniCrosswordsTask()
    else:
        raise NotImplementedError

# run.py line 48 — Current choices
args.add_argument('--task', type=str, required=True,
                  choices=['game24', 'text', 'crosswords'])

Import

from tot.tasks import get_task

I/O Contract

Inputs (Registration)

Name	Type	Required	Description
task_name	str	Yes	New task identifier to add to factory and CLI choices
TaskClass	class	Yes	Implemented Task subclass to be instantiated

Outputs (After Registration)

Name	Type	Description
get_task(name)	Task	Factory now returns the new task when called with the new name
CLI validation	bool	argparse now accepts the new task name without error

Usage Examples

Adding a New Task Registration

# Step 1: Edit src/tot/tasks/__init__.py
# Add before the `else: raise NotImplementedError` line:
def get_task(name):
    if name == 'game24':
        from tot.tasks.game24 import Game24Task
        return Game24Task()
    elif name == 'text':
        from tot.tasks.text import TextTask
        return TextTask()
    elif name == 'crosswords':
        from tot.tasks.crosswords import MiniCrosswordsTask
        return MiniCrosswordsTask()
    elif name == 'sudoku':                          # NEW
        from tot.tasks.sudoku import SudokuTask     # NEW
        return SudokuTask()                         # NEW
    else:
        raise NotImplementedError

# Step 2: Edit run.py line 48
args.add_argument('--task', type=str, required=True,
                  choices=['game24', 'text', 'crosswords', 'sudoku'])  # ADD 'sudoku'

Verifying Registration

# After registration, verify the task can be instantiated
from tot.tasks import get_task

task = get_task('sudoku')
print(type(task))  # <class 'tot.tasks.sudoku.SudokuTask'>
print(len(task))   # Number of puzzles
print(task.steps)  # BFS depth

Related Pages

Implements Principle

• Principle:Princeton_nlp_Tree_of_thought_llm_Task_Registration

Requires Environment

• Environment:Princeton_nlp_Tree_of_thought_llm_Python_OpenAI