mixle.task.plan module

distill_planner – train tiny models to break a problem into steps and carry them out.

Decomposition, the mixle way: a plan is an autoregressive chain of calibrated tool calls ending in STOP. The teacher (a frontier 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 – and “predict the next call” is exactly 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 honesty contract extends step-wise: a step is emitted only when the selector is confident AND 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 half-executed guessed plan is never returned silently – 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-strength decomposition (the students are tiny); free-form novel planning still wants trace-SFT on the causal LM – the same trace format feeds it when that rung lands.

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:
selector: Any
extractors: dict[str, Any]
tools: dict[str, ToolSpec]
teacher: Callable[[str], list[dict]]
plan_agreement: float
max_steps: int = 8
n_requests: int = 0
n_escalated: int = 0
harvested: list[tuple[str, list[dict]]]
try_plan(request, *, execute=None)[source]

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

No teacher, no stats — this is what a server without the frontier can run.

Parameters:
Return type:

dict[str, Any] | None

report()[source]
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