mixle.task.plan module

distill_planner trains local models to decompose requests into tool steps.

The plan representation is an autoregressive chain of calibrated tool calls ending in STOP. The teacher (an LLM, an agent loop, or a rule) shows plans for example requests; each trace flattens into (context, next-call) pairs where the context is the request plus the steps taken so far. “Predict the next call” is the problem toolcall already solves: a conformal selector for which tool comes next (STOP is just another action) and a per-tool extractor for its arguments, both reading the rendered context.

The safety contract is stepwise: a step is emitted only when the selector is confident, the required arguments extract, and, when an execute map is given, the call actually runs. Any failure escalates the whole request to the teacher; a partially executed guessed plan is not returned as local success, and the escalation is harvested as a fresh trace for the next distillation round.

teacher(request) -> [{“tool”: …, “args”: {…}}, …] # the plan planner = distill_planner(teacher, requests, tools) planner(request) # {“plan”, “escalate”} planner(request, execute={“lookup”: fn, …}) # + per-step “results”, verified

This is template-oriented decomposition. For free-form generated plans, use the trace-SFT planner on the same trace format.

class Planner(selector, extractors, tools, teacher, plan_agreement, max_steps=8, n_requests=0, n_escalated=0, harvested=<factory>)[source]

Bases: object

A distilled decomposer: emit verified steps until STOP, or escalate the whole problem.

Parameters:
try_plan(request, *, execute=None)[source]

The local decomposition alone: a complete verified plan, or None (= must escalate).

This method does not call the teacher.

Parameters:
Return type:

dict[str, Any] | None

report()[source]

Return plan agreement, escalation, and harvested-trace metrics.

Return type:

dict[str, Any]

save(path)[source]

Persist selector + per-tool extractors + specs as one artifact directory; load() restores.

Parameters:

path (str)

Return type:

str

classmethod load(path, teacher, *, device='cpu')[source]

Reconstitute a serving Planner from save() output plus the teacher fallback.

Parameters:
Return type:

Planner

distill_planner(teacher, requests, tools, *, holdout=0.2, seed=0, max_steps=8, selector_kw=None, extractor_kw=None)[source]

Distill the teacher’s multi-step plans into next-step students (see module docstring).

Plan-level verification is measured on held-out requests the students never trained on: a plan agrees when every step’s tool and required arguments match the teacher’s plan exactly, in order.

Parameters:
Return type:

Planner