trinity.common.workflows.thinking_budget_workflow module

OpenAI-compatible workflow with vLLM thinking-budget support.

class trinity.common.workflows.thinking_budget_workflow.ThinkingBudgetWorkflow(*, task: Task, model: ModelWrapper, auxiliary_models: List[ModelWrapper] | None = None)

Bases: Workflow

Run raw OpenAI messages with a per-request reasoning-token budget.

vLLM forces the configured end-of-reasoning token sequence when the budget is exhausted. The workflow locates that sequence in the returned completion token IDs and masks it out so it does not contribute to the policy loss.

can_repeat: bool = True

DEFAULT_REASONING_END_STR = '</think>'

__init__(*, task: Task, model: ModelWrapper, auxiliary_models: List[ModelWrapper] | None = None)

set_repeat_times(repeat_times: int, run_id_base: int) → None

Set the number of times to repeat the workflow. :param repeat_times: number of times to repeat the workflow (if repeatable). :type repeat_times: int :param run_id_base: base run_id for setting run_id in experiences. :type run_id_base: int

run() → List[Experience]

Run workflow and return a list of experiences.