Skip to content

Commit d544b7b

Browse files
ampagentIsaac Miller
andcommitted
Make Decide a Module and Parameter rather than a Predict
Amp-Thread-ID: https://ampcode.com/threads/T-01a0bf8d-a1dd-7603-906c-ea06c46ff7f7 Co-authored-by: Isaac Miller <isaac@cmpnd.ai>
1 parent 528067d commit d544b7b

3 files changed

Lines changed: 95 additions & 19 deletions

File tree

‎docs/docs/api/experimental/Decide.md‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -140,8 +140,11 @@ endpoint default to `TYPESAFE_DEFAULT_MODEL` and `TYPESAFE_BASE_URL`, or `jev-la
140140
and `https://api.typesafe.ai`. A client passed to `Decide(..., client=...)` takes
141141
precedence over `dspy.settings.system_one`; no text LM is needed.
142142

143-
`Decide` participates in module composition, `named_predictors()`, callbacks,
144-
traces, `batch`, and `acall`. It validates required inputs and rejects unknown
143+
`Decide` inherits from `Module` and `Parameter`, not `Predict`. It participates in
144+
module composition, `named_parameters()`, callbacks, traces, `batch`, and `acall`.
145+
`named_predictors()` excludes it, so predictor-specific optimizers can target the
146+
`Predict` leaves of a mixed program without attaching demonstrations to `Decide`.
147+
It validates required inputs and rejects unknown
145148
inputs and unsupported outputs. Signature task instructions and field descriptions
146149
are included in each provider question. `Decide` does not support demonstrations:
147150
passing `demos=` or attaching nonempty `.demos` raises before inference. Optimizers
@@ -162,6 +165,10 @@ restored.load("assess.json")
162165
```
163166

164167
State-only loading requires the same signature architecture, as with `Predict`.
168+
Saved state contains only `signature`, `thresholds`, `weights`, and `client`.
169+
`reset()` preserves configuration because there is no demonstration or training
170+
state to clear. `reset_copy()` makes an independent copy with the same settings,
171+
so predictor-specific optimizers do not erase tuned decision parameters.
165172
Parameters survive this round trip; demonstrations are not saved, and loading
166173
state with nonempty demonstrations raises rather than silently discarding them. Explicit TypeSafe
167174
client settings are saved without API keys; credentials come from the environment

‎dspy/predict/decide.py‎

Lines changed: 28 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,17 @@
1010
from dspy.adapters.utils import get_field_description_string
1111
from dspy.clients.typesafe import TypeSafe
1212
from dspy.dsp.utils.settings import settings
13-
from dspy.predict.predict import Predict, _sanitize_lm_state, serialize_object
13+
from dspy.predict.parameter import Parameter
14+
from dspy.predict.predict import _sanitize_lm_state, serialize_object
15+
from dspy.primitives.module import Module
16+
from dspy.primitives.prediction import Prediction
1417
from dspy.signatures.signature import Signature, ensure_signature
1518
from dspy.utils.annotation import experimental
1619
from dspy.utils.callback import BaseCallback
1720

1821

1922
@experimental
20-
class Decide(Predict):
23+
class Decide(Module, Parameter):
2124
"""Answer a signature's closed-set outputs in one System One request.
2225
2326
Declare outputs using Noul, Score[(number, description), ...], and
@@ -47,7 +50,8 @@ def __init__(
4750
client: TypeSafe | None = None,
4851
callbacks: list[BaseCallback] | None = None,
4952
):
50-
super().__init__(signature, callbacks=callbacks)
53+
super().__init__(callbacks=callbacks)
54+
self.signature = ensure_signature(signature)
5155
self.client = client
5256
types = self._output_types(self.signature)
5357
self.thresholds = {name: 0.5 for name, kind in types.items() if kind is Noul}
@@ -56,6 +60,9 @@ def __init__(
5660
{name: {str(v): 1.0 for v, _ in kind.options} for name, kind in types.items() if issubclass(kind, Choice)}
5761
)
5862

63+
def reset(self):
64+
"""Keep configuration intact: Decide has no demonstration or training state to reset."""
65+
5966
@staticmethod
6067
def _output_types(signature):
6168
if not signature.output_fields:
@@ -128,7 +135,7 @@ def _questions(self, signature, types):
128135
def _prepare(self, kwargs):
129136
trace = kwargs.pop("_trace", True)
130137
signature = ensure_signature(kwargs.pop("signature", self.signature))
131-
if "demos" in kwargs or self.demos:
138+
if "demos" in kwargs or getattr(self, "demos", None):
132139
raise ValueError("Decide does not support demonstrations.")
133140
client = self.client if self.client is not None else settings.system_one
134141
if client is None:
@@ -204,7 +211,12 @@ def _decode(self, answers, signature, types, inputs, trace):
204211
probabilities=answer["probabilities"],
205212
)
206213
outputs[name] = result if field.annotation is kind else result.value
207-
return self._forward_postprocess([outputs], signature, _trace=trace, **inputs)
214+
prediction = Prediction.from_completions([outputs], signature=signature)
215+
if trace and settings.trace is not None and settings.max_trace_size > 0:
216+
if len(settings.trace) >= settings.max_trace_size:
217+
settings.trace.pop(0)
218+
settings.trace.append((self, inputs, prediction))
219+
return prediction
208220

209221
def forward(self, **kwargs):
210222
client, signature, types, inputs, questions, trace = self._prepare(kwargs)
@@ -217,26 +229,28 @@ async def aforward(self, **kwargs):
217229
return self._decode(answers, signature, types, inputs, trace)
218230

219231
def dump_state(self, json_mode=True):
220-
if self.demos:
232+
if getattr(self, "demos", None):
221233
raise ValueError("Decide does not support demonstrations.")
222-
state = super().dump_state(json_mode=json_mode)
223-
state.pop("demos")
224-
state["thresholds"] = copy.deepcopy(self.thresholds)
225-
state["weights"] = copy.deepcopy(self.weights)
226234
if self.client is not None and not isinstance(self.client, TypeSafe):
227235
raise TypeError(
228236
"Saving an explicit Decide client requires dspy.experimental.TypeSafe; configure custom clients in settings."
229237
)
230-
state["client"] = self.client.dump_state() if self.client is not None else None
231-
return state
238+
return {
239+
"signature": self.signature.dump_state(),
240+
"thresholds": copy.deepcopy(self.thresholds),
241+
"weights": copy.deepcopy(self.weights),
242+
"client": self.client.dump_state() if self.client is not None else None,
243+
}
232244

233245
def load_state(self, state, *, allow_unsafe_lm_state=False):
234246
state = copy.deepcopy(state)
235247
if state.pop("demos", None):
236248
raise ValueError("Decide does not support demonstrations.")
237249
client_state = state.pop("client", None)
238-
super().load_state(state, allow_unsafe_lm_state=allow_unsafe_lm_state)
239-
self.demos = []
250+
self.signature = self.signature.load_state(state["signature"])
251+
self.thresholds = state["thresholds"]
252+
self.weights = state["weights"]
253+
self.__dict__.pop("demos", None)
240254
self.client = TypeSafe(**_sanitize_lm_state(client_state, allow_unsafe_lm_state)) if client_state else None
241255
self._validate_parameters(self._output_types(self.signature))
242256
return self

‎tests/predict/test_decide.py‎

Lines changed: 58 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -331,7 +331,7 @@ def forward(self, text):
331331
program = Pipeline()
332332
client = FakeClient()
333333
callback = Callback()
334-
assert program.named_predictors() == [("decide", program.decide)]
334+
assert program.named_predictors() == []
335335
assert program.named_parameters() == [("decide", program.decide)]
336336
trace = []
337337
with dspy.context(system_one=client, trace=trace, max_trace_size=1, callbacks=[callback]):
@@ -393,7 +393,7 @@ def test_copy_and_json_state_preserve_config(tmp_path):
393393
restored.load(path)
394394
assert restored.weights == module.weights
395395
assert restored.thresholds == module.thresholds
396-
assert "demos" not in module.dump_state()
396+
assert set(module.dump_state()) == {"signature", "thresholds", "weights", "client"}
397397
with dspy.context(system_one=FakeClient()):
398398
assert restored(text="next").toDict() == before.toDict()
399399

@@ -419,10 +419,65 @@ async def test_reject_demonstrations_before_inference_or_saving():
419419
module.dump_state()
420420
assert client.calls == []
421421
module.load_state({**state, "demos": []})
422-
assert module.demos == []
422+
assert not hasattr(module, "demos")
423423
assert module(text="x").flag is True
424424

425425

426+
def test_mixed_program_discovery_optimizer_and_persistence(tmp_path):
427+
from dspy.utils.dummies import DummyLM
428+
429+
class Pipeline(dspy.Module):
430+
def __init__(self):
431+
self.nodes = {"decision": Decide("text -> flag: bool")}
432+
self.explain = dspy.Predict("flag: bool -> explanation: str")
433+
434+
def forward(self, text):
435+
return self.explain(flag=self.nodes["decision"](text=text).flag)
436+
437+
program = Pipeline()
438+
decision = program.nodes["decision"]
439+
decision.thresholds["flag"] = 0.9
440+
assert not isinstance(decision, dspy.Predict)
441+
assert decision.named_parameters() == [("self", decision)]
442+
assert not hasattr(decision, "demos")
443+
assert program.named_parameters() == [("nodes['decision']", decision), ("explain", program.explain)]
444+
assert program.named_predictors() == [("explain", program.explain)]
445+
trained = dspy.LabeledFewShot(k=1).compile(
446+
program, trainset=[dspy.Example(flag=False, explanation="Below threshold").with_inputs("flag")]
447+
)
448+
assert len(trained.explain.demos) == 1
449+
assert not hasattr(trained.nodes["decision"], "demos")
450+
assert trained.nodes["decision"].thresholds == {"flag": 0.9}
451+
path = tmp_path / "mixed.json"
452+
trained.save(path)
453+
restored = Pipeline()
454+
restored.load(path)
455+
assert restored.nodes["decision"].thresholds == {"flag": 0.9}
456+
assert len(restored.explain.demos) == 1
457+
trace = []
458+
with dspy.context(system_one=FakeClient(), lm=DummyLM([{"explanation": "Below threshold"}]), trace=trace):
459+
assert restored(text="x").explanation == "Below threshold"
460+
assert [step[0] for step in trace] == [restored.nodes["decision"], restored.explain]
461+
assert trace[1][1] == {"flag": False}
462+
463+
464+
def test_reset_copy_preserves_configuration_without_aliasing():
465+
module = decide(True, client=TypeSafe("jev-test"))
466+
module.thresholds["flag"] = 0.9
467+
module.weights["rating"] = [-2, 0, 8]
468+
module.weights["label"] = {"other": 3}
469+
reset = module.reset_copy()
470+
assert reset.thresholds == {"flag": 0.9}
471+
assert reset.weights == {"rating": [-2, 0, 8], "label": {"other": 3}}
472+
assert reset.signature is module.signature
473+
assert reset.client.model == "jev-test"
474+
reset.thresholds["flag"] = 0.5
475+
reset.weights["rating"][1] = 1
476+
reset.weights["label"]["other"] = 1
477+
assert module.thresholds == {"flag": 0.9}
478+
assert module.weights == {"rating": [-2, 0, 8], "label": {"other": 3}}
479+
480+
426481
def test_explicit_client_state_omits_key_and_gates_endpoint(tmp_path):
427482
module = decide(client=TypeSafe("jev-test", api_key="test-credential", base_url="https://example.test"))
428483
path = tmp_path / "decide.json"

0 commit comments

Comments
 (0)