shikumi-optimize 0.2.0.0 → 0.2.1.0

raw patch · 26 files changed

+1072/−263 lines, 26 filesdep ~baikaidep ~shikumidep ~shikumi-compilePVP: major bump suggested

API removals or changes: PVP suggests a major version bump

Dependency ranges changed: baikai, shikumi, shikumi-compile, shikumi-eval, shikumi-optimize, shikumi-trace

API changes (from Hackage documentation)

Files

• CHANGELOG.md +20/−0
• shikumi-optimize.cabal +13/−11
• src/Shikumi/Optimize/Bootstrap.hs +41/−12
• src/Shikumi/Optimize/COPRO.hs +34/−33
• src/Shikumi/Optimize/Ensemble.hs +27/−8
• src/Shikumi/Optimize/GEPA.hs +47/−48
• src/Shikumi/Optimize/Instruction.hs +33/−31
• src/Shikumi/Optimize/KNN.hs +10/−3
• src/Shikumi/Optimize/LabeledFewShot.hs +13/−6
• src/Shikumi/Optimize/MIPRO.hs +136/−71
• src/Shikumi/Optimize/RandomSearch.hs +11/−7
• src/Shikumi/Optimize/Search.hs +139/−4
• src/Shikumi/Optimize/Types.hs +24/−5
• test/AcceptanceSpec.hs +60/−4
• test/BootstrapSpec.hs +24/−3
• test/EnsembleSpec.hs +16/−2
• test/GepaSpec.hs +15/−2
• test/InstructionSpec.hs +15/−2
• test/KNNSpec.hs +19/−2
• test/LabeledFewShotSpec.hs +14/−2
• test/Main.hs +4/−0
• test/Miprov2Spec.hs +23/−1
• test/RandomSearchSpec.hs +19/−3
• test/SearchSpec.hs +114/−0
• test/SeedingSpec.hs +164/−0
• test/StubLM.hs +37/−3

@@ -2,6 +2,26 @@  ## Unreleased +## 0.2.1.0 - 2026-07-05++### Added++- Shared optimizer budget metering helpers, effective-instruction inspection, and+  exact LLM-call counting for opaque optimizer runs.+- Budgeted ensemble search via `ensembleSearchWith` and exported+  `bootstrapKeptDemos` for metered bootstrap demo recovery.++### Changed++- Optimizers now reserve predicted LM-call cost before scoring/proposing and+  return the best candidate found so far when the next spend would exceed budget.+- Instruction search and seeding preserve effective instructions without+  serializing redundant overrides.+- Ensemble and KNN documentation now calls out structure-changing artifacts and+  how to load saved state against matching templates.+- Refreshed internal `shikumi`, `shikumi-compile`, `shikumi-eval`, and+  `shikumi-trace` bounds for the current package set.+ ## 0.2.0.0 - 2026-06-28  ### Added

@@ -1,6 +1,6 @@ cabal-version:   3.4 name:            shikumi-optimize-version:         0.2.0.0+version:         0.2.1.0 synopsis:        The optimizer framework for shikumi LM programs (EP-10) category:        AI description:@@ -66,10 +66,10 @@     , effectful     , generic-lens     ^>=2.2     , lens             ^>=5.3-    , shikumi          ^>=0.2.0.0-    , shikumi-compile  ^>=0.1.1.0-    , shikumi-eval     ^>=0.1.1.0-    , shikumi-trace    ^>=0.1.1.0+    , shikumi          ^>=0.3.0.0+    , shikumi-compile  ^>=0.2.0.0+    , shikumi-eval     ^>=0.2.0.0+    , shikumi-trace    ^>=0.2.0.0     , text             ^>=2.1     , vector @@ -92,21 +92,23 @@     OptimizeSpec     ProposeSpec     RandomSearchSpec+    SearchSpec+    SeedingSpec     StubLM    build-depends:     , aeson-    , baikai            >=0.2      && <0.3+    , baikai            >=0.3      && <0.4     , base     , containers     , effectful     , generic-lens     , lens-    , shikumi           ^>=0.2.0.0-    , shikumi-compile   ^>=0.1.1.0-    , shikumi-eval      ^>=0.1.1.0-    , shikumi-optimize  ^>=0.2.0.0-    , shikumi-trace     ^>=0.1.1.0+    , shikumi           ^>=0.3.0.0+    , shikumi-compile   ^>=0.2.0.0+    , shikumi-eval      ^>=0.2.0.0+    , shikumi-optimize  ^>=0.2.1.0+    , shikumi-trace     ^>=0.2.0.0     , tasty     , tasty-hunit     , text

@@ -21,6 +21,7 @@ module Shikumi.Optimize.Bootstrap   ( bootstrapFewShot,     bootstrapFewShotWith,+    bootstrapKeptDemos,     BootstrapConfig (..),     defaultBootstrapConfig,     recoverDemo,@@ -28,14 +29,17 @@ where  import Data.Aeson (ToJSON, toJSON)-import Effectful.Error.Static (catchError)+import Effectful (Eff, (:>))+import Effectful.Error.Static (Error, catchError)+import Effectful.Prim (Prim) import GHC.Generics (Generic) import Shikumi.Error (ShikumiError)-import Shikumi.Eval (Example (..), datasetExamples, prediction, unScore)+import Shikumi.Eval (Dataset, Example (..), Metric, datasetExamples, prediction, unScore)+import Shikumi.LLM (LLM) import Shikumi.Optimize.LabeledFewShot (withDemos)-import Shikumi.Optimize.Search (freezeProgram)+import Shikumi.Optimize.Search (BudgetMeter, freezeProgram, newBudgetMeter, tryCharge) import Shikumi.Optimize.Types (Budget (..), Optimizer (..))-import Shikumi.Program (Demo (..), Program, runProgram)+import Shikumi.Program (Demo (..), Program, foldParams, runProgram)  -- | Tunables for a bootstrap search. data BootstrapConfig = BootstrapConfig@@ -64,7 +68,10 @@  -- | Bootstrap few-shot with an explicit configuration. The @teacher@ may be a -- stronger or chain-of-thought variant of the student, or the student itself; it--- must share the student's input/output types.+-- must share the student's input/output types. Each teacher run reserves one+-- predicted LM completion per teacher predict node before it runs; when the next+-- teacher run does not fit the 'Budget', demo recovery stops and the demos found so+-- far are attached. bootstrapFewShotWith ::   (ToJSON i, ToJSON o) =>   BootstrapConfig ->@@ -73,11 +80,22 @@   Budget ->   Optimizer i o bootstrapFewShotWith cfg teacher budget = Optimizer $ \train metric student -> do-  -- Bound the number of teacher runs by the budget (each run is >= 1 LM call).-  let exs = take (max 0 (maxLmCalls budget)) (datasetExamples train)-      -- Run the teacher on one example; if it succeeds and the metric passes,-      -- emit its recovered demo, else emit nothing. Recovery is total: a teacher-      -- error yields no demo rather than aborting the search.+  meter <- newBudgetMeter budget+  kept <- bootstrapKeptDemos cfg meter teacher train metric+  pure (freezeProgram (withDemos kept student))++-- | Recover metric-passing demos from teacher runs under a shared budget meter.+bootstrapKeptDemos ::+  (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es, Prim :> es) =>+  BootstrapConfig ->+  BudgetMeter ->+  Program i o ->+  Dataset i o ->+  Metric o ->+  Eff es [Demo]+bootstrapKeptDemos cfg meter teacher train metric = do+  let teacherCost = max 1 (length (foldParams teacher))+      cap = max 0 (maxBootstrappedDemos cfg)       keepIfPassing (Example inp expd) =         ( do             out <- runProgram teacher inp@@ -85,5 +103,16 @@             pure [recoverDemo inp out | s >= passThreshold cfg]         )           `catchError` \_ (_ :: ShikumiError) -> pure []-  kept <- concat <$> mapM keepIfPassing exs-  pure (freezeProgram (withDemos (take (maxBootstrappedDemos cfg) kept) student))+      collect kept []+        | length kept >= cap = pure (take cap kept)+        | otherwise = pure kept+      collect kept (ex : rest)+        | length kept >= cap = pure (take cap kept)+        | otherwise = do+            fits <- tryCharge meter teacherCost+            if not fits+              then pure kept+              else do+                newKept <- keepIfPassing ex+                collect (kept ++ newKept) rest+  collect [] (datasetExamples train)

@@ -9,7 +9,8 @@ -- COPRO consumes EP-19's grounded proposer ('Shikumi.Optimize.Propose.proposeInstructions') -- directly: each round's call passes the node's current instruction and its scored -- 'PastInstruction' history, and the proposer returns ranked candidates with the--- current instruction always retained (the safety property — a node never degrades).+-- current effective instruction always retained. Keeping that candidate writes no+-- redundant override, preserving the safety property that a node never degrades. -- -- Output is V1's 'Shikumi.Compile.Types.CompiledProgram' via 'freezeProgram', invoked -- through 'Shikumi.Optimize.optimize' and serialized unchanged (integration point #4).@@ -22,7 +23,6 @@  import Control.Monad (foldM) import Data.Aeson (ToJSON)-import Data.Maybe (fromMaybe) import Effectful (Eff, (:>)) import Effectful.Concurrent (Concurrent) import Effectful.Error.Static (Error)@@ -30,7 +30,7 @@ import GHC.Generics (Generic) import Shikumi.Effect.Time (Time) import Shikumi.Error (ShikumiError)-import Shikumi.Eval (Dataset, Metric, datasetSize)+import Shikumi.Eval (Dataset, Metric) import Shikumi.LLM (LLM) import Shikumi.Optimize.Propose   ( PastInstruction (..),@@ -38,7 +38,7 @@     ProposeResult (..),     proposeInstructions,   )-import Shikumi.Optimize.Search (freezeProgram, instructionAt, scoreOn, setNodeInstr)+import Shikumi.Optimize.Search (BudgetMeter, effectiveInstructionAt, freezeProgram, meteredScore, newBudgetMeter, setNodeInstrIfNew, tryCharge) import Shikumi.Optimize.Types (Budget (..), Optimizer (..), defaultBudget) import Shikumi.Program (Program, foldParams) @@ -59,19 +59,22 @@  -- | Coordinate-ascent instruction optimization. Visits each node in @foldParams@ -- order, optimizing it over @depth@ rounds against the already-improved earlier--- nodes, threading one running LM-call count so the 'Budget' bounds the whole search.+-- nodes. Proposer calls and candidate scoring reserve their predicted cost through+-- one shared 'Budget', so the search returns the best-so-far before the next spend+-- would exceed either ceiling. copro :: (ToJSON i, ToJSON o) => CoproConfig -> Optimizer i o copro cfg = Optimizer $ \train metric student -> do+  meter <- newBudgetMeter (budget cfg)   let nNodes = length (foldParams student)-  (final, _calls) <-+  final <-     foldM-      (\acc idx -> optimizeNode cfg train metric idx acc)-      (student, 0)+      (\acc idx -> optimizeNode cfg meter train metric idx acc)+      student       [0 .. nNodes - 1]   pure (freezeProgram final)  -- | Optimize node @idx@ over @depth@ rounds. Returns the program with node @idx@ set--- to its best-found instruction, plus the updated LM-call count. Each round proposes+-- to its best-found instruction. Each round proposes -- @breadth - 1@ fresh candidates (plus the retained current instruction) via the -- grounded proposer fed the scored attempt history, scores the not-yet-seen ones on -- the whole training set, records @(instruction, best-score)@, and sets the node to@@ -79,27 +82,26 @@ optimizeNode ::   (ToJSON i, ToJSON o, LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) =>   CoproConfig ->+  BudgetMeter ->   Dataset i o ->   Metric o ->   Int ->-  (Program i o, Int) ->-  Eff es (Program i o, Int)-optimizeNode cfg train metric idx (prog0, calls0) = goRound 1 prog0 calls0 []+  Program i o ->+  Eff es (Program i o)+optimizeNode cfg meter train metric idx prog0 = goRound 1 prog0 []   where-    dsSize = datasetSize train     bdth = max 2 (breadth cfg)     dpth = max 1 (depth cfg)     proposerCost = 4 + (bdth - 1)-    maxCalls = maxLmCalls (budget cfg)-    maxCands = maxCandidates (budget cfg) -    goRound r prog calls evald-      | r > dpth = pure (setBest prog evald, calls)+    goRound r prog evald+      | r > dpth = pure (setBest prog evald)       | otherwise = do-          let cur = fromMaybe "" (instructionAt idx prog)+          let cur = effectiveInstructionAt idx prog               hist = [PastInstruction i s | (i, s) <- evald]-          (cands, calls1) <--            if calls + proposerCost <= maxCalls+          fitsProposer <- tryCharge meter proposerCost+          cands <-+            if fitsProposer               then do                 ProposeResult cs <-                   proposeInstructions@@ -114,22 +116,21 @@                         tipIndex = 1,                         viewBatch = 2                       }-                pure (cs, calls + proposerCost)-              else pure ([cur], calls)-          (evald', calls2) <- scoreNew prog calls1 evald (dedupNew evald cands)+                pure cs+              else pure [cur]+          evald' <- scoreNew prog evald (dedupNew evald cands)           -- set the node to the best instruction found so far, then continue           let prog' = setBest prog evald'-          goRound (r + 1) prog' calls2 evald'+          goRound (r + 1) prog' evald' -    -- Score the not-yet-evaluated candidates, threading the call count and stopping+    -- Score the not-yet-evaluated candidates, using the shared meter and stopping     -- before either Budget ceiling would be exceeded.-    scoreNew _ calls evald [] = pure (evald, calls)-    scoreNew prog calls evald (c : cs)-      | calls + dsSize > maxCalls = pure (evald, calls)-      | length evald >= maxCands = pure (evald, calls)-      | otherwise = do-          s <- scoreOn train metric (setNodeInstr idx c prog)-          scoreNew prog (calls + dsSize) (evald ++ [(c, s)]) cs+    scoreNew _ evald [] = pure evald+    scoreNew prog evald (c : cs) = do+      ms <- meteredScore meter train metric (setNodeInstrIfNew idx c prog)+      case ms of+        Nothing -> pure evald+        Just s -> scoreNew prog (evald ++ [(c, s)]) cs      -- Candidates not already scored, de-duplicated against each other (order-preserving).     dedupNew evald = go (map fst evald)@@ -142,6 +143,6 @@     -- Set node idx to the highest-scoring instruction recorded (earliest on ties).     setBest prog evald = case evald of       [] -> prog-      (e : es) -> setNodeInstr idx (fst (foldl' pick e es)) prog+      (e : es) -> setNodeInstrIfNew idx (fst (foldl' pick e es)) prog       where         pick best c = if snd c > snd best then c else best

@@ -9,8 +9,16 @@ -- -- Resampling is deterministic (a fixed linear-congruential stream seeded by the -- member index), so the search is reproducible run to run.+--+-- The returned artifact is structure-changing: it is an @Ensemble@ combinator over+-- the optimized members, not just a parameter rewrite of the student. The member+-- parameters are persisted by 'Shikumi.Compile.Serialize.encodeCompiled', but the+-- reducer closure and exact member structure live in the program template held in+-- code. Load saved state onto the matching ensemble template, not onto the plain+-- student. module Shikumi.Optimize.Ensemble   ( ensembleSearch,+    ensembleSearchWith,     majorityReducer,   ) where@@ -19,18 +27,29 @@ import Shikumi.Combinator (ensemble) import Shikumi.Compile.Types (compiledProgram) import Shikumi.Eval (Dataset, dataset, datasetExamples)-import Shikumi.Optimize.Search (freezeProgram)-import Shikumi.Optimize.Types (Optimizer (..))+import Shikumi.Optimize.Search (freezeProgram, withLmCallCount)+import Shikumi.Optimize.Types (Budget (..), Optimizer (..), defaultBudget)  -- | Build an @size@-member ensemble: run @inner@ on @size@ bootstrap resamples of--- the training set and combine the resulting programs by majority vote.+-- the training set and combine the resulting programs by majority vote. This+-- changes structure by returning an @Ensemble@; saved state must be decoded onto a+-- matching ensemble template. ensembleSearch :: (Eq o) => Int -> Optimizer i o -> Optimizer i o-ensembleSearch size inner = Optimizer $ \train metric student -> do+ensembleSearch = ensembleSearchWith defaultBudget++-- | Build a budgeted ensemble. The budget is enforced between members using exact+-- LLM-call counting; at least one member always runs, so the final spend may exceed+-- the bound by one member's cost.+ensembleSearchWith :: (Eq o) => Budget -> Int -> Optimizer i o -> Optimizer i o+ensembleSearchWith budget size inner = Optimizer $ \train metric student -> do   let seeds = [1 .. max 1 size]-  members <--    mapM-      (\seed -> compiledProgram <$> runOptimizer inner (resample seed train) metric student)-      seeds+      go _total acc [] = pure (reverse acc)+      go total acc (seed : rest)+        | total >= maxLmCalls budget && not (null acc) = pure (reverse acc)+        | otherwise = do+            (cp, calls) <- withLmCallCount (runOptimizer inner (resample seed train) metric student)+            go (total + calls) (compiledProgram cp : acc) rest+  members <- go 0 [] seeds   pure (freezeProgram (ensemble members majorityReducer))  -- | Sample a dataset with replacement, deterministically, seeded by @seed@. An

@@ -33,10 +33,7 @@   ) where -import Control.Lens ((&), (?~)) import Control.Monad (forM, forM_, when)-import Data.Generics.Labels ()-import Data.Maybe (fromMaybe) import Data.Text (Text) import Data.Text qualified as T import Effectful (Eff, (:>))@@ -61,14 +58,12 @@ import Shikumi.LLM (LLM) import Shikumi.Module (predict) import Shikumi.Optimize.Pareto (Candidate (..), paretoFrontier, sampleParent)-import Shikumi.Optimize.Search (freezeProgram)+import Shikumi.Optimize.Search (effectiveInstructionAt, freezeProgram, newBudgetMeter, scoringCost, setNodeInstrIfNew, tryCharge) import Shikumi.Optimize.Types (Budget (..), Optimizer (..)) import Shikumi.Program   ( NodeFields (..),-    Params (..),     Program,     foldParams,-    mapParamsAt,     nodeFieldsIndexed,     runProgram,     setProgramParams,@@ -176,18 +171,12 @@        in if null crits             then pure prog             else do-              let cur = fromMaybe "" (instructionOverride (paramsAt idx prog))+              let cur = effectiveInstructionAt idx prog                   fldSummary = renderFields (drop idx fields)                   fb = T.intercalate "\n" crits               ReflectOut newInstr <-                 runProgram proposer (ReflectIn cur fb progSummary dataSummary fldSummary)-              pure (mapParamsAt idx (\ps -> ps & #instructionOverride ?~ newInstr) prog)---- | The 'Params' at a node index (empty if out of range).-paramsAt :: Int -> Program i o -> Params-paramsAt idx prog = case drop idx (foldParams prog) of-  (ps : _) -> ps-  [] -> Params Nothing []+              pure (setNodeInstrIfNew idx newInstr prog)  -- | Render a node's field names for the proposer prompt. renderFields :: [NodeFields] -> Text@@ -203,54 +192,66 @@  -- | The reflective evolutionary optimizer. Takes its reflective proposer and feedback -- metric explicitly (so it is testable under a stub LM) and returns V1's--- 'Optimizer'.+-- 'Optimizer'. GEPA gates its seed evaluation before any LM call; if the budget is+-- too small to score the student once, it returns the student unscored. Each+-- evolution step reserves a conservative full-step cost before capture, reflection,+-- and child scoring. gepa ::   Program ReflectIn ReflectOut ->   FeedbackMetric o ->   Budget ->   Optimizer i o gepa proposer fbMetric budget = Optimizer $ \train metric student -> do+  meter <- newBudgetMeter budget   let paths = programNodePaths student       fields = nodeFieldsIndexed student       nNodes = max 1 (length paths)-      n = max 1 (datasetSize train)       progSummary = fallbackProgramSummary (length paths)       dataSummary = fallbackDatasetSummary (datasetSize train)-      maxCalls = maxLmCalls budget       maxCands = maxCandidates budget       rebuild cand = either (const student) id (setProgramParams (params cand) student)+      seedCost = scoringCost train student -  seedRpt <- evaluatePure train metric student-  let seedCand = Candidate (foldParams student) (perEx seedRpt) (aggregateScore seedRpt)+  seedFits <- tryCharge meter seedCost+  if not seedFits+    then pure (freezeProgram student)+    else do+      seedRpt <- evaluatePure train metric student+      let seedCand = Candidate (foldParams student) (perEx seedRpt) (aggregateScore seedRpt) -      -- A full step costs: capture (n) + mutate (1) + child eval (n) = 2n + 1.-      stepCost = 2 * n + 1-      stepCap = maxCands + 4+          -- A full step costs: capture + child evaluation over the whole dataset,+          -- plus one reflective proposer call.+          stepCost = 2 * seedCost + 1+          stepCap = maxCands + 4 -      loop step calls cands seed frontier-        | step >= stepCap = pure (bestOf seedCand frontier)-        | calls + stepCost > maxCalls = pure (bestOf seedCand frontier)-        | length cands >= maxCands = pure (bestOf seedCand frontier)-        | otherwise = case sampleParent seed (paretoFrontier frontier) of-            Nothing -> pure (bestOf seedCand frontier)-            Just (parent, seed') -> do-              let parentProg = rebuild parent-                  idx = step `mod` nNodes-              (fblog, _) <- captureFeedback train fbMetric parentProg-              case drop idx paths of-                (path : _)-                  | null (feedbackFor path fblog) ->-                      -- nothing to reflect on at this node; advance (capture cost only)-                      loop (step + 1) (calls + n) cands seed' frontier-                _ -> do-                  child <- mutateNode proposer progSummary dataSummary fields fblog paths idx parentProg-                  rpt <- evaluatePure train metric child-                  let childCand = Candidate (foldParams child) (perEx rpt) (aggregateScore rpt)-                      frontier' = paretoFrontier (childCand : frontier)-                  loop (step + 1) (calls + stepCost) (childCand : cands) seed' frontier'+          loop step cands seed frontier+            | step >= stepCap = pure (bestOf seedCand frontier)+            | length cands >= maxCands = pure (bestOf seedCand frontier)+            | otherwise = do+                fitsStep <- tryCharge meter stepCost+                if not fitsStep+                  then pure (bestOf seedCand frontier)+                  else case sampleParent seed (paretoFrontier frontier) of+                    Nothing -> pure (bestOf seedCand frontier)+                    Just (parent, seed') -> do+                      let parentProg = rebuild parent+                          idx = step `mod` nNodes+                      (fblog, _) <- captureFeedback train fbMetric parentProg+                      case drop idx paths of+                        (path : _)+                          | null (feedbackFor path fblog) ->+                              -- nothing to reflect on at this node; the reserved+                              -- full-step budget is a conservative upper bound.+                              loop (step + 1) cands seed' frontier+                        _ -> do+                          child <- mutateNode proposer progSummary dataSummary fields fblog paths idx parentProg+                          rpt <- evaluatePure train metric child+                          let childCand = Candidate (foldParams child) (perEx rpt) (aggregateScore rpt)+                              frontier' = paretoFrontier (childCand : frontier)+                          loop (step + 1) (childCand : cands) seed' frontier' -  best <- loop 0 n [seedCand] 1 [seedCand]-  pure (freezeProgram (rebuild best))+      best <- loop 0 [seedCand] 1 [seedCand]+      pure (freezeProgram (rebuild best))  -- | The per-example score vector from a report, in dataset order. perEx :: Report -> [Double]@@ -259,9 +260,7 @@ -- | The frontier candidate with the highest aggregate (earliest on ties); falls back -- to the seed if the frontier is somehow empty. bestOf :: Candidate -> [Candidate] -> Candidate-bestOf seedCand cs = case cs of-  [] -> seedCand-  (x : xs) -> foldl' (\b c -> if aggregate c > aggregate b then c else b) x xs+bestOf seedCand = foldl' (\b c -> if aggregate c > aggregate b then c else b) seedCand  -- | A minimal program summary (EP-19's program describer is the richer source). fallbackProgramSummary :: Int -> Text

@@ -9,16 +9,18 @@ -- The proposer and its signal-gatherers are themselves ordinary shikumi 'Program's, so -- they are typed, cached, traced, and testable with the same stub-LM machinery as -- everything else — the optimizer is written in the framework it optimizes. The--- /current/ instruction is always retained as a candidate (the proposer guarantees it),--- so a node can never end up worse than where it started.+-- /current/ effective instruction (override, or signature base when no override is+-- present) is always retained as a candidate (the proposer guarantees it), and keeping+-- that candidate writes no redundant override, so a node can never end up worse than+-- where it started. ----- __Budget.__ The grounded proposer makes @4 + proposalsPerNode@ LM calls per node--- (dataset summary, program describe, module describe, and one generation per proposal);--- scoring one candidate costs one LM call per dataset example. The search threads a--- running raw-call count and stops — returning the best found /so far/ — before either--- bound in the 'Budget' would be exceeded, so the recorded LM-call count never exceeds--- @maxLmCalls@. When the remaining budget cannot cover a node's full proposal, that node--- keeps its current instruction (no proposer call) rather than partially proposing.+-- __Budget.__ The grounded proposer reserves @4 + proposalsPerNode@ predicted LM+-- completions per node (dataset summary, program describe, module describe, and one+-- generation per proposal); scoring one candidate reserves one completion per dataset+-- example per predict node. The search stops — returning the best found /so far/ —+-- before either bound in the 'Budget' would be exceeded. When the remaining budget+-- cannot cover a node's full proposal, that node keeps its current instruction (no+-- proposer call) rather than partially proposing. -- -- This module re-points V1's blind proposer at EP-19's grounded surface; the old -- @ProposeIn@/@ProposeOut@/@proposeInstruction@ predictor is removed (the grounded@@ -31,14 +33,12 @@  import Control.Monad (foldM) import Data.Aeson (ToJSON)-import Data.Maybe (fromMaybe)-import Shikumi.Eval (datasetSize) import Shikumi.Optimize.Propose   ( ProposeRequest (..),     ProposeResult (..),     proposeInstructions,   )-import Shikumi.Optimize.Search (freezeProgram, instructionAt, scoreOn, setNodeInstr)+import Shikumi.Optimize.Search (effectiveInstructionAt, freezeProgram, meteredScore, newBudgetMeter, setNodeInstrIfNew, tryCharge) import Shikumi.Optimize.Types (Budget (..), Optimizer (..)) import Shikumi.Program (foldParams) @@ -46,29 +46,31 @@ -- an explicit LM-call budget, using the grounded proposer to generate candidates. instructionSearch :: (ToJSON i, ToJSON o) => Int -> Budget -> Optimizer i o instructionSearch proposalsPerNode budget = Optimizer $ \train metric student -> do-  let dsSize = datasetSize train-      nNodes = length (foldParams student)+  meter <- newBudgetMeter budget+  let nNodes = length (foldParams student)       -- The grounded proposer's fixed per-node LM-call cost (see module header).       proposerCost = 4 + max 0 proposalsPerNode -      -- Score candidate instructions for node @idx@, threading the call count and+      -- Score candidate instructions for node @idx@, using the shared meter and       -- the best-so-far; stop before the next scoring would exceed the budget.-      scoreCands calls best idx prog cs = case cs of-        [] -> pure (best, calls)-        (c : rest)-          | calls + dsSize > maxLmCalls budget -> pure (best, calls)-          | otherwise -> do-              s <- scoreOn train metric (setNodeInstr idx c prog)+      scoreCands best idx prog cs = case cs of+        [] -> pure best+        (c : rest) -> do+          ms <- meteredScore meter train metric (setNodeInstrIfNew idx c prog)+          case ms of+            Nothing -> pure best+            Just s -> do               let best' = case best of                     Nothing -> Just (c, s)                     Just (_, bs) -> if s > bs then Just (c, s) else best-              scoreCands (calls + dsSize) best' idx prog rest+              scoreCands best' idx prog rest        -- Optimize one node, holding the others fixed.-      stepNode (prog, calls) idx = do-        let curInstr = fromMaybe "" (instructionAt idx prog)-        (cands, calls1) <--          if calls + proposerCost <= maxLmCalls budget+      stepNode prog idx = do+        let curInstr = effectiveInstructionAt idx prog+        fitsProposer <- tryCharge meter proposerCost+        cands <-+          if fitsProposer             then do               ProposeResult cs <-                 proposeInstructions@@ -83,10 +85,10 @@                       tipIndex = 0,                       viewBatch = 2                     }-              pure (cs, calls + proposerCost)-            else pure ([curInstr], calls)-        (best, calls2) <- scoreCands calls1 Nothing idx prog cands-        pure (setNodeInstr idx (maybe curInstr fst best) prog, calls2)+              pure cs+            else pure [curInstr]+        best <- scoreCands Nothing idx prog cands+        pure (setNodeInstrIfNew idx (maybe curInstr fst best) prog) -  (final, _) <- foldM stepNode (student, 0) [0 .. nNodes - 1]+  final <- foldM stepNode student [0 .. nNodes - 1]   pure (freezeProgram final)

@@ -13,7 +13,8 @@ -- EP-15's pure @runEmbedding@ argument), not the @Embedding@ effect: an 'Embed' -- body's row is fixed to @(LLM, Error ShikumiError)@, so it cannot call @embedText@; -- all embedding-effect work happens at the caller, outside the node. The run-time--- form carries no @Params@ (it serializes as the empty vector, exactly like @react@).+-- form carries no @Params@ (it serializes as an @Embed@ shape with an empty vector,+-- exactly like @react@). module Shikumi.Optimize.KNN   ( -- * Run-time form (the faithful analog)     knnFewShot,@@ -100,7 +101,13 @@    in embed $ \i -> runProgram (withDemos (nearestDemos embedder k exs (toPrompt i)) student) i  -- | Run-time KNN as an 'Optimizer': selection is by embedding geometry, not by--- score, so it consults neither the metric nor the LM at optimize time.+-- score, so it consults neither the metric nor the LM at optimize time and spends+-- zero optimizer LM calls. The result is a structure-changing @Embed@ wrapper+-- around the student. Its run-time selector closure is not persisted by+-- 'Shikumi.Compile.Serialize.encodeCompiled', so the saved state must be decoded+-- onto the same @knnDemos@ template; decode onto the plain student fails with a+-- shape mismatch. Use 'knnFewShotCentroid' when the desired artifact is only+-- serializable node parameters. knnFewShot ::   (ToJSON i, ToJSON o, ToPrompt i) =>   (Text -> Vector Double) ->@@ -111,7 +118,7 @@  -- | Compile-time fallback: bake the @k@ training examples nearest the training -- centroid as a fixed demo set on every node, and freeze. A plain @Params@ artifact,--- no run-time embedder needed.+-- no run-time embedder needed, and zero optimizer LM calls are spent. knnFewShotCentroid ::   (ToJSON i, ToJSON o, ToPrompt i) =>   (Text -> Vector Double) ->

@@ -5,10 +5,12 @@ -- -- No LM calls are made beyond scoring. Candidate sets are enumerated -- deterministically (all size-@k@ combinations of the training demos, in a fixed--- order, bounded by the budget) rather than randomly sampled, so the result is--- reproducible run to run — the tests rely on this.+-- order) rather than randomly sampled. 'labeledFewShotWith' reserves one scoring+-- cost per candidate before evaluation and stops at the shared 'Budget'; the+-- default 'labeledFewShot' uses 'defaultBudget'. module Shikumi.Optimize.LabeledFewShot   ( labeledFewShot,+    labeledFewShotWith,     labeledCandidateSets,     withDemos,   )@@ -18,16 +20,21 @@ import Data.Aeson (ToJSON, toJSON) import Data.Generics.Labels () import Shikumi.Eval (Dataset, Example (..), datasetExamples)-import Shikumi.Optimize.Search (freezeProgram, scoreOn, selectBest)-import Shikumi.Optimize.Types (Optimizer (..), Scored (..), defaultBudget)+import Shikumi.Optimize.Search (freezeProgram, meteredScore, newBudgetMeter, selectBestMetered)+import Shikumi.Optimize.Types (Budget, Optimizer (..), Scored (..), defaultBudget) import Shikumi.Program (Demo (..), Program, mapParams)  -- | Select the best size-@k@ set of labelled demonstrations from the training set. -- Multi-node programs receive the same demo set at every node (the DSPy default). labeledFewShot :: (ToJSON i, ToJSON o) => Int -> Optimizer i o-labeledFewShot k = Optimizer $ \train metric prog -> do+labeledFewShot = labeledFewShotWith defaultBudget++-- | Select the best size-@k@ labelled demo set under an explicit budget.+labeledFewShotWith :: (ToJSON i, ToJSON o) => Budget -> Int -> Optimizer i o+labeledFewShotWith budget k = Optimizer $ \train metric prog -> do+  meter <- newBudgetMeter budget   let sets = labeledCandidateSets k train-  best <- selectBest defaultBudget (\ds -> scoreOn train metric (withDemos ds prog)) sets+  best <- selectBestMetered meter (\ds -> meteredScore meter train metric (withDemos ds prog)) sets   pure $ case best of     Nothing -> freezeProgram prog     Just sc -> freezeProgram (withDemos (candidate sc) prog)

@@ -5,8 +5,9 @@ -- one node at a time, MIPROv2 searches /both/ axes — instruction × demo set — across -- all nodes jointly, in three phases: -----   1. __bootstrap__ candidate demo sets (per node) from the teacher's passing runs---      plus the labelled training pairs ('bootstrapDemoCandidates');+--   1. __bootstrap__ candidate demo sets (per node), retaining the node's current+--      demos at index 0, then adding the empty set, teacher passing runs, and the+--      labelled training pairs ('bootstrapDemoCandidates'); --   2. __propose__ candidate instructions (per node) through EP-19's grounded --      proposer ('proposeInstructionCandidates'); and --   3. __search__ the joint @(instruction × demoset)@ grid with cheap minibatch@@ -44,7 +45,6 @@ import Control.Monad (forM) import Data.Aeson (ToJSON) import Data.Aeson.Text (encodeToLazyText)-import Data.Maybe (fromMaybe) import Data.Text (Text) import Data.Text.Lazy qualified as TL import Effectful (Eff, (:>))@@ -71,7 +71,7 @@     ProposeResult (..),     proposeInstructions,   )-import Shikumi.Optimize.Search (freezeProgram, scoreOn)+import Shikumi.Optimize.Search (BudgetMeter, effectiveInstructionAt, freezeProgram, meteredScore, newBudgetMeter, setNodeInstrIfNew, tryCharge) import Shikumi.Optimize.Types (Budget (..), Optimizer (..), defaultBudget) import Shikumi.Program   ( Demo (..),@@ -95,7 +95,7 @@ data Miprov2Config = Miprov2Config   { -- | instruction candidates proposed per node (incl. the current at index 0)     numInstructCandidates :: !Int,-    -- | demo-set candidates per node (incl. the empty set at index 0)+    -- | demo-set candidates per node (incl. the node's current demos at index 0)     numDemoCandidates :: !Int,     -- | minibatch trials the search performs     numTrials :: !Int,@@ -109,7 +109,7 @@     maxBootstrappedDemos :: !Int,     -- | min metric score for a teacher run to contribute demos     bootstrapThreshold :: !Double,-    -- | hard LM-call / candidate ceiling (V1's 'Budget')+    -- | hard predicted LM-completion / candidate ceiling (V1's 'Budget')     budget :: !Budget   }   deriving stock (Eq, Show, Generic)@@ -134,9 +134,10 @@ -- | MIPROv2 with an explicit config and an explicit teacher program. miprov2With :: (ToJSON i, ToJSON o) => Miprov2Config -> Program i o -> Optimizer i o miprov2With cfg teacher = Optimizer $ \train metric student -> do-  demoCands <- bootstrapDemoCandidates cfg teacher train metric student-  instrCands <- proposeInstructionCandidates cfg student train demoCands-  best <- searchJoint cfg train metric student instrCands demoCands+  meter <- newBudgetMeter (budget cfg)+  demoCands <- bootstrapDemoCandidatesWith meter cfg teacher train metric student+  instrCands <- proposeInstructionCandidatesWith meter cfg student train demoCands+  best <- searchJointWith meter cfg train metric student instrCands demoCands   pure (freezeProgram best)  -- ---------------------------------------------------------------------------@@ -144,14 +145,12 @@ -- ---------------------------------------------------------------------------  -- | For each node (in @foldParams@ order), a list of candidate demo sets. Candidate--- 0 is always the empty set (so "no demos" is always reachable). The remaining sets--- come from the teacher's metric-passing runs (bootstrapped) and from the labelled--- training pairs (labelled demos, DSPy's @max_labeled_demos@) — recovered at the--- program-I/O level and attached to every node (the documented degradation while--- per-node trace recovery via EP-16 is deferred; for single-node programs the two--- coincide).+-- 0 is the node's current demos, making the baseline vector an identity. The empty+-- set remains reachable after that, followed by teacher metric-passing runs+-- (bootstrapped) and labelled training pairs (DSPy's @max_labeled_demos@) recovered at+-- the program-I/O level and attached to every node. bootstrapDemoCandidates ::-  (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es) =>+  (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es, Prim :> es) =>   Miprov2Config ->   -- | teacher   Program i o ->@@ -161,7 +160,23 @@   Program i o ->   Eff es [[[Demo]]] bootstrapDemoCandidates cfg teacher train metric student = do+  meter <- newBudgetMeter (budget cfg)+  bootstrapDemoCandidatesWith meter cfg teacher train metric student++bootstrapDemoCandidatesWith ::+  (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es, Prim :> es) =>+  BudgetMeter ->+  Miprov2Config ->+  -- | teacher+  Program i o ->+  Dataset i o ->+  Metric o ->+  -- | student (for the node count)+  Program i o ->+  Eff es [[[Demo]]]+bootstrapDemoCandidatesWith meter cfg teacher train metric student = do   let exs = datasetExamples train+      teacherCost = max 1 (length (foldParams teacher))       keepIfPassing (Example inp expd) =         ( do             out <- runProgram teacher inp@@ -169,13 +184,23 @@             pure [recoverDemo inp out | s >= bootstrapThreshold cfg]         )           `catchError` \_ (_ :: ShikumiError) -> pure []-  bootstrapped <- concat <$> mapM keepIfPassing exs+      collect [] = pure []+      collect (ex : rest) = do+        fits <- tryCharge meter teacherCost+        if not fits+          then pure []+          else do+            kept <- keepIfPassing ex+            (kept ++) <$> collect rest+  bootstrapped <- collect exs   let cap = max 1 (maxBootstrappedDemos cfg)       labeledSet = take cap (map (\(Example i o) -> recoverDemo i o) exs)       bootSet = take cap bootstrapped-      sets = take (max 1 (numDemoCandidates cfg)) ([] : filter (not . null) [labeledSet, bootSet])-      nNodes = length (foldParams student)-  pure (replicate nNodes sets)+      nodeSets ps =+        take+          (max 1 (numDemoCandidates cfg))+          (dedup (demos ps : [] : filter (not . null) [labeledSet, bootSet]))+  pure (map nodeSets (foldParams student))  -- --------------------------------------------------------------------------- -- Phase 2 — propose instruction candidates@@ -186,30 +211,48 @@ -- Uses EP-19's 'proposeInstructions'; the per-candidate tip starts at the /creative/ -- tip (@tipIndex = 1@) so a node always gets a strong proposal even at small N. proposeInstructionCandidates ::-  (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es) =>+  (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es, Prim :> es) =>   Miprov2Config ->   Program i o ->   Dataset i o ->   [[[Demo]]] ->   Eff es [[Text]]-proposeInstructionCandidates cfg student train demoCands =-  forM (zip3 [0 ..] (foldParams student) demoCands) $ \(k, ps, nodeDemoSets) -> do-    let curInstr = fromMaybe "" (instructionOverride ps)+proposeInstructionCandidates cfg student train demoCands = do+  meter <- newBudgetMeter (budget cfg)+  proposeInstructionCandidatesWith meter cfg student train demoCands++proposeInstructionCandidatesWith ::+  (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es, Prim :> es) =>+  BudgetMeter ->+  Miprov2Config ->+  Program i o ->+  Dataset i o ->+  [[[Demo]]] ->+  Eff es [[Text]]+proposeInstructionCandidatesWith meter cfg student train demoCands =+  forM (zip [0 ..] demoCands) $ \(k, nodeDemoSets) -> do+    let curInstr = effectiveInstructionAt k student         demoTexts = map renderDemo (concat (take 1 (filter (not . null) nodeDemoSets)))-    ProposeResult cs <--      proposeInstructions-        train-        ProposeRequest-          { program = student,-            targetNode = k,-            currentInstruction = curInstr,-            history = [],-            bootstrappedDemos = demoTexts,-            numCandidates = max 0 (numInstructCandidates cfg - 1),-            tipIndex = 1,-            viewBatch = 2-          }-    pure cs+        numProposals = max 0 (numInstructCandidates cfg - 1)+        proposerCost = 4 + numProposals+    fits <- tryCharge meter proposerCost+    if not fits+      then pure [curInstr]+      else do+        ProposeResult cs <-+          proposeInstructions+            train+            ProposeRequest+              { program = student,+                targetNode = k,+                currentInstruction = curInstr,+                history = [],+                bootstrappedDemos = demoTexts,+                numCandidates = numProposals,+                tipIndex = 1,+                viewBatch = 2+              }+        pure cs  -- | Render a recovered 'Demo' as @<input-json> => <output-json>@ for the proposal -- prompt's demo signal.@@ -227,9 +270,11 @@  -- | Search the joint per-node @(instruction × demoset)@ grid by greedy coordinate -- descent with minibatch screening, returning the best program found within the--- 'Budget'. Each trial screens the one-coordinate neighbours of the running best on a--- seeded minibatch, then full-evaluates the best-screened neighbour and accepts it--- only if it strictly improves the full score.+-- 'Budget'. Bootstrap teacher runs, grounded proposer calls, minibatch scoring, and+-- full scoring all reserve predicted cost against one meter in 'miprov2With'. Each+-- trial screens the one-coordinate neighbours of the running best on a seeded+-- minibatch, then full-evaluates the best-screened neighbour and accepts it only if+-- it strictly improves the full score. searchJoint ::   (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) =>   Miprov2Config ->@@ -242,53 +287,73 @@   [[[Demo]]] ->   Eff es (Program i o) searchJoint cfg train metric student instrCands demoCands = do+  meter <- newBudgetMeter (budget cfg)+  searchJointWith meter cfg train metric student instrCands demoCands++-- | Meter-aware implementation of 'searchJoint'. Later MIPRO phases share a meter+-- with this helper so all three phases consume one budget.+searchJointWith ::+  (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) =>+  BudgetMeter ->+  Miprov2Config ->+  Dataset i o ->+  Metric o ->+  Program i o ->+  -- | per-node instruction candidates+  [[Text]] ->+  -- | per-node demo-set candidates+  [[[Demo]]] ->+  Eff es (Program i o)+searchJointWith meter cfg train metric student instrCands demoCands = do   let nNodes = length (foldParams student)       dsSize = datasetSize train       mbSize = max 1 (min dsSize (minibatchSize cfg))-      maxCalls = maxLmCalls (budget cfg)       apply = applyVec instrCands demoCands student       baseVec = replicate nNodes (0, 0)-  if dsSize > maxCalls-    then pure (apply baseVec)-    else do-      baseFull <- scoreOn train metric (apply baseVec)-      let screen _t calls acc [] = pure (reverse acc, calls)-          screen t calls acc (v : vs)-            | calls + mbSize > maxCalls = pure (reverse acc, calls)-            | otherwise = do-                s <- scoreOn (minibatchAt mbSize (t * 7 + length acc) train) metric (apply v)-                screen t (calls + mbSize) ((v, s) : acc) vs-          go t calls bestVec bestFull+  mBaseFull <- meteredScore meter train metric (apply baseVec)+  case mBaseFull of+    Nothing -> pure (apply baseVec)+    Just baseFull -> do+      let screen _t acc [] = pure (reverse acc)+          screen t acc (v : vs) = do+            ms <- meteredScore meter (minibatchAt mbSize (t * 7 + length acc) train) metric (apply v)+            case ms of+              Nothing -> pure (reverse acc)+              Just s -> screen t ((v, s) : acc) vs+          go t bestVec bestFull             | t >= numTrials cfg = pure bestVec             | otherwise = do-                (scored, calls1) <- screen t calls [] (oneCoordNeighbors instrCands demoCands bestVec)+                scored <- screen t [] (oneCoordNeighbors instrCands demoCands bestVec)                 case bestByScore scored of                   Nothing -> pure bestVec-                  Just propVec-                    | calls1 + dsSize > maxCalls -> pure bestVec-                    | otherwise -> do-                        fs <- scoreOn train metric (apply propVec)+                  Just propVec -> do+                    mFull <- meteredScore meter train metric (apply propVec)+                    case mFull of+                      Nothing -> pure bestVec+                      Just fs ->                         if fs > bestFull-                          then go (t + 1) (calls1 + dsSize) propVec fs-                          else go (t + 1) (calls1 + dsSize) bestVec bestFull-      best <- go 0 dsSize baseVec baseFull+                          then go (t + 1) propVec fs+                          else go (t + 1) bestVec bestFull+      best <- go 0 baseVec baseFull       pure (apply best)  -- | Apply a joint vector to the student: set each node's instruction and demos to the--- selected candidates.+-- selected candidates. Index 0 on either axis is an identity by construction. applyVec :: [[Text]] -> [[[Demo]]] -> Program i o -> JointVec -> Program i o applyVec instrCands demoCands prog vec = foldl step prog (zip [0 ..] vec)   where     step p (k, (i, d)) =-      mapParamsAt-        k-        ( \ps ->-            ps-              { instructionOverride = Just (instrCands !! k !! i),-                demos = demoCands !! k !! d-              }-        )-        p+      setNodeInstrIfNew k (instrCands !! k !! i) $+        mapParamsAt k (\ps -> ps {demos = demoCands !! k !! d}) p++-- | Order-preserving de-duplication.+dedup :: (Eq a) => [a] -> [a]+dedup = go []+  where+    go _ [] = []+    go seen (x : xs)+      | x `elem` seen = go seen xs+      | otherwise = x : go (x : seen) xs  -- | The one-coordinate neighbours of a vector: each move changes exactly one node's -- instruction or demo index to a non-current alternative.

@@ -20,10 +20,10 @@ import Data.Generics.Labels () import Data.List (sortBy) import Data.Ord (comparing)-import Shikumi.Compile.Types (compiledProgram) import Shikumi.Eval (Dataset, dataset, datasetExamples)-import Shikumi.Optimize.Bootstrap (bootstrapFewShotWith, defaultBootstrapConfig)-import Shikumi.Optimize.Search (freezeProgram, scoreOn, selectBest)+import Shikumi.Optimize.Bootstrap (bootstrapKeptDemos, defaultBootstrapConfig)+import Shikumi.Optimize.LabeledFewShot (withDemos)+import Shikumi.Optimize.Search (freezeProgram, meteredScore, newBudgetMeter, selectBestMetered) import Shikumi.Optimize.Types (Budget (..), Optimizer (..), Scored (..)) import Shikumi.Program (Program) @@ -66,7 +66,10 @@         (x : _) -> lo + (x `mod` span')         [] -> lo --- | 'bootstrapRandomSearch' with explicit tunables.+-- | 'bootstrapRandomSearch' with explicit tunables. One 'Budget' covers all seed+-- bootstrap teacher runs and the final candidate scoring pass; when the meter is+-- exhausted, later seeds or scoring candidates are skipped and the best scored+-- candidate so far is returned. bootstrapRandomSearchWith ::   (ToJSON i, ToJSON o) =>   RandomSearchConfig ->@@ -77,14 +80,15 @@   Budget ->   Optimizer i o bootstrapRandomSearchWith cfg teacher numCandidates budget = Optimizer $ \train metric student -> do+  meter <- newBudgetMeter budget   let seeds = [1 .. max 1 numCandidates]       candidateFor seed = do         let cfg' = defaultBootstrapConfig & #maxBootstrappedDemos .~ sizeFor cfg seed-            opt = bootstrapFewShotWith cfg' teacher budget-        compiledProgram <$> runOptimizer opt (shuffle seed train) metric student+        demos <- bootstrapKeptDemos cfg' meter teacher (shuffle seed train) metric+        pure (withDemos demos student)   seeded <- mapM candidateFor seeds   let cands = student : seeded -- zero-shot baseline first-  best <- selectBest budget (scoreOn train metric) cands+  best <- selectBestMetered meter (\p -> meteredScore meter train metric p) cands   pure $ case best of     Nothing -> freezeProgram student     Just sc -> freezeProgram (candidate sc)

@@ -10,26 +10,37 @@ module Shikumi.Optimize.Search   ( selectBest,     scoreOn,+    BudgetMeter (..),+    newBudgetMeter,+    tryCharge,+    scoringCost,+    meteredScore,+    selectBestMetered,+    withLmCallCount,     freezeProgram,     setNodeInstr,+    setNodeInstrIfNew,     instructionAt,+    effectiveInstructionAt,   ) where  import Control.Lens ((&), (?~)) import Data.Generics.Labels ()+import Data.Maybe (fromMaybe) import Data.Text (Text) import Effectful (Eff, (:>)) import Effectful.Concurrent (Concurrent)+import Effectful.Dispatch.Dynamic (interpose) import Effectful.Error.Static (Error)-import Effectful.Prim (Prim)+import Effectful.Prim.IORef (IORef, Prim, atomicModifyIORef', newIORef, readIORef) import Shikumi.Compile.Types (CompiledProgram (..)) import Shikumi.Effect.Time (Time) import Shikumi.Error (ShikumiError)-import Shikumi.Eval (Dataset, Metric, Report (aggregateScore), evaluatePure)-import Shikumi.LLM (LLM)+import Shikumi.Eval (Dataset, Metric, Report (aggregateScore), datasetSize, evaluatePure)+import Shikumi.LLM (LLM (..), complete, stream) import Shikumi.Optimize.Types (Budget (..), Scored (..))-import Shikumi.Program (Params (..), Program, foldParams, mapParamsAt)+import Shikumi.Program (Params (..), Program, foldParams, mapParamsAt, nodeInstructionsIndexed)  -- | Score every candidate (left to right), stopping once the candidate budget is -- hit, and return the best by score (ties: earliest wins). The scorer is the@@ -54,6 +65,112 @@     bestOf [] = Nothing     bestOf (x : xs) = Just (foldl' (\b c -> if score c > score b then c else b) x xs) +-- | Mutable spend-tracking for one optimizer run. The call counter tracks+-- predicted LM completions already reserved; the candidate counter tracks+-- scored candidates.+data BudgetMeter = BudgetMeter+  { meterBudget :: !Budget,+    meterCalls :: !(IORef Int),+    meterCands :: !(IORef Int)+  }++-- | Create a fresh meter for one optimizer invocation.+newBudgetMeter :: (Prim :> es) => Budget -> Eff es BudgetMeter+newBudgetMeter budget = do+  calls <- newIORef 0+  cands <- newIORef 0+  pure BudgetMeter {meterBudget = budget, meterCalls = calls, meterCands = cands}++-- | Charge @n@ predicted LM calls only when doing so keeps the run inside+-- 'maxLmCalls'. Non-positive charges are free and succeed.+tryCharge :: (Prim :> es) => BudgetMeter -> Int -> Eff es Bool+tryCharge _ n | n <= 0 = pure True+tryCharge meter n =+  atomicModifyIORef' (meterCalls meter) $ \calls ->+    if calls + n <= maxLmCalls (meterBudget meter)+      then (calls + n, True)+      else (calls, False)++-- | The predicted LM-call cost of scoring @p@ over @ds@ once: one call per+-- example per 'Predict' node. Wrappers that re-run the LM can spend more.+scoringCost :: Dataset i o -> Program i o -> Int+scoringCost ds p = datasetSize ds * max 1 (length (foldParams p))++-- | Score one candidate program under a meter. Returns 'Nothing' without+-- spending when either the candidate ceiling is reached or the predicted scoring+-- cost no longer fits.+meteredScore ::+  (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) =>+  BudgetMeter ->+  Dataset i o ->+  Metric o ->+  Program i o ->+  Eff es (Maybe Double)+meteredScore meter ds metric prog = do+  fitsCandidate <- tryBumpCandidate meter+  if not fitsCandidate+    then pure Nothing+    else do+      fitsCalls <- tryCharge meter (scoringCost ds prog)+      if fitsCalls+        then Just <$> scoreOn ds metric prog+        else do+          unbumpCandidate meter+          pure Nothing++-- | Fold candidates left-to-right under a scorer that can stop the search by+-- returning 'Nothing'. Ties keep the earliest candidate.+selectBestMetered ::+  BudgetMeter ->+  (cand -> Eff es (Maybe Double)) ->+  [cand] ->+  Eff es (Maybe (Scored cand))+selectBestMetered _ scorer = go Nothing+  where+    go best [] = pure best+    go best (cand : rest) = do+      ms <- scorer cand+      case ms of+        Nothing -> pure best+        Just s ->+          let scored = Scored cand s+              best' = case best of+                Nothing -> Just scored+                Just old -> Just (if score scored > score old then scored else old)+           in go best' rest++-- | Run an action while counting every LLM operation exactly. This is used where+-- the spender is opaque to the budget predictor.+withLmCallCount :: (LLM :> es, Prim :> es) => Eff es a -> Eff es (a, Int)+withLmCallCount act = do+  ref <- newIORef 0+  result <-+    interpose+      ( \_ -> \case+          Complete m c o -> do+            bump ref+            complete m c o+          Stream m c o -> do+            bump ref+            stream m c o+      )+      act+  count <- readIORef ref+  pure (result, count)+  where+    bump ref = atomicModifyIORef' ref (\n -> (n + 1, ()))++tryBumpCandidate :: (Prim :> es) => BudgetMeter -> Eff es Bool+tryBumpCandidate meter =+  atomicModifyIORef' (meterCands meter) $ \cands ->+    if cands < max 0 (maxCandidates (meterBudget meter))+      then (cands + 1, True)+      else (cands, False)++unbumpCandidate :: (Prim :> es) => BudgetMeter -> Eff es ()+unbumpCandidate meter =+  atomicModifyIORef' (meterCands meter) $ \cands -> (max 0 (cands - 1), ())+ -- | Score a candidate program against a dataset and pure metric: run EP-8's -- @evaluate@ and take the aggregate. One call scores one program over the whole -- dataset (one LM call per example).@@ -76,6 +193,24 @@ -- @instructionSearch@ and @copro@. setNodeInstr :: Int -> Text -> Program i o -> Program i o setNodeInstr idx instr = mapParamsAt idx (\ps -> ps & #instructionOverride ?~ instr)++-- | The instruction node @idx@ actually runs with: its 'instructionOverride' when+-- set, otherwise its signature's base instruction. This mirrors the runtime+-- precedence in 'Shikumi.Program.effectiveSignature'. Out-of-range indices yield+-- the empty string.+effectiveInstructionAt :: Int -> Program i o -> Text+effectiveInstructionAt idx prog =+  case drop idx (zip (nodeInstructionsIndexed prog) (foldParams prog)) of+    ((base, ps) : _) -> fromMaybe base (instructionOverride ps)+    [] -> ""++-- | Set node @idx@'s instruction override unless @instr@ is already its effective+-- instruction. Keeping the current signature instruction therefore leaves+-- 'instructionOverride' as 'Nothing' instead of serializing a redundant override.+setNodeInstrIfNew :: Int -> Text -> Program i o -> Program i o+setNodeInstrIfNew idx instr prog+  | instr == effectiveInstructionAt idx prog = prog+  | otherwise = setNodeInstr idx instr prog  -- | The instruction override stored at node @idx@ (in @foldParams@ order), if any. instructionAt :: Int -> Program i o -> Maybe Text

@@ -10,6 +10,15 @@ -- never changes a program's structure or types — only its parameters — so the -- optimized program is the same typed function, merely better-behaved. --+-- Two optimizers are explicit structure-changing exceptions. 'Shikumi.Optimize.KNN.knnFewShot'+-- returns an @Embed@ wrapper that selects demos at run time from an opaque closure;+-- persist the underlying student or use 'Shikumi.Optimize.KNN.knnFewShotCentroid'+-- when a plain parameter artifact is required. 'Shikumi.Optimize.Ensemble.ensembleSearch'+-- returns an @Ensemble@ over member programs; its saved state can only be loaded+-- onto the matching ensemble template, because the reducer closure and member+-- structure are part of the program held in code rather than the saved parameter+-- state.+-- -- __Why a record of functions, not a typeclass.__ The four optimizers differ in -- what extra inputs they carry (a teacher program, a budget, an inner optimizer -- for ensembling) but share one driver signature. A @newtype@ whose field is the@@ -66,11 +75,21 @@       Eff es (CompiledProgram i o)   } --- | A hard, explicit bound on search cost. The optimizers that issue LM calls--- count the calls they make (proposer calls + per-candidate scoring evaluations,--- each scoring evaluation costing one LM call per dataset example) and stop —--- returning the best candidate found /so far/ — before any bound is exceeded, so--- they never silently produce an unscored program and never blow a cost ceiling.+-- | A hard, explicit bound on optimizer search cost. The accounting unit is a+-- predicted LM completion: scoring one candidate costs one completion per dataset+-- example per 'Shikumi.Program.Predict' node; a teacher run costs one completion+-- per 'Shikumi.Program.Predict' node; and a grounded proposer call costs @4 + N@+-- completions for @N@ requested proposals. Optimizers reserve this predicted cost+-- before spending it and return the best candidate found so far when the next+-- spend would exceed a ceiling.+--+-- This prediction is exact for plain predict/compose programs. Wrappers that can+-- re-invoke the LM internally, such as retry, majority-vote, ensemble, or embed+-- bodies, make the prediction a lower bound. 'Shikumi.Optimize.Ensemble.ensembleSearchWith'+-- is the exception that counts member optimizers exactly, but enforces its ceiling+-- only between members, so it may overshoot by one member's cost. If a budget is+-- too small to score even the seed candidate, optimizers return the input program+-- unscored. data Budget = Budget   { -- | ceiling on raw LM calls the optimizer may make (proposals + scoring)     maxLmCalls :: !Int,

@@ -12,6 +12,9 @@ -- reproducible. module AcceptanceSpec (tests) where +import Control.Monad (when)+import Data.IORef (newIORef, readIORef)+import Data.Text qualified as T import Effectful (Eff, IOE, runEff) import Effectful.Concurrent (Concurrent, runConcurrent) import Effectful.Error.Static (Error, runErrorNoCallStack)@@ -20,13 +23,22 @@ import Shikumi.Compile.ZeroShot (zeroShot) import Shikumi.Effect.Time (Time, runTime) import Shikumi.Error (ShikumiError)-import Shikumi.Eval (Dataset, dataset, exactMatch, example)+import Shikumi.Eval (Dataset, Prediction, Score, boolScore, dataset, exactMatch, example, predictionPrimary) import Shikumi.LLM (LLM) import Shikumi.Optimize import Shikumi.Program (Program)-import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, sentimentProg)+import StubLM+  ( Label (..),+    Sentence (..),+    ruleInstruction,+    ruled,+    runStubLM,+    runStubLMCapturing,+    sentimentPipeline,+    sentimentProg,+  ) import Test.Tasty (TestTree, testGroup)-import Test.Tasty.HUnit (assertBool, assertFailure, testCase)+import Test.Tasty.HUnit (Assertion, assertBool, assertFailure, testCase, (@?=))  runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act@@ -86,7 +98,30 @@             assertBool ("before should be low, got " <> show before) (before < 0.5)             improved "bootstrapFewShot" before afterBoot             improved "labeledFewShot" before afterLabeled-            improved "instructionSearch" before afterInstr+            improved "instructionSearch" before afterInstr,+      testGroup+        "seeding over multi-node programs"+        [ testCase "instructionSearch optimizes node 0 without changing node 1" $+            checkPipeline "instructionSearch" (instructionSearch 3 defaultBudget) True,+          testCase "copro optimizes node 0 without changing node 1" $+            checkPipeline "copro" (copro defaultCoproConfig) True,+          testCase "miprov2 keeps the echo node's effective instruction" $+            checkPipeline "miprov2" (miprov2 Miprov2Light) False,+          testCase "gepa keeps the echo node's effective instruction" $+            checkPipeline "gepa" (gepa reflectiveProposer fbMetric defaultBudget) False,+          testCase "instructionSearch scores the true student as a candidate" $ do+            ref <- newIORef []+            res <-+              runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $+                runStubLMCapturing ref (optimize (instructionSearch 1 defaultBudget) trainset exactMatch ruled)+            case res of+              Left e -> assertFailure ("unexpected error: " <> show e)+              Right _ -> pure ()+            captured <- readIORef ref+            assertBool+              "expected a scored request rendered with the RULE signature instruction"+              (any (\txt -> ruleInstruction `T.isInfixOf` txt && "text: good" `T.isInfixOf` txt) captured)+        ]     ]   where     -- Strictly higher than before, and at least a floor of 0.75, so the test fails@@ -98,3 +133,24 @@       assertBool         (name <> ": expected " <> show after <> " >= 0.75")         (after >= 0.75)++fbMetric :: Label -> Prediction Label -> (Score, T.Text)+fbMetric expd p =+  let correct = expd == predictionPrimary p+   in (boolScore correct, if correct then "" else "be more specific")++checkPipeline :: String -> Optimizer Sentence Label -> Bool -> Assertion+checkPipeline name opt requireFloor = do+  res <- runStub $ do+    before <- scoreOn heldout exactMatch sentimentPipeline+    cp <- optimize opt trainset exactMatch sentimentPipeline+    let out = compiledProgram cp+    after <- scoreOn heldout exactMatch out+    pure (before, after, effectiveInstructionAt 1 out)+  case res of+    Left e -> assertFailure ("unexpected error: " <> show e)+    Right (before, after, node1Instruction) -> do+      assertBool (name <> ": expected " <> show after <> " >= " <> show before) (after >= before)+      when requireFloor $+        assertBool (name <> ": expected " <> show after <> " >= 0.75") (after >= 0.75)+      node1Instruction @?= "Echo the sentiment label unchanged."

@@ -9,6 +9,7 @@ -- teacher correct on, and the mislabelled one is excluded. module BootstrapSpec (tests) where +import Data.IORef (newIORef, readIORef) import Effectful (Eff, IOE, runEff) import Effectful.Concurrent (Concurrent, runConcurrent) import Effectful.Error.Static (Error, runErrorNoCallStack)@@ -22,9 +23,9 @@ import Shikumi.Optimize import Shikumi.Program (Demo (..), Params (..), Program, programParams) import Shikumi.Schema (fromModel)-import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, sentimentProg)+import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, runStubLMCounting, sentimentPipeline, sentimentProg) import Test.Tasty (TestTree, testGroup)-import Test.Tasty.HUnit (assertFailure, testCase, (@?=))+import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=))  runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act@@ -44,6 +45,15 @@       example (Sentence "good soup") (Label "negative")     ] +budgetTrainset :: Dataset Sentence Label+budgetTrainset =+  dataset+    [ example (Sentence "good film") (Label "positive"),+      example (Sentence "bad film") (Label "negative"),+      example (Sentence "good show") (Label "positive"),+      example (Sentence "bad play") (Label "negative")+    ]+ -- | The demos baked into a compiled single-node program. compiledDemos :: CompiledProgram Sentence Label -> [Demo] compiledDemos cp = case programParams (compiledProgram cp) of@@ -68,5 +78,16 @@                 outs = traverse (\d -> fromModel (output d)) ds :: Either ShikumiError [Label]             length ds @?= 2             ins @?= Right [Sentence "good film", Sentence "bad film"]-            outs @?= Right [Label "positive", Label "negative"]+            outs @?= Right [Label "positive", Label "negative"],+      testCase "bootstrapFewShot charges a teacher run by predict-node count" $ do+        ref <- newIORef (0 :: Int)+        let budget = Budget {maxLmCalls = 6, maxCandidates = 32}+        res <-+          runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $+            runStubLMCounting ref (optimize (bootstrapFewShot sentimentPipeline budget) budgetTrainset exactMatch sentimentProg)+        case res of+          Left e -> assertFailure ("unexpected error: " <> show e)+          Right _ -> pure ()+        count <- readIORef ref+        assertBool ("expected 0 < calls <= 6, got " <> show count) (count > 0 && count <= 6)     ]

@@ -10,6 +10,7 @@ module EnsembleSpec (tests) where  import Data.Aeson (toJSON)+import Data.IORef (newIORef, readIORef) import Data.Text (Text) import Effectful (Eff, IOE, runEff) import Effectful.Concurrent (Concurrent, runConcurrent)@@ -23,7 +24,7 @@ import Shikumi.LLM (LLM) import Shikumi.Optimize import Shikumi.Program (Demo (..), Program, ProgramShape (..), programShape)-import StubLM (Label (..), Sentence (..), runStubLM, sentimentProg)+import StubLM (Label (..), Sentence (..), runStubLM, runStubLMCounting, sentimentProg) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) @@ -88,5 +89,18 @@           Left e -> assertFailure ("unexpected error: " <> show e)           Right cp -> case programShape (compiledProgram cp) of             ShapeEnsemble subs -> length subs @?= 3-            other -> assertFailure ("expected a 3-member ShapeEnsemble, got " <> show other)+            other -> assertFailure ("expected a 3-member ShapeEnsemble, got " <> show other),+      testCase "ensembleSearchWith enforces the budget between members" $ do+        ref <- newIORef (0 :: Int)+        let budget = Budget {maxLmCalls = 1, maxCandidates = 32}+        res <-+          runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $+            runStubLMCounting ref (optimize (ensembleSearchWith budget 5 (labeledFewShot 1)) trainset exactMatch sentimentProg)+        case res of+          Left e -> assertFailure ("unexpected error: " <> show e)+          Right cp -> case programShape (compiledProgram cp) of+            ShapeEnsemble subs -> length subs @?= 1+            other -> assertFailure ("expected a 1-member ShapeEnsemble, got " <> show other)+        count <- readIORef ref+        count @?= 4     ]

@@ -16,7 +16,8 @@ import Shikumi.Eval (Dataset, boolScore, dataset, exactMatch, example, predictionPrimary) import Shikumi.LLM (LLM) import Shikumi.Optimize-  ( Candidate (..),+  ( Budget (..),+    Candidate (..),     FeedbackMetric,     captureFeedback,     defaultBudget,@@ -28,6 +29,7 @@     reflectiveProposer,     sampleParent,     scoreOn,+    withLmCallCount,   ) import Shikumi.Program (Params (..), foldParams, nodeFieldsIndexed, programParams) import Shikumi.Trace.Feedback (feedbackFor)@@ -65,7 +67,7 @@ runGepa act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runGepaStubLM act  tests :: TestTree-tests = testGroup "Gepa" [paretoPure, feedbackCapture, reflectiveMutation, heldoutLift, roundTrips]+tests = testGroup "Gepa" [paretoPure, feedbackCapture, reflectiveMutation, heldoutLift, budgetGate, roundTrips]  -- --------------------------------------------------------------------------- -- Pure Pareto-frontier tests@@ -153,6 +155,17 @@         assertBool ("before should be low, got " <> show before) (before < 0.5)         assertBool ("after " <> show after <> " should exceed before " <> show before) (after > before)         assertBool ("after should reach the floor, got " <> show after) (after >= 0.75)++budgetGate :: TestTree+budgetGate =+  testCase "gepa returns the student without seed evaluation when the budget is too small" $ do+    let tiny = Budget {maxLmCalls = 1, maxCandidates = 4}+    res <- runGepa (withLmCallCount (optimize (gepa reflectiveProposer fbMetric tiny) trainset exactMatch sentimentProg))+    case res of+      Left e -> assertFailure ("unexpected error: " <> show e)+      Right (cp, calls) -> do+        calls @?= 0+        programParams (compiledProgram cp) @?= programParams sentimentProg  roundTrips :: TestTree roundTrips =

@@ -20,7 +20,7 @@ import Shikumi.LLM (LLM) import Shikumi.Optimize import Shikumi.Program (Params (..), programParams)-import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, runStubLMCounting, sentimentProg)+import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, runStubLMCounting, sentimentPipeline, sentimentProg) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) @@ -57,5 +57,18 @@         count <- readIORef ref         assertBool           ("expected 0 < calls <= 6, got " <> show count)-          (count > 0 && count <= 6)+          (count > 0 && count <= 6),+      testCase "charges scoring by dataset size times predict nodes" $ do+        ref <- newIORef (0 :: Int)+        let budget = Budget {maxLmCalls = 11, maxCandidates = 32}+        res <-+          runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $+            runStubLMCounting ref (optimize (instructionSearch 1 budget) trainset exactMatch sentimentPipeline)+        case res of+          Left e -> assertFailure ("unexpected error: " <> show e)+          Right _ -> pure ()+        count <- readIORef ref+        assertBool+          ("expected 0 < calls <= 11 for two-node scoring, got " <> show count)+          (count > 0 && count <= 11)     ]

@@ -23,7 +23,7 @@ import Shikumi.Error (ShikumiError) import Shikumi.Eval (Dataset, dataset, datasetExamples, exactMatch, example) import Shikumi.LLM (LLM)-import Shikumi.Optimize (knnDemos, knnFewShotCentroid, nearestDemos, optimize)+import Shikumi.Optimize (knnDemos, knnFewShot, knnFewShotCentroid, nearestDemos, optimize, withLmCallCount) import Shikumi.Optimize.Search (freezeProgram) import Shikumi.Program (Demo (..), Params (..), programParams, runProgram) import StubLM (Label (..), Sentence (..), runStubLM, sentimentProg)@@ -74,7 +74,7 @@ runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act  tests :: TestTree-tests = testGroup "KNN" [inputNearest, centroidNearest, runtimeRuns, roundTrips]+tests = testGroup "KNN" [inputNearest, centroidNearest, runtimeRuns, optimizeBudget, roundTrips]  -- --------------------------------------------------------------------------- -- M1 — selection@@ -105,6 +105,18 @@     res <- runStub (runProgram prog (Sentence "france capital quiz"))     assertBool "the Embed-wrapped KNN program runs" (either (const False) (const True) res) +optimizeBudget :: TestTree+optimizeBudget =+  testCase "KNN optimizers make zero LM calls at optimize time" $ do+    res <-+      runStub $ do+        (_, runtimeCalls) <- withLmCallCount (optimize (knnFewShot stubEmbed 2) balanced exactMatch sentimentProg)+        (_, centroidCalls) <- withLmCallCount (optimize (knnFewShotCentroid stubEmbed 2) skewed exactMatch sentimentProg)+        pure (runtimeCalls, centroidCalls)+    case res of+      Left e -> assertFailure ("unexpected error: " <> show e)+      Right counts -> counts @?= (0, 0)+ -- --------------------------------------------------------------------------- -- M3 — serialization round-trip -- ---------------------------------------------------------------------------@@ -119,6 +131,11 @@         case decodeCompiledOnto (knnDemos stubEmbed 2 balanced sentimentProg) (encodeCompiled compiled) of           Left err -> assertFailure ("decode failed: " <> err)           Right cp' -> programParams (compiledProgram cp') @?= [],+      testCase "run-time KNN decode onto the plain student fails with a shape message" $ do+        let compiled = freezeProgram (knnDemos stubEmbed 2 balanced sentimentProg)+        case decodeCompiledOnto sentimentProg (encodeCompiled compiled) of+          Left err -> assertBool ("shape mismatch mentioned: " <> err) ("shape mismatch" `T.isInfixOf` T.pack err)+          Right _ -> assertFailure "expected decode onto the plain student to fail",       testCase "centroid KNN baked demos round-trip onto the bare student" $ do         res <- runStub (optimize (knnFewShotCentroid stubEmbed 2) skewed exactMatch sentimentProg)         case res of

@@ -9,6 +9,7 @@ -- is a perfect 1.0. module LabeledFewShotSpec (tests) where +import Data.IORef (newIORef, readIORef) import Effectful (Eff, IOE, runEff) import Effectful.Concurrent (Concurrent, runConcurrent) import Effectful.Error.Static (Error, runErrorNoCallStack)@@ -20,7 +21,7 @@ import Shikumi.LLM (LLM) import Shikumi.Optimize import Shikumi.Program (Demo, Params (..), programParams)-import StubLM (Label (..), Sentence (..), runStubLM, sentimentProg)+import StubLM (Label (..), Sentence (..), runStubLM, runStubLMCounting, sentimentProg) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) @@ -70,5 +71,16 @@               (maximum candScores > minimum candScores)             assertBool               ("best candidate should be perfect, got " <> show chosenScore)-              (chosenScore == 1.0)+              (chosenScore == 1.0),+      testCase "labeledFewShotWith respects maxLmCalls while scoring candidates" $ do+        ref <- newIORef (0 :: Int)+        let budget = Budget {maxLmCalls = 6, maxCandidates = 100}+        res <-+          runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $+            runStubLMCounting ref (optimize (labeledFewShotWith budget 2) trainset exactMatch sentimentProg)+        case res of+          Left e -> assertFailure ("unexpected error: " <> show e)+          Right _ -> pure ()+        count <- readIORef ref+        assertBool ("expected 0 < calls <= 6, got " <> show count) (count > 0 && count <= 6)     ]

@@ -17,6 +17,8 @@ import OptimizeSpec qualified import ProposeSpec qualified import RandomSearchSpec qualified+import SearchSpec qualified+import SeedingSpec qualified import Test.Tasty (defaultMain, testGroup)  main :: IO ()@@ -35,5 +37,7 @@         GepaSpec.tests,         KNNSpec.tests,         RandomSearchSpec.tests,+        SearchSpec.tests,+        SeedingSpec.tests,         AcceptanceSpec.tests       ]

@@ -28,6 +28,7 @@     bootstrapDemoCandidates,     defaultBudget,     instructionSearch,+    miprov2,     miprov2Auto,     miprov2With,     optimize,@@ -179,7 +180,28 @@             assertBool               ("miprov2 " <> show afterMipro <> " should beat instructionSearch " <> show afterInstr)               (afterMipro > afterInstr)-            assertBool ("miprov2 should reach the floor, got " <> show afterMipro) (afterMipro >= 0.75)+            assertBool ("miprov2 should reach the floor, got " <> show afterMipro) (afterMipro >= 0.75),+      testCase "miprov2With shares one budget across bootstrap, proposal, and search" $ do+        ref <- newIORef (0 :: Int)+        let cfg = cfgLight & #budget .~ Budget {maxLmCalls = 4, maxCandidates = 32}+        res <-+          runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $+            runJointStubLMCounting ref (optimize (miprov2With cfg sentimentProg) jointTrain exactMatch sentimentProg)+        case res of+          Left e -> assertFailure ("unexpected error: " <> show e)+          Right _ -> pure ()+        count <- readIORef ref+        assertBool ("expected 0 < calls <= 4, got " <> show count) (count > 0 && count <= 4),+      testCase "miprov2 preset stays within the default budget" $ do+        ref <- newIORef (0 :: Int)+        res <-+          runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $+            runJointStubLMCounting ref (optimize (miprov2 Miprov2Light) jointTrain exactMatch sentimentProg)+        case res of+          Left e -> assertFailure ("unexpected error: " <> show e)+          Right _ -> pure ()+        count <- readIORef ref+        assertBool ("expected 0 < calls <= 200, got " <> show count) (count > 0 && count <= maxLmCalls defaultBudget)     ]  serialize :: TestTree

@@ -5,6 +5,7 @@ -- the best compiled output round-trips through serialization. module RandomSearchSpec (tests) where +import Data.IORef (newIORef, readIORef) import Effectful (Eff, IOE, runEff) import Effectful.Concurrent (Concurrent, runConcurrent) import Effectful.Error.Static (Error, runErrorNoCallStack)@@ -16,7 +17,8 @@ import Shikumi.Eval (Dataset, dataset, exactMatch, example) import Shikumi.LLM (LLM) import Shikumi.Optimize-  ( bootstrapFewShot,+  ( Budget (..),+    bootstrapFewShot,     bootstrapRandomSearch,     defaultBudget,     optimize,@@ -24,7 +26,7 @@     setNodeInstr,   ) import Shikumi.Program (Program)-import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, sentimentProg)+import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, runStubLMCounting, sentimentProg) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) @@ -57,7 +59,7 @@ runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act  tests :: TestTree-tests = testGroup "RandomSearch" [bestOfNAtLeastSingle, reproducible, roundTrips]+tests = testGroup "RandomSearch" [bestOfNAtLeastSingle, reproducible, budgetRespected, roundTrips]  bestOfNAtLeastSingle :: TestTree bestOfNAtLeastSingle =@@ -89,6 +91,20 @@     case res of       Left e -> assertFailure ("unexpected error: " <> show e)       Right (sa, sb) -> sa @?= sb++budgetRespected :: TestTree+budgetRespected =+  testCase "bootstrapRandomSearch shares one budget across seeds and scoring" $ do+    ref <- newIORef (0 :: Int)+    let budget = Budget {maxLmCalls = 20, maxCandidates = 100}+    res <-+      runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $+        runStubLMCounting ref (optimize (bootstrapRandomSearch teacher 5 budget) trainset exactMatch sentimentProg)+    case res of+      Left e -> assertFailure ("unexpected error: " <> show e)+      Right _ -> pure ()+    count <- readIORef ref+    assertBool ("expected 0 < calls <= 20, got " <> show count) (count > 0 && count <= 20)  roundTrips :: TestTree roundTrips =

@@ -0,0 +1,114 @@+{-# LANGUAGE TypeApplications #-}++-- | Unit tests for the shared optimizer search helpers.+module SearchSpec (tests) where++import Data.IORef (newIORef, readIORef)+import Effectful (Eff, IOE, runEff)+import Effectful.Concurrent (Concurrent, runConcurrent)+import Effectful.Error.Static (Error, runErrorNoCallStack)+import Effectful.Prim (Prim, runPrim)+import Effectful.Prim.IORef qualified as EIORef+import Shikumi.Combinator ((>>>))+import Shikumi.Effect.Time (Time, runTime)+import Shikumi.Error (ShikumiError)+import Shikumi.Eval (Dataset, dataset, exactMatch, example)+import Shikumi.LLM (LLM)+import Shikumi.Module (predict)+import Shikumi.Optimize.Search+  ( BudgetMeter (..),+    effectiveInstructionAt,+    instructionAt,+    meteredScore,+    newBudgetMeter,+    scoreOn,+    selectBestMetered,+    setNodeInstr,+    setNodeInstrIfNew,+    tryCharge,+    withLmCallCount,+  )+import Shikumi.Optimize.Types (Budget (..), Scored (..))+import Shikumi.Program (Program)+import Shikumi.Signature (Signature, mkSignature)+import StubLM (Label (..), Sentence (..), ruleInstruction, ruled, runStubLM, runStubLMCounting, sentimentProg)+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=))++runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a)+runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act++trainset :: Dataset Sentence Label+trainset =+  dataset+    [ example (Sentence "good film") (Label "positive"),+      example (Sentence "bad film") (Label "negative")+    ]++echoSig :: Signature Label Label+echoSig = mkSignature "Echo the sentiment label unchanged."++twoNode :: Program Sentence Label+twoNode = sentimentProg >>> predict echoSig++tests :: TestTree+tests =+  testGroup+    "Search"+    [ testCase "effectiveInstructionAt reads signature base, override, and out-of-range" $ do+        effectiveInstructionAt 0 ruled @?= ruleInstruction+        effectiveInstructionAt 0 (setNodeInstr 0 "override" ruled) @?= "override"+        effectiveInstructionAt 99 ruled @?= "",+      testCase "setNodeInstrIfNew does not write redundant override" $ do+        let kept = setNodeInstrIfNew 0 ruleInstruction ruled+            changed = setNodeInstrIfNew 0 "different" ruled+        instructionAt 0 kept @?= Nothing+        instructionAt 0 changed @?= Just "different",+      testCase "effectiveInstructionAt is index-aligned on a two-node program" $ do+        effectiveInstructionAt 0 twoNode @?= ""+        effectiveInstructionAt 1 twoNode @?= "Echo the sentiment label unchanged."+        effectiveInstructionAt 1 (setNodeInstr 1 "changed" twoNode) @?= "changed",+      testCase "tryCharge refuses an over-budget charge without changing the counter" $ do+        (ok1, ok2, ok3, calls) <-+          runEff . runPrim $ do+            meter <- newBudgetMeter Budget {maxLmCalls = 3, maxCandidates = 32}+            ok1 <- tryCharge meter 2+            ok2 <- tryCharge meter 2+            ok3 <- tryCharge meter 1+            calls <- EIORef.readIORef (meterCalls meter)+            pure (ok1, ok2, ok3, calls)+        (ok1, ok2, ok3, calls) @?= (True, False, True, 3),+      testCase "meteredScore returns Nothing without invoking the LM when cost does not fit" $ do+        ref <- newIORef (0 :: Int)+        res <-+          runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $+            runStubLMCounting ref $ do+              meter <- newBudgetMeter Budget {maxLmCalls = 1, maxCandidates = 32}+              meteredScore meter trainset exactMatch sentimentProg+        case res of+          Left e -> assertFailure ("unexpected error: " <> show e)+          Right Nothing -> pure ()+          Right (Just s) -> assertFailure ("expected no score, got " <> show s)+        count <- readIORef ref+        count @?= 0,+      testCase "selectBestMetered stops at first Nothing and returns best so far" $ do+        best <-+          runEff . runPrim $ do+            meter <- newBudgetMeter Budget {maxLmCalls = 100, maxCandidates = 32}+            selectBestMetered+              meter+              ( \n ->+                  pure $+                    if n >= (4 :: Int)+                      then Nothing+                      else Just (fromIntegral n)+              )+              [1 :: Int, 3, 2, 4, 99]+        best @?= Just (Scored 3 3.0),+      testCase "withLmCallCount counts one completion per scored example" $ do+        res <- runStub $ withLmCallCount (scoreOn trainset exactMatch sentimentProg)+        case res of+          Left e -> assertFailure ("unexpected error: " <> show e)+          Right (_score, count) ->+            assertBool ("expected two completions, got " <> show count) (count == 2)+    ]

@@ -0,0 +1,164 @@+{-# LANGUAGE TypeApplications #-}++-- | Regression tests for optimizer instruction seeding.+module SeedingSpec (tests) where++import Control.Lens ((&), (.~))+import Data.Generics.Labels ()+import Data.IORef (newIORef, readIORef)+import Data.Text qualified as T+import Effectful (Eff, IOE, runEff)+import Effectful.Concurrent (Concurrent, runConcurrent)+import Effectful.Error.Static (Error, runErrorNoCallStack)+import Effectful.Prim (Prim, runPrim)+import Shikumi.Compile.Types (compiledProgram)+import Shikumi.Effect.Time (Time, runTime)+import Shikumi.Error (ShikumiError)+import Shikumi.Eval (Dataset, dataset, exactMatch, example)+import Shikumi.Eval qualified as Eval+import Shikumi.LLM (LLM)+import Shikumi.Optimize+  ( Budget (..),+    CoproConfig (..),+    Miprov2Auto (..),+    Miprov2Config (..),+    bootstrapDemoCandidates,+    copro,+    defaultBudget,+    gepa,+    instructionAt,+    instructionSearch,+    miprov2Auto,+    miprov2With,+    optimize,+    recoverDemo,+    reflectiveProposer,+    scoreOn,+    withDemos,+  )+import Shikumi.Program (Demo)+import StubLM (Label (..), Sentence (..), ruleInstruction, ruled, runGepaStubLMCapturing, runJointStubLM, runStubLM)+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=))++runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a)+runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act++runJoint :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a)+runJoint act =+  runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runJointStubLM act++trainset :: Dataset Sentence Label+trainset =+  dataset+    [ example (Sentence "good film") (Label "positive"),+      example (Sentence "good book") (Label "positive"),+      example (Sentence "bad film") (Label "negative"),+      example (Sentence "bad book") (Label "negative")+    ]++heldout :: Dataset Sentence Label+heldout =+  dataset+    [ example (Sentence "good movie") (Label "positive"),+      example (Sentence "bad movie") (Label "negative")+    ]++jointTrain :: Dataset Sentence Label+jointTrain =+  dataset+    [ example (Sentence "good movie") (Label "positive"),+      example (Sentence "bad movie") (Label "negative"),+      example (Sentence "great film") (Label "positive"),+      example (Sentence "terrible film") (Label "negative")+    ]++jointHeldout :: Dataset Sentence Label+jointHeldout =+  dataset+    [ example (Sentence "good show") (Label "positive"),+      example (Sentence "bad show") (Label "negative"),+      example (Sentence "great show") (Label "positive"),+      example (Sentence "terrible show") (Label "negative")+    ]++coveringDemos :: [Demo]+coveringDemos =+  [ recoverDemo (Sentence "great film") (Label "positive"),+    recoverDemo (Sentence "terrible film") (Label "negative")+  ]++tests :: TestTree+tests =+  testGroup+    "instruction seeding"+    [ testCase "instructionSearch never degrades a solved signature instruction" $ do+        res <-+          runStub $ do+            before <- scoreOn heldout exactMatch ruled+            cp <- optimize (instructionSearch 1 defaultBudget) trainset exactMatch ruled+            after <- scoreOn heldout exactMatch (compiledProgram cp)+            pure (before, after, instructionAt 0 (compiledProgram cp))+        case res of+          Left e -> assertFailure ("unexpected error: " <> show e)+          Right (before, after, override) -> do+            before @?= 1.0+            assertBool ("instructionSearch: expected " <> show after <> " >= " <> show before) (after >= before)+            override @?= Nothing,+      testCase "copro never degrades a solved student when proposer calls are unaffordable" $ do+        let cfg = CoproConfig {breadth = 2, depth = 1, budget = Budget {maxLmCalls = 4, maxCandidates = 32}}+        res <-+          runStub $ do+            before <- scoreOn heldout exactMatch ruled+            cp <- optimize (copro cfg) trainset exactMatch ruled+            after <- scoreOn heldout exactMatch (compiledProgram cp)+            pure (before, after, instructionAt 0 (compiledProgram cp))+        case res of+          Left e -> assertFailure ("unexpected error: " <> show e)+          Right (before, after, override) -> do+            before @?= 1.0+            assertBool ("copro: expected " <> show after <> " >= " <> show before) (after >= before)+            override @?= Nothing,+      testCase "miprov2 never returns worse than its solved input baseline" $ do+        let student = withDemos coveringDemos ruled+            cfg :: Miprov2Config+            cfg = miprov2Auto Miprov2Light & #budget .~ Budget {maxLmCalls = 4, maxCandidates = 32}+        res <-+          runJoint $ do+            before <- scoreOn jointHeldout exactMatch student+            cp <- optimize (miprov2With cfg student) jointTrain exactMatch student+            after <- scoreOn jointHeldout exactMatch (compiledProgram cp)+            pure (before, after)+        case res of+          Left e -> assertFailure ("unexpected error: " <> show e)+          Right (before, after) -> do+            before @?= 1.0+            assertBool ("miprov2: expected " <> show after <> " >= " <> show before) (after >= before),+      testCase "miprov2 demo candidate 0 is the node's current demos" $ do+        let student = withDemos coveringDemos ruled+            cfg :: Miprov2Config+            cfg = miprov2Auto Miprov2Light & #numDemoCandidates .~ 3+        res <- runJoint (bootstrapDemoCandidates cfg student jointTrain exactMatch student)+        case res of+          Left e -> assertFailure ("unexpected error: " <> show e)+          Right [sets] -> do+            take 1 sets @?= [coveringDemos]+            assertBool "empty demo set remains available" ([] `elem` sets)+          Right other -> assertFailure ("expected one node, got " <> show (length other)),+      testCase "gepa reflection prompt sees the effective signature instruction" $ do+        ref <- newIORef []+        let fbMetric _ _ = (Eval.boolScore False, "be more specific")+        res <-+          runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $+            runGepaStubLMCapturing ref (optimize (gepa reflectiveProposer fbMetric defaultBudget) trainset exactMatch ruled)+        case res of+          Left e -> assertFailure ("unexpected error: " <> show e)+          Right _ -> pure ()+        captured <- readIORef ref+        case filter ("proposedInstruction" `T.isInfixOf`) captured of+          firstReflection : _ ->+            assertBool+              "expected first GEPA reflection request to include the effective RULE instruction"+              (("currentInstruction: " <> ruleInstruction) `T.isInfixOf` firstReflection)+          [] -> assertFailure "expected at least one GEPA reflection request"+    ]

@@ -39,6 +39,9 @@     -- * Signature and program     sentimentSig,     sentimentProg,+    ruled,+    echoSig,+    sentimentPipeline,      -- * Ground truth and helpers     goldLabel,@@ -56,6 +59,7 @@      -- * The reflective task (EP-22, GEPA)     runGepaStubLM,+    runGepaStubLMCapturing,   ) where @@ -82,6 +86,7 @@ import Effectful.Dispatch.Dynamic (interpret) import GHC.Generics (Generic) import Shikumi.Adapter (ToPrompt)+import Shikumi.Combinator ((>>>)) import Shikumi.LLM (LLM (..)) import Shikumi.Module (predict) import Shikumi.Program (Program)@@ -132,6 +137,19 @@ sentimentProg :: Program Sentence Label sentimentProg = predict sentimentSig +-- | A single-node sentiment program whose signature instruction already solves the+-- task under the stub. Optimizers must preserve it when "keep current" wins.+ruled :: Program Sentence Label+ruled = predict (mkSignature ruleInstruction)++-- | A second-stage identity node for multi-node optimizer tests.+echoSig :: Signature Label Label+echoSig = mkSignature "Echo the sentiment label unchanged."++-- | A two-node sentiment pipeline: classify a sentence, then echo the label.+sentimentPipeline :: Program Sentence Label+sentimentPipeline = sentimentProg >>> predict echoSig+ -- --------------------------------------------------------------------------- -- Ground truth and helpers -- ---------------------------------------------------------------------------@@ -206,6 +224,7 @@ -- | The joint task's classification rule (see 'runJointStubLM'). answerJoint :: Context -> Text answerJoint ctx+  | Just lbl <- parseEcho (lastUserText ctx) = lbl   | "good" `elem` ws || "bad" `elem` ws =       -- region A: correct only with a RULE instruction       if instructionHasRule ctx then goldLabel s else "neutral"@@ -262,6 +281,15 @@   Complete _ ctx _ -> pure (mkResponse (respondGepa ctx))   Stream {} -> pure [] +-- | Like 'runGepaStubLM' but records each request's rendered text for prompt-signal+-- assertions.+runGepaStubLMCapturing :: (IOE :> es) => IORef [Text] -> Eff (LLM : es) a -> Eff es a+runGepaStubLMCapturing ref = interpret $ \_ -> \case+  Complete _ ctx _ -> do+    liftIO (modifyIORef' ref (++ [fullRequestText ctx]))+    pure (mkResponse (respondGepa ctx))+  Stream {} -> pure []+ respondGepa :: Context -> Text respondGepa ctx   | sysHasMarker "proposedInstruction" ctx =@@ -283,9 +311,11 @@ -- | Classify the actual input given the demos and instruction in the context. answerSentiment :: Context -> Text answerSentiment ctx =-  case demos of-    [] -> if instructionHasRule ctx then goldLabel s else "neutral"-    _ -> nnLabel s demos+  case parseEcho (lastUserText ctx) of+    Just lbl -> lbl+    Nothing -> case demos of+      [] -> if instructionHasRule ctx then goldLabel s else "neutral"+      _ -> nnLabel s demos   where     msgs = V.toList (ctx ^. #messages)     demos = demoPairs msgs@@ -364,6 +394,10 @@ parseSentence t = T.strip (fromMaybe stripped (T.stripPrefix "text:" stripped))   where     stripped = T.strip t++-- | Parse the rendered input to the echo node, whose input type is 'Label'.+parseEcho :: Text -> Maybe Text+parseEcho t = T.strip <$> T.stripPrefix "sentiment:" (T.strip t)  -- | Read the @sentiment@ marker section out of a rendered demo output. parseLabel :: Text -> Text