diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,20 @@ ## Unreleased +## 0.3.0.0 — 2026-09-08 + +- Raise the internal bounds to `shikumi ^>=0.4.0.0`, `shikumi-compile ^>=0.2.1.0`, `shikumi-eval ^>=0.3.0.0`, and `shikumi-trace ^>=0.3.0.0`. + +- Upgrade the dependency on `mori://shinzui/baikai/packages/baikai` to `>=0.7.0.0 && <0.8`. + +- Add `structureSearchWith` for validation-only finite recipe selection through shared admission/objectives, plus optional candidate identity metadata in reports and lifecycle events. + +- Add validated GEPA execution, named objective selection, hard LLM-operation admission, bounded candidate execution, and diagnostic lifecycle reports. + +- Add indexed failure-aware feedback, validated node/invocation attribution, effectful GEPA callbacks, bounded redacted local reflection, and typed error classification. Legacy critiques are now labeled program-scoped rather than broadcast to nodes; existing entry-point signatures remain supported. + +- Recover validated node-local bootstrap pools with explicit mapping, per-node caps, seeded independent selection and diagnostics. RandomSearch and MIPRO use node pools. Composite bootstrap requires captured leaves; `bootstrapKeptDemos` is now single-node only. Existing BootstrapConfig construction and legacy single-predictor bootstrap remain supported. + ## 0.2.1.3 — 2026-08-29 ### Changed diff --git a/shikumi-optimize.cabal b/shikumi-optimize.cabal --- a/shikumi-optimize.cabal +++ b/shikumi-optimize.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: shikumi-optimize -version: 0.2.1.3 +version: 0.3.0.0 synopsis: The optimizer framework for shikumi LM programs (EP-10) category: AI description: @@ -44,6 +44,8 @@ Shikumi.Optimize.Bootstrap Shikumi.Optimize.COPRO Shikumi.Optimize.Ensemble + Shikumi.Optimize.Execution + Shikumi.Optimize.Feedback Shikumi.Optimize.GEPA Shikumi.Optimize.Instruction Shikumi.Optimize.KNN @@ -56,7 +58,9 @@ Shikumi.Optimize.Propose.Tips Shikumi.Optimize.Propose.Types Shikumi.Optimize.RandomSearch + Shikumi.Optimize.Report Shikumi.Optimize.Search + Shikumi.Optimize.Structure Shikumi.Optimize.Types build-depends: @@ -66,10 +70,10 @@ , effectful >=2.5 && <2.7 , generic-lens >=2.2 && <2.4 , lens ^>=5.3 - , shikumi ^>=0.3.0.0 - , shikumi-compile ^>=0.2.0.0 - , shikumi-eval ^>=0.2.0.0 - , shikumi-trace ^>=0.2.0.0 + , shikumi ^>=0.4.0.0 + , shikumi-compile ^>=0.2.1.0 + , shikumi-eval ^>=0.3.0.0 + , shikumi-trace ^>=0.3.0.0 , text ^>=2.1 , vector >=0.13 && <0.14 @@ -84,31 +88,36 @@ BootstrapSpec CoproSpec EnsembleSpec + ExecutionSpec + FeedbackSpec GepaSpec InstructionSpec KNNSpec LabeledFewShotSpec Miprov2Spec + NodeBootstrapSpec + ObjectiveSpec OptimizeSpec ProposeSpec RandomSearchSpec SearchSpec SeedingSpec + StructureSpec StubLM build-depends: , aeson - , baikai >=0.6 && <0.7 + , baikai >=0.7.0.0 && <0.8 , base , containers , effectful , generic-lens , lens - , 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 + , shikumi ^>=0.4.0.0 + , shikumi-compile ^>=0.2.1.0 + , shikumi-eval ^>=0.3.0.0 + , shikumi-optimize ^>=0.3.0.0 + , shikumi-trace ^>=0.3.0.0 , tasty , tasty-hunit , text diff --git a/src/Shikumi/Optimize.hs b/src/Shikumi/Optimize.hs --- a/src/Shikumi/Optimize.hs +++ b/src/Shikumi/Optimize.hs @@ -11,8 +11,11 @@ -- it over the dataset via 'scoreOn'). This is the "thread candidates explicitly" -- discipline the MasterPlan mandates. module Shikumi.Optimize - ( -- * The driver + ( module Shikumi.Optimize.Structure, + + -- * The driver optimize, + optimizeWith, -- * Re-exports module Shikumi.Optimize.Types, @@ -32,7 +35,7 @@ import Effectful (Eff, (:>)) import Effectful.Concurrent (Concurrent) -import Effectful.Error.Static (Error) +import Effectful.Error.Static (Error, throwError) import Effectful.Prim (Prim) import Shikumi.Compile.Types (CompiledProgram) import Shikumi.Effect.Time (Time) @@ -42,6 +45,7 @@ import Shikumi.Optimize.Bootstrap import Shikumi.Optimize.COPRO import Shikumi.Optimize.Ensemble +import Shikumi.Optimize.Execution import Shikumi.Optimize.GEPA import Shikumi.Optimize.Instruction import Shikumi.Optimize.KNN @@ -49,7 +53,9 @@ import Shikumi.Optimize.MIPRO import Shikumi.Optimize.Pareto import Shikumi.Optimize.RandomSearch +import Shikumi.Optimize.Report (CandidateStatus (Unscored), OptimizationReport (..), RunStatus (BudgetStopped)) import Shikumi.Optimize.Search +import Shikumi.Optimize.Structure import Shikumi.Optimize.Types import Shikumi.Program (Program) @@ -64,3 +70,20 @@ Program i o -> Eff es (CompiledProgram i o) optimize opt train metric prog = runOptimizer opt train metric prog + +-- | Run a configured strategy with diagnostic reporting and hard operation admission. +optimizeWith :: + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + RunConfig -> + ConfiguredOptimizer i o -> + Dataset i o -> + Metric o -> + Program i o -> + Eff es (CompiledProgram i o, OptimizationReport) +optimizeWith cfg opt ds metric prog = do + (result, report) <- runSearchSession cfg $ \session -> runConfiguredOptimizer opt session ds metric prog + case result of + Right compiled -> pure (compiled, report) + Left e + | runStatus report == BudgetStopped -> pure (freezeProgram prog, report {selectedCandidate = Nothing, resultStatus = Just Unscored, selectionReason = "Unscored baseline: strategy stopped before returning a result"}) + | otherwise -> throwError e diff --git a/src/Shikumi/Optimize/Bootstrap.hs b/src/Shikumi/Optimize/Bootstrap.hs --- a/src/Shikumi/Optimize/Bootstrap.hs +++ b/src/Shikumi/Optimize/Bootstrap.hs @@ -1,45 +1,46 @@ +{-# LANGUAGE GADTs #-} {-# LANGUAGE ScopedTypeVariables #-} --- | M2 — bootstrap few-shot. Run a /teacher/ program over the training set, keep --- the runs the metric judged correct, and attach those input/output pairs as --- demonstrations to the /student/. This "bootstraps" high-quality demos from the --- program's own successful behaviour (DSPy's @BootstrapFewShot@). --- --- __Adapted to the delivered substrate.__ The plan envisioned recovering a demo --- for every /internal/ node by reading EP-7's trace tree, keyed by a per-node --- @NodePath@, via a @runProgramTraced@. The delivered EP-7 --- (@docs/plans/7-…replay.md@) records LM-call spans by opaque @SpanId@ with the --- canonical-request and raw-response JSON, but provides neither a --- @NodePath@↔program-node correlation nor a @runProgramTraced@, and the recorded --- prompt is the rendered wire request, not the structured typed input. So this --- plan recovers demos at the __program-I/O level__: a demo is the pair of the --- example's input and the teacher's produced output ('recoverDemo'), and the kept --- demos are attached to every node (the DSPy default for multi-module programs). --- This is faithful to bootstrap's user-visible behaviour and is exactly what the --- single-node acceptance test (M5) verifies; per-internal-node recovery waits on --- EP-7/EP-4 exposing a node-correlated trace. See the plan's Decision Log. +-- | Recover typed demonstrations from accepted teacher invocations, keyed by +-- student node. Composite programs require explicit capture codecs. module Shikumi.Optimize.Bootstrap ( bootstrapFewShot, bootstrapFewShotWith, bootstrapKeptDemos, + bootstrapNodeDemos, + bootstrapDemosFor, + withNodeDemos, + NodeBootstrapConfig (..), + defaultNodeBootstrapConfig, + BootstrapReport (..), BootstrapConfig (..), defaultBootstrapConfig, recoverDemo, ) where -import Data.Aeson (ToJSON, toJSON) +import Control.Monad (unless, when) +import Data.Aeson (ToJSON, Value (Bool), toJSON) +import Data.List (nub, sortOn) +import Data.Map.Strict qualified as Map +import Data.Maybe (isNothing) +import Data.Text (Text) +import Data.Text qualified as T import Effectful (Eff, (:>)) -import Effectful.Error.Static (Error, catchError) +import Effectful.Error.Static (Error, throwError) import Effectful.Prim (Prim) import GHC.Generics (Generic) -import Shikumi.Error (ShikumiError) +import Shikumi.Error (ShikumiError (..)) import Shikumi.Eval (Dataset, Example (..), Metric, datasetExamples, prediction, unScore) import Shikumi.LLM (LLM) -import Shikumi.Optimize.LabeledFewShot (withDemos) import Shikumi.Optimize.Search (BudgetMeter, freezeProgram, newBudgetMeter, tryCharge) import Shikumi.Optimize.Types (Budget (..), Optimizer (..)) -import Shikumi.Program (Demo (..), Program, foldParams, runProgram) +import Shikumi.Program (CaptureCodec (..), Demo (..), Params (..), Program (..), ProgramShape (..), foldParams, mapParamsAt, programShape) +import Shikumi.Schema (FromModel, deriveSchema, fromModel) +import Shikumi.Schema.Types (fieldName, objectSchema) +import Shikumi.Signature (Signature, inputFields) +import Shikumi.Trace.Node (NodePath (..), programNodePaths, renderNodePath) +import Shikumi.Trace.Observation (NodeObservation (..), runProgramObserved) -- | Tunables for a bootstrap search. data BootstrapConfig = BootstrapConfig @@ -81,38 +82,182 @@ Optimizer i o bootstrapFewShotWith cfg teacher budget = Optimizer $ \train metric student -> do meter <- newBudgetMeter budget - kept <- bootstrapKeptDemos cfg meter teacher train metric - pure (freezeProgram (withDemos kept student)) + (kept, _) <- bootstrapDemosFor (defaultNodeBootstrapConfig {nodeBootstrapConfig = cfg}) meter teacher train metric student + pure (freezeProgram (withNodeDemos 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 -> +-- | Separate configuration keeps existing BootstrapConfig record construction +-- source compatible. Explicit mappings are teacher-to-student; omitted nodes +-- receive no recovered demos. Duplicate targets require explicit merge consent. +data NodeBootstrapConfig = NodeBootstrapConfig + { nodeBootstrapConfig :: !BootstrapConfig, + nodeMapping :: !(Maybe [(NodePath, NodePath)]), + mergeTargetMappings :: !Bool, + nodeSeed :: !(Maybe Int) + } + deriving stock (Eq, Show) + +defaultNodeBootstrapConfig :: NodeBootstrapConfig +defaultNodeBootstrapConfig = NodeBootstrapConfig defaultBootstrapConfig Nothing False Nothing + +-- | Runtime rejection diagnostics, including metric failures and opaque limits. +newtype BootstrapReport = BootstrapReport {bootstrapDiagnostics :: [Text]} + deriving stock (Eq, Show) + +data LeafEvidence = LeafEvidence (Maybe (Value, Value)) (Demo -> Either ShikumiError ()) + +leafEvidence :: forall i o. (FromModel i, FromModel o) => Maybe (CaptureCodec i o) -> Signature i o -> LeafEvidence +leafEvidence codec _ = LeafEvidence ((\c -> (captureInputSchema c, captureOutputSchema c)) <$> codec) $ \(Demo i o) -> do + _ <- (fromModel i :: Either ShikumiError i) + _ <- (fromModel o :: Either ShikumiError o) + pure () + +evidence :: Program i o -> [LeafEvidence] +evidence = go + where + go :: Program x y -> [LeafEvidence] + go (Predict sig _) = [leafEvidence Nothing sig] + go (PredictCaptured codec sig _) = [leafEvidence (Just codec) sig] + go (Compose a b) = go a ++ go b + go (FMap _ p) = go p + go (Map _ p) = go p + go (Parallel a b) = go a ++ go b + go (Retry _ p) = go p + go (RetryWhen _ _ p) = go p + go (Validate _ p) = go p + go (MajorityVote _ _ _ p) = go p + go (Ensemble ps _) = concatMap go ps + go (Embed _) = [] + +-- Opaque limits are reported even if no teacher invocation is necessary. +hasOpaque :: ProgramShape -> Bool +hasOpaque ShapeEmbed = True +hasOpaque (ShapePredict _) = False +hasOpaque (ShapeCompose a b) = hasOpaque a || hasOpaque b +hasOpaque (ShapeFMap p) = hasOpaque p +hasOpaque (ShapeMap _ p) = hasOpaque p +hasOpaque (ShapeParallel a b) = hasOpaque a || hasOpaque b +hasOpaque (ShapeRetry _ p) = hasOpaque p +hasOpaque (ShapeRetryWhen _ p) = hasOpaque p +hasOpaque (ShapeValidate p) = hasOpaque p +hasOpaque (ShapeMajorityVote _ _ p) = hasOpaque p +hasOpaque (ShapeEnsemble ps) = any hasOpaque ps + +-- | Install only the pool addressed to each node, preserving template codecs. +withNodeDemos :: Map.Map NodePath [Demo] -> Program i o -> Program i o +withNodeDemos pools program = foldl install program (zip [0 ..] (programNodePaths program)) + where + install p (idx, path) = mapParamsAt idx (\ps -> ps {demos = Map.findWithDefault [] path pools}) p + +-- | Capture-capable entry point. Mapping/schema failures happen before budget +-- charges or LM calls; declared schemas are also checked by target decoding. +bootstrapNodeDemos :: + (LLM :> es, Error ShikumiError :> es, Prim :> es) => + NodeBootstrapConfig -> 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 - let s = unScore (metric expd (prediction out)) - pure [recoverDemo inp out | s >= passThreshold cfg] - ) - `catchError` \_ (_ :: ShikumiError) -> pure [] - collect kept [] - | length kept >= cap = pure (take cap kept) - | otherwise = pure kept - collect kept (ex : rest) - | length kept >= cap = pure (take cap kept) + Program j k -> + Eff es (Map.Map NodePath [Demo], BootstrapReport) +bootstrapNodeDemos cfg meter teacher train metric student = do + let teachers = Map.fromList (zip (programNodePaths teacher) (evidence teacher)) + students = Map.fromList (zip (programNodePaths student) (evidence student)) + failWith = throwError . ValidationFailure + settings = nodeBootstrapConfig cfg + cap = max 0 (maxBootstrappedDemos settings) + pairs <- case nodeMapping cfg of + Just mappings -> pure mappings + Nothing -> do + unless (programShape teacher == programShape student) $ + failWith "bootstrap: different teacher/student structures require an explicit nodeMapping" + pure [(p, p) | p <- Map.keys students] + unless (mergeTargetMappings cfg || length (map snd pairs) == length (nub (map snd pairs))) $ + failWith "bootstrap: duplicate student targets require mergeTargetMappings" + when (any (\(LeafEvidence schema _) -> isNothing schema) (Map.elems teachers ++ Map.elems students)) $ + failWith "bootstrap: composite programs require predictCaptured at every predictor (missing capture codec/schema)" + mappings <- + traverse + ( \(src, dst) -> case (Map.lookup src teachers, Map.lookup dst students) of + (Just (LeafEvidence (Just ts) _), Just (LeafEvidence (Just ss) decode)) + | ts == ss -> pure (src, dst, decode) + | otherwise -> failWith ("bootstrap: incompatible schemas for " <> renderNodePath src <> " -> " <> renderNodePath dst) + _ -> failWith "bootstrap: unknown mapping path or missing capture schema" + ) + pairs + let emptyPools = Map.fromList [(p, []) | p <- Map.keys students] + recover (pools, notes) obs = + foldl + (attach obs) + (pools, notes) + [(dst, decode) | (src, dst, decode) <- mappings, src == observationPath obs] + attach obs (pools, notes) (dst, decode) + | not (observationEligible obs) = (pools, notes) + | otherwise = case (observationInput obs, observationOutput obs) of + (Just i, Just o) -> + let demo = Demo i o + in case decode demo of + Left err -> (pools, notes ++ ["bootstrap: rejected captured demo at " <> renderNodePath dst <> ": " <> T.pack (show err)]) + Right () -> (Map.adjust (++ [demo]) dst pools, notes) + _ -> (pools, notes ++ ["bootstrap: missing encoded invocation at " <> renderNodePath dst]) + full pools = all (\(_, dst, _) -> length (Map.findWithDefault [] dst pools) >= cap) mappings + collect acc [] = pure acc + collect acc@(pools, notes) (Example inp expd : rest) + | null mappings || cap == 0 || (isNothing (nodeSeed cfg) && full pools) = pure acc | otherwise = do - fits <- tryCharge meter teacherCost + fits <- tryCharge meter (max 1 (length (foldParams teacher))) if not fits - then pure kept + then pure acc else do - newKept <- keepIfPassing ex - collect (kept ++ newKept) rest - collect [] (datasetExamples train) + (result, obs) <- runProgramObserved teacher inp + let next = case result of + Left err -> (pools, notes ++ ["bootstrap: teacher example failed: " <> T.pack (show err)]) + Right out | unScore (metric expd (prediction out)) >= passThreshold settings -> foldl recover (pools, notes) obs + Right _ -> (pools, notes ++ ["bootstrap: teacher example below metric threshold"]) + collect next rest + let limits = ["bootstrap: Embed is opaque; internal predictors cannot supply demonstrations" | hasOpaque (programShape teacher) || hasOpaque (programShape student)] + (pools, notes) <- collect (emptyPools, limits) (datasetExamples train) + pure (Map.mapWithKey (select cap (nodeSeed cfg)) pools, BootstrapReport notes) + +-- Independent, reproducible streams use the stable target path, never map index. +select :: Int -> Maybe Int -> NodePath -> [Demo] -> [Demo] +select cap Nothing _ = take cap +select cap (Just seed) path = take cap . map snd . sortOn fst . zip keys + where + start = T.foldl' (\n c -> (n * 131 + toInteger (fromEnum c)) `mod` 2147483648) (toInteger seed) (renderNodePath path) + keys = drop 1 (iterate (\n -> (1103515245 * n + 12345) `mod` 2147483648) start) + +-- | Legacy outer encoders apply only to a bare single prediction. Signature +-- field metadata supplies matching evidence for this same-typed compatibility +-- case; all composite programs use explicit CaptureCodec schema evidence. +bootstrapDemosFor :: + (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es, Prim :> es) => + NodeBootstrapConfig -> + BudgetMeter -> + Program i o -> + Dataset i o -> + Metric o -> + Program i o -> + Eff es (Map.Map NodePath [Demo], BootstrapReport) +bootstrapDemosFor cfg meter teacher train metric student = case (teacher, student) of + (Predict {}, Predict {}) -> bootstrapNodeDemos cfg meter (legacy teacher) train metric (legacy student) + (Predict sig ps, PredictCaptured codec _ _) -> bootstrapNodeDemos cfg meter (PredictCaptured (codec {encodeCaptureInput = toJSON, encodeCaptureOutput = toJSON}) sig ps) train metric student + (PredictCaptured codec _ _, Predict sig ps) -> bootstrapNodeDemos cfg meter teacher train metric (PredictCaptured (codec {encodeCaptureInput = toJSON, encodeCaptureOutput = toJSON}) sig ps) + _ -> bootstrapNodeDemos cfg meter teacher train metric student + where + legacy :: forall x y. (ToJSON x, ToJSON y) => Program x y -> Program x y + legacy (Predict sig ps) = PredictCaptured (CaptureCodec toJSON toJSON (objectSchema [(name, Bool True) | name <- map fieldName (inputFields sig)] (map fieldName (inputFields sig))) (deriveSchema @y)) sig ps + legacy p = p + +-- | Compatibility pool of outer input/output demos, for a bare single predictor +-- only. Composite callers must migrate to bootstrapNodeDemos/bootstrapDemosFor. +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 = case teacher of + Predict {} -> recover teacher + PredictCaptured _ sig ps -> recover (Predict sig ps) + _ -> throwError (ValidationFailure "bootstrapKeptDemos: composite programs require node-local bootstrapDemosFor") + where + recover outerTeacher = do + (pools, _) <- bootstrapDemosFor (defaultNodeBootstrapConfig {nodeBootstrapConfig = cfg}) meter outerTeacher train metric outerTeacher + pure (Map.findWithDefault [] (NodePath []) pools) diff --git a/src/Shikumi/Optimize/Execution.hs b/src/Shikumi/Optimize/Execution.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Execution.hs @@ -0,0 +1,332 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Shared execution and accounting, independent of any mutation strategy. +module Shikumi.Optimize.Execution + ( RunLimits (..), + RunConfig (..), + defaultRunConfig, + validateRunConfig, + SearchSession, + CandidateId, + sessionLimits, + runSearchSession, + markLegacy, + setSelection, + sessionStopped, + canStartCandidate, + addPredictedWork, + reserveCandidate, + annotateCandidate, + remainingCandidates, + evaluateCandidate, + evaluateCandidates, + ObjectiveMetric, + ExampleMeasurement (..), + scalarObjectives, + ) +where + +import Control.Monad (forM, unless, when) +import Data.Aeson (FromJSON, ToJSON) +import Data.List (sortOn) +import Data.Map.Strict qualified as Map +import Data.Set qualified as Set +import Data.Text (Text) +import Effectful (Eff, Effect, (:>)) +import Effectful.Concurrent (Concurrent, ThreadId, myThreadId) +import Effectful.Concurrent.Async (mapConcurrently) +import Effectful.Concurrent.QSem (QSem, newQSem, signalQSem, waitQSem) +import Effectful.Dispatch.Dynamic (interpose) +import Effectful.Error.Static (Error, throwError) +import Effectful.Exception qualified as E +import Effectful.Prim.IORef (IORef, Prim, atomicModifyIORef', newIORef, readIORef) +import GHC.Generics (Generic) +import Shikumi.Effect.Time (Time, getMonotonicTimeNSec) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.Eval (Dataset, Example (..), Metric, datasetExamples, datasetSize, prediction, unScore) +import Shikumi.Eval.Evaluate (scoreExecution, tryShikumi) +import Shikumi.Eval.Report (FailurePolicy (..)) +import Shikumi.Eval.Usage (withUsageTotals) +import Shikumi.LLM (LLM (..), complete, stream) +import Shikumi.Optimize.Feedback (EvaluationEvidence (..)) +import Shikumi.Optimize.Report +import Shikumi.Trace.Observation (NodeObservation) + +-- | Serializable controls; the executable observer is deliberately separate. +data RunLimits = RunLimits + { operationLimit :: !Int, + candidateLimit :: !Int, + evaluationConcurrency :: !Int, + deterministicSeed :: !Int + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +data RunConfig = RunConfig + { runLimits :: !RunLimits, + eventSink :: forall es. (Concurrent :> es, Prim :> es) => OptimizationEvent -> Eff es () + } + +defaultRunConfig :: RunConfig +defaultRunConfig = RunConfig (RunLimits 200 32 1 1) (const (pure ())) + +validateRunConfig :: RunConfig -> Either ShikumiError () +validateRunConfig cfg = + unless + (operationLimit l >= 0 && candidateLimit l >= 0 && evaluationConcurrency l > 0) + (Left (ValidationFailure "execution limits must be nonnegative, concurrency positive")) + where + l = runLimits cfg + +data SearchSession (es :: [Effect]) = SearchSession + { config :: !RunConfig, + state :: !(IORef OptimizationReport), + reserved :: !(IORef Int), + operationCollectors :: !(IORef (Map.Map ThreadId [IORef Int])), + startedCandidates :: !(IORef (Set.Set Int)), + stopped :: !(IORef Bool), + dispatchPermits :: !QSem, + observerLock :: !QSem + } + +-- | An opaque reservation tied to exactly one run. It cannot be forged from a +-- report's numeric ID, reused, or transferred to another session. +data CandidateId = CandidateId !(IORef Int) !Int deriving stock (Eq) + +sessionLimits :: SearchSession es -> RunLimits +sessionLimits = runLimits . config + +modifyReport :: (Prim :> es) => SearchSession es -> (OptimizationReport -> OptimizationReport) -> Eff es () +modifyReport s f = atomicModifyIORef' (state s) (\r -> (f r, ())) + +emit :: (Concurrent :> es, Prim :> es) => SearchSession es -> EventKind -> Eff es () +emit s kind = E.bracket_ (waitQSem (observerLock s)) (signalQSem (observerLock s)) $ do + ev <- atomicModifyIORef' (state s) $ \r -> + let e = OptimizationEvent (length (events r)) kind in (r {events = events r ++ [e]}, e) + eventSink (config s) ev `E.catchSync` \_ -> modifyReport s (\r -> r {observerFailures = observerFailures r + 1}) + +sessionStopped :: (Prim :> es) => SearchSession es -> Eff es Bool +sessionStopped = readIORef . stopped + +stop :: (Prim :> es) => SearchSession es -> Eff es () +stop s = atomicModifyIORef' (stopped s) (const (True, ())) + +-- | The action's typed error is returned alongside diagnostics. Cancellation and +-- host exceptions propagate after terminal bookkeeping; observers are best effort. +runSearchSession :: + forall es a. + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Prim :> es) => + RunConfig -> (SearchSession es -> Eff es a) -> Eff es (Either ShikumiError a, OptimizationReport) +runSearchSession cfg action = do + either throwError pure (validateRunConfig cfg) + ref <- newIORef (OptimizationReport 1 Completed 0 0 (Map.fromList [("operationLimit", operationLimit (runLimits cfg)), ("candidateLimit", candidateLimit (runLimits cfg)), ("evaluationConcurrency", evaluationConcurrency (runLimits cfg)), ("deterministicSeed", deterministicSeed (runLimits cfg))]) Nothing [] Map.empty [] [] Nothing (Just Unscored) "Unscored baseline" True "unspecified" 0 []) + count <- newIORef 0 + started <- newIORef Set.empty + collectors <- newIORef Map.empty + halted <- newIORef False + terminal <- newIORef False + permits <- newQSem (evaluationConcurrency (runLimits cfg)) + lock <- newQSem 1 + let s = SearchSession cfg ref count collectors started halted permits lock + finish status = do + first <- atomicModifyIORef' terminal (\done -> (True, not done)) + when first $ do + when (status == BudgetStopped) (emit s BudgetStop) + allocated <- readIORef count + startedIds <- readIORef started + modifyReport s (\r -> r {runStatus = status, candidates = sortOn candidateId (candidates r), unexecutedReservations = [ix | ix <- [0 .. allocated - 1], Set.notMember ix startedIds]}) + emit s (RunFinished status) + dispatch :: forall x. Eff es x -> Eff es x + dispatch op = E.bracket_ (waitQSem (dispatchPermits s)) (signalQSem (dispatchPermits s)) $ E.mask $ \restore -> do + admitted <- atomicModifyIORef' ref $ \r -> + if admittedOperations r < operationLimit (runLimits cfg) + then (r {admittedOperations = admittedOperations r + 1}, True) + else (r, False) + unless admitted (stop s >> throwError (BudgetExceeded "optimizer operation admission exhausted")) + tid <- myThreadId + localCollectors <- Map.findWithDefault [] tid <$> readIORef collectors + mapM_ (\counter -> atomicModifyIORef' counter (\n -> (n + 1, ()))) localCollectors + restore op + + result <- + ( do + emit s RunStarted + when (operationLimit (runLimits cfg) == 0 || candidateLimit (runLimits cfg) == 0) (stop s) + outcome <- + tryShikumi $ + interpose + ( \_ -> \case + Complete m c o -> dispatch (complete m c o) + Stream m c o -> dispatch (stream m c o) + ) + (action s) + exhausted <- sessionStopped s + current <- readIORef ref + let onlyFailures = not (null (candidates current)) && all ((== CandidateFailed) . candidateStatus) (candidates current) + finish + ( case outcome of + Left e + | exhausted && e == BudgetExceeded "optimizer operation admission exhausted" -> BudgetStopped + | otherwise -> Failed + Right _ + | exhausted || runStatus current == BudgetStopped -> BudgetStopped + | onlyFailures -> Failed + | otherwise -> Completed + ) + pure outcome + ) + `E.withException` (\(e :: E.SomeException) -> finish (if E.isAsyncException e then Cancelled else Failed)) + report <- readIORef ref + pure (result, report) + +markLegacy :: (Prim :> es) => SearchSession es -> Eff es () +markLegacy s = modifyReport s (\r -> r {candidateDetailAvailable = False, resultStatus = Nothing, selectionReason = "Opaque legacy optimizer; candidate details unavailable", validationMode = "legacy optimizer controlled"}) + +setSelection :: (Prim :> es) => SearchSession es -> Text -> ObjectivePolicy -> Eff es () +setSelection s mode policy = modifyReport s $ \r -> + let winner = selectObjectiveWinner policy (candidates r) + in r + { frontier = map candidateId (objectiveFrontier policy (candidates r)), + selectedCandidate = candidateId <$> winner, + resultStatus = Just (maybe Unscored (const CandidateCompleted) winner), + selectionReason = maybe "Unscored baseline: no eligible completed candidate" (const ("Pareto frontier; primary objective " <> primaryObjective policy <> "; ordered ties; creation order")) winner, + reportedPolicy = Just policy, + validationMode = mode + } + +addPredictedWork :: (Prim :> es) => SearchSession es -> Int -> Eff es () +addPredictedWork s n = modifyReport s (\r -> r {predictedWork = predictedWork r + max 0 n}) + +-- | Unreserved candidate slots in this session. +remainingCandidates :: (Prim :> es) => SearchSession es -> Eff es Int +remainingCandidates s = (\n -> max 0 (candidateLimit (sessionLimits s) - n)) <$> readIORef (reserved s) + +-- | Reserve IDs in scheduling order before spawning any workers. +reserveCandidate :: (Prim :> es) => SearchSession es -> Eff es (Maybe CandidateId) +reserveCandidate s = do + halted <- sessionStopped s + if halted + then pure Nothing + else do + ix <- atomicModifyIORef' (reserved s) $ \n -> if n < candidateLimit (sessionLimits s) then (n + 1, Just n) else (n, Nothing) + when (ix == Nothing) (modifyReport s (\r -> r {runStatus = BudgetStopped})) + pure (CandidateId (reserved s) <$> ix) + +data ExampleMeasurement o = ExampleMeasurement + {evidence :: !(EvaluationEvidence o), operations :: !Int, latencySeconds :: !Double} + deriving stock (Eq, Show) + +type ObjectiveMetric es o = o -> ExampleMeasurement o -> Eff es ObjectiveValues + +scalarObjectives :: (Applicative m) => Metric o -> o -> ExampleMeasurement o -> m ObjectiveValues +scalarObjectives metric expected measured = pure $ Map.singleton "quality" $ case executionResult (evidence measured) of + Left _ -> 0 + Right out -> unScore (metric expected (prediction out)) + +-- | All validation positions are required. The runner preserves root errors and +-- observations; isolated collectors measure each example, including nested calls. +evaluateCandidate :: + forall es i o. + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + SearchSession es -> + CandidateId -> + Dataset i o -> + (i -> Eff es (Either ShikumiError o, [NodeObservation])) -> + (ShikumiError -> FailurePolicy) -> + Metric o -> + ObjectivePolicy -> + ObjectiveMetric es o -> + Eff es CandidateReport +evaluateCandidate s (CandidateId owner ident) ds runner classify metric policy objective = E.mask $ \restore -> do + unless (owner == reserved s) (throwError (ValidationFailure "candidate reservation belongs to another session")) + fresh <- atomicModifyIORef' (startedCandidates s) (\seen -> (Set.insert ident seen, Set.notMember ident seen)) + unless fresh (throwError (ValidationFailure "candidate reservation has already executed")) + rowsRef <- newIORef [] + calls <- newIORef 0 + ended <- newIORef False + let finish status reason = do + first <- atomicModifyIORef' ended (\b -> (True, not b)) + rows <- reverse <$> readIORef rowsRef + n <- readIORef calls + let aggregates = aggregateObjectives policy [vs | (_, _, vs) <- rows] + report = CandidateReport ident status (datasetSize ds) (length rows) [(ix, sc) | (ix, sc, _) <- rows] (either (const Map.empty) id aggregates) n reason + when first $ do + modifyReport s (\r -> r {candidates = report : candidates r}) + emit s (CandidateEnded ident status) + pure report + -- Register the collector for this dispatch, including inherited handlers + -- in child threads. Only the admission boundary increments it: waiting, + -- denied and cancelled-before-admission calls are never counted. + counted :: forall x. Eff es x -> Eff es x + counted op = do + tid <- myThreadId + E.bracket_ + (atomicModifyIORef' (operationCollectors s) (\m -> (Map.insertWith (++) tid [calls] m, ()))) + (atomicModifyIORef' (operationCollectors s) (\m -> (Map.update (\xs -> case drop 1 xs of [] -> Nothing; rest -> Just rest) tid m, ()))) + op + count act = + interpose + ( \_ -> \case + Complete m c o -> counted (complete m c o) + Stream m c o -> counted (stream m c o) + ) + act + body = forM (zip [0 ..] (datasetExamples ds)) $ \(ix, Example inp expected) -> do + before <- readIORef calls + start <- getMonotonicTimeNSec + ((out, obs), usage) <- withUsageTotals (count (runner inp)) + end <- getMonotonicTimeNSec + after <- readIORef calls + let ev = EvaluationEvidence ix out obs usage + failurePolicy BudgetExceeded {} = FailAbort + failurePolicy e = classify e + (_, (score, _)) <- scoreExecution failurePolicy executionResult (pure ev) (pure . metric expected . prediction) + vs <- objective expected (ExampleMeasurement ev (after - before) (fromIntegral (end - start) / 1e9)) + exhausted <- sessionStopped s + unless exhausted $ atomicModifyIORef' rowsRef (\rows -> ((ix, unScore score, vs) : rows, ())) + when exhausted (throwError (BudgetExceeded "optimizer operation admission exhausted")) + outcome <- (emit s (CandidateStarted ident) >> restore (tryShikumi body)) `E.onException` (finish CandidateIncomplete (Just "execution interrupted") >> pure ()) + halted <- sessionStopped s + case outcome of + Left e + | halted && e == BudgetExceeded "optimizer operation admission exhausted" -> finish CandidateIncomplete (Just "operation budget exhausted") + | otherwise -> do + r <- finish CandidateFailed (Just "typed execution failure") + case e of + BudgetExceeded {} -> throwError e + _ -> case classify e of FailAbort -> throwError e; _ -> pure r + Right _ | halted -> finish CandidateIncomplete (Just "operation budget exhausted") + Right _ -> do + rows <- readIORef rowsRef + case aggregateObjectives policy [vs | (_, _, vs) <- rows] of + Left reason -> finish CandidateFailed (Just reason) + Right _ -> finish CandidateCompleted Nothing + +-- | Finite batches bound candidate jobs; examples within each job are sequential. +-- The independent dispatch semaphore also bounds nested Program/Embed concurrency. +evaluateCandidates :: (Concurrent :> es) => SearchSession es -> (a -> Eff es b) -> [a] -> Eff es [b] +evaluateCandidates s f = go + where + go [] = pure [] + go xs = do + let (batch, rest) = splitAt (evaluationConcurrency (sessionLimits s)) xs + done <- mapConcurrently f batch + (done ++) <$> go rest + +-- | Attach caller-owned non-sensitive identity before starting a candidate. +annotateCandidate :: (Concurrent :> es, Prim :> es, Error ShikumiError :> es) => SearchSession es -> CandidateId -> Map.Map Text Text -> Eff es () +annotateCandidate s (CandidateId owner ident) metadata = do + unless (owner == reserved s) (throwError (ValidationFailure "candidate reservation belongs to another session")) + modifyReport s (\r -> r {candidateMetadata = Map.insert ident metadata (candidateMetadata r)}) + emit s (CandidateMetadata ident metadata) + +-- | Check operation capacity before starting another scheduled candidate. +-- Completed candidates are unaffected when the last slot was used exactly. +canStartCandidate :: (Prim :> es) => SearchSession es -> Eff es Bool +canStartCandidate s = do + halted <- sessionStopped s + current <- readIORef (state s) + let available = not halted && admittedOperations current < operationLimit (sessionLimits s) + unless available (modifyReport s (\r -> r {runStatus = BudgetStopped})) + pure available diff --git a/src/Shikumi/Optimize/Feedback.hs b/src/Shikumi/Optimize/Feedback.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Feedback.hs @@ -0,0 +1,161 @@ +-- | Indexed execution evidence and explicitly attributed, bounded feedback. +module Shikumi.Optimize.Feedback + ( FeedbackMetric, + EvidenceMetric, + EvaluationEvidence (..), + NodeFeedback (..), + Provenance (..), + FeedbackResult (..), + FeedbackConfig (..), + defaultFeedbackConfig, + candidateFailurePolicy, + legacyFeedback, + validateFeedback, + validateFeedbackConfig, + boundText, + captureEvidence, + ) +where + +import Control.Monad (forM, unless) +import Data.Text (Text) +import Data.Text qualified as T +import Effectful (Eff, (:>)) +import Effectful.Error.Static (Error, throwError) +import Effectful.Prim (Prim) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.Eval (Dataset, Example (..), Prediction, Score, datasetExamples, prediction, scoreZero) +import Shikumi.Eval.Evaluate (scoreExecution, tryShikumi) +import Shikumi.Eval.Report (FailurePolicy (..), FailureReason (..), UsageTotals) +import Shikumi.Eval.Usage (withUsageTotals) +import Shikumi.LLM (LLM) +import Shikumi.Program (Program) +import Shikumi.Trace.Node (NodePath, programNodePaths) +import Shikumi.Trace.Observation (NodeObservation (..), runProgramObserved) + +type FeedbackMetric o = o -> Prediction o -> (Score, Text) + +type EvidenceMetric es o = o -> EvaluationEvidence o -> Eff es FeedbackResult + +data EvaluationEvidence o = EvaluationEvidence + { exampleIndex :: !Int, + executionResult :: !(Either ShikumiError o), + observations :: ![NodeObservation], + executionUsage :: !UsageTotals + } + deriving stock (Eq, Show) + +data Provenance = Caller | Model | LegacyProgram deriving stock (Eq, Show) + +data NodeFeedback = NodeFeedback + { feedbackExample :: !Int, + feedbackPath :: !NodePath, + feedbackInvocation :: !(Maybe Int), + critique :: !Text, + provenance :: !Provenance + } + deriving stock (Eq, Show) + +data FeedbackResult = FeedbackResult + { overallScore :: !Score, + programCritique :: !(Maybe (Provenance, Text)), + nodeCritiques :: ![NodeFeedback] + } + deriving stock (Eq, Show) + +-- | Redaction runs on the entire rendered reflection payload before truncation. +-- It does not erase the raw evidence returned to the caller. +data FeedbackConfig = FeedbackConfig + { critiqueCharacters :: !Int, + reflectionExamples :: !Int, + reflectionCharacters :: !Int, + includeProgramCritique :: !Bool, + redactEvidence :: Text -> Text, + failureClassification :: ShikumiError -> FailurePolicy + } + +defaultFeedbackConfig :: FeedbackConfig +defaultFeedbackConfig = FeedbackConfig 2000 4 8000 False id (candidateFailurePolicy scoreZero) + +-- | Only candidate/output errors are scored by default. Budget exhaustion always +-- escapes, even when a caller supplies a broader classifier. +candidateFailurePolicy :: Score -> ShikumiError -> FailurePolicy +candidateFailurePolicy s = \case + InvalidJSON {} -> FailScore s + MissingField {} -> FailScore s + SchemaMismatch {} -> FailScore s + ValidationFailure {} -> FailScore s + _ -> FailAbort + +-- | Legacy critiques remain program-scoped, including for single predictors. +legacyFeedback :: (Applicative m) => FeedbackMetric o -> o -> EvaluationEvidence o -> m FeedbackResult +legacyFeedback metric expected ev = pure $ case executionResult ev of + Left _ -> FeedbackResult scoreZero Nothing [] + Right out -> + let (s, t) = metric expected (prediction out) + in FeedbackResult s (if T.null t then Nothing else Just (LegacyProgram, t)) [] + +validateFeedbackConfig :: FeedbackConfig -> Either ShikumiError () +validateFeedbackConfig cfg = + unless + (all (>= 0) [critiqueCharacters cfg, reflectionExamples cfg, reflectionCharacters cfg]) + (Left (ValidationFailure "feedback: character and example bounds must be nonnegative")) + +-- | Unicode-safe prefix with an in-budget truncation marker, including odd/zero limits. +boundText :: Int -> Text -> Text +boundText n t + | n <= 0 = "" + | T.length t <= n = t + | otherwise = T.take (n - 1) t <> "…" + +validateFeedback :: FeedbackConfig -> [NodePath] -> EvaluationEvidence o -> FeedbackResult -> Either ShikumiError FeedbackResult +validateFeedback cfg paths ev result = do + validateFeedbackConfig cfg + critiques <- forM (nodeCritiques result) $ \fb -> do + unless + (feedbackExample fb == exampleIndex ev && feedbackPath fb `elem` paths && any (matches fb) (observations ev)) + (Left (ValidationFailure ("feedback: target was not executed in example " <> T.pack (show (exampleIndex ev)) <> ": " <> T.pack (show fb)))) + unless + (provenance fb /= LegacyProgram) + (Left (ValidationFailure "feedback: LegacyProgram cannot assert node attribution")) + pure fb {critique = boundText (critiqueCharacters cfg) (critique fb)} + pure result {nodeCritiques = critiques, programCritique = fmap (\(p, t) -> (p, boundText (critiqueCharacters cfg) t)) (programCritique result)} + where + matches fb obs = + not (observationOpaque obs) + && observationPath obs == feedbackPath fb + && maybe True (== observationInvocation obs) (feedbackInvocation fb) + +-- | Sequential, position-stable capture. Each envelope retains the original root +-- error and retry lineage. Metric errors are separately identified; configuration +-- errors escape. Host exceptions (including cancellation) are never intercepted. +captureEvidence :: + (LLM :> es, Error ShikumiError :> es, Prim :> es) => + FeedbackConfig -> + Dataset i o -> + EvidenceMetric es o -> + Program i o -> + Eff es [(EvaluationEvidence o, FeedbackResult, Maybe FailureReason)] +captureEvidence cfg ds metric prog = do + either throwError pure (validateFeedbackConfig cfg) + forM (zip [0 ..] (datasetExamples ds)) $ \(ix, Example inp expected) -> do + ((out, obs), usage) <- withUsageTotals (runProgramObserved prog inp) + let ev = EvaluationEvidence ix out obs usage + policy BudgetExceeded {} = FailAbort + policy e = failureClassification cfg e + (_, (rootScore, rootFailure)) <- scoreExecution policy executionResult (pure ev) (const (pure scoreZero)) + -- Keep the full feedback value while reusing the evaluator's metric boundary. + (judged, (metricScore, metricFailure)) <- + scoreExecution + policy + id + (tryShikumi (metric expected ev)) + (pure . overallScore) + case judged of + Left _ -> pure (ev, FeedbackResult metricScore Nothing [], fmap asMetric metricFailure) + Right fb -> do + valid <- either throwError pure (validateFeedback cfg (programNodePaths prog) ev fb) + pure (ev, case out of Left _ -> valid {overallScore = rootScore}; Right _ -> valid, rootFailure) + where + asMetric (ProgramError t) = MetricError t + asMetric reason = reason diff --git a/src/Shikumi/Optimize/GEPA.hs b/src/Shikumi/Optimize/GEPA.hs --- a/src/Shikumi/Optimize/GEPA.hs +++ b/src/Shikumi/Optimize/GEPA.hs @@ -1,27 +1,8 @@ +{-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} --- | GEPA (EP-22): a reflective, evolutionary instruction optimizer. Where greedy --- coordinate ascent is blind, GEPA is /reflective/: it runs the program while --- capturing, per node, a short natural-language critique ("feedback") of how that --- node performed, then reflects on those critiques to propose a rewritten --- instruction. Where greedy search keeps one best program, GEPA keeps a __Pareto --- frontier__ (see "Shikumi.Optimize.Pareto") of candidates none strictly worse than --- another across the per-example score vector, samples a parent from it, mutates one --- node by reflection, scores the child, and folds it back in — until a 'Budget' is --- spent. --- --- GEPA consumes EP-16's per-node feedback channel ('attachFeedback'/'feedbackFor' --- keyed by 'NodePath', with node identity from 'programNodePaths') and EP-19's --- summaries (here via small in-package fallbacks). The 'Trace'/'Feedback' effects are --- discharged /internally/ (via 'runFeedback' against the ambient 'Prim'), so the --- public 'Optimizer' row is unchanged (MasterPlan integration point #4/#5). Feedback --- is attached at the program level to every node (the DSPy default and the M1 --- baseline); node-specific critique from per-node sub-traces is a documented deferral. --- --- Output is V1's 'Shikumi.Compile.Types.CompiledProgram' via 'freezeProgram'; the --- frontier is internal bookkeeping, not part of the returned type. GEPA reuses V1's --- @Metric@/@Score@ plus a critique @Text@ (its 'FeedbackMetric') rather than a --- parallel reward type (MasterPlan integration point #1). +-- | Reflective evolution with failure-aware, node-grounded evidence. Legacy +-- callbacks produce explicitly program-scoped critiques, never node attribution. module Shikumi.Optimize.GEPA ( FeedbackMetric, ReflectIn (..), @@ -30,36 +11,46 @@ captureFeedback, mutateNode, gepa, + gepaWith, + GEPAConfig (..), + ObjectiveCallback (..), + defaultGEPAConfig, + gepaWithFeedback, + FeedbackCallback (..), + mutateFromEvidence, ) where import Control.Monad (forM, forM_, when) +import Data.Either (isRight) +import Data.List (find, findIndex, sortOn) +import Data.Maybe (catMaybes) import Data.Text (Text) import Data.Text qualified as T import Effectful (Eff, (:>)) -import Effectful.Error.Static (Error) +import Effectful.Concurrent (Concurrent) +import Effectful.Error.Static (Error, throwError) import Effectful.Prim (Prim) import GHC.Generics (Generic) import Shikumi.Adapter (ToPrompt) -import Shikumi.Error (ShikumiError) +import Shikumi.Effect.Time (Time) +import Shikumi.Error (ShikumiError (..)) import Shikumi.Eval ( Dataset, - Example (..), - ExampleResult (..), - Prediction, - Report (..), - Score, + dataset, datasetExamples, datasetSize, - evaluatePure, - prediction, unScore, ) +import Shikumi.Eval.Evaluate (tryShikumi) import Shikumi.LLM (LLM) import Shikumi.Module (predict) +import Shikumi.Optimize.Execution qualified as X +import Shikumi.Optimize.Feedback import Shikumi.Optimize.Pareto (Candidate (..), paretoFrontier, sampleParent) +import Shikumi.Optimize.Report qualified as R import Shikumi.Optimize.Search (effectiveInstructionAt, freezeProgram, newBudgetMeter, scoringCost, setNodeInstrIfNew, tryCharge) -import Shikumi.Optimize.Types (Budget (..), Optimizer (..)) +import Shikumi.Optimize.Types (Budget (..), ConfiguredOptimizer (..), Optimizer (..)) import Shikumi.Program ( NodeFields (..), Program, @@ -71,11 +62,8 @@ import Shikumi.Schema (FromModel, ToSchema, Validatable) import Shikumi.Signature (mkSignature) import Shikumi.Trace.Feedback (FeedbackLog, attachFeedback, feedbackFor, runFeedback) -import Shikumi.Trace.Node (NodePath, programNodePaths) - --- | A feedback metric: like V1's @Metric@ but also emits a short critique. Reuses --- @Score@ (EP-18's reward vocabulary reduces to this) plus a critique @Text@. -type FeedbackMetric o = o -> Prediction o -> (Score, Text) +import Shikumi.Trace.Node (NodePath (..), programNodePaths) +import Shikumi.Trace.Observation (NodeObservation (..), runProgramObserved) -- --------------------------------------------------------------------------- -- The reflective proposer @@ -124,25 +112,17 @@ -- M1 — feedback capture -- --------------------------------------------------------------------------- --- | Run the program over the whole dataset, attaching the feedback metric's critique --- (when non-empty) to every node keyed by its 'NodePath', and returning the --- 'FeedbackLog' alongside the per-example score vector (for the Pareto frontier). +-- | Compatibility projection. Program critique is stored once at the root key, +-- labeled with its provenance; use 'captureEvidence' to retain full attribution. captureFeedback :: (LLM :> es, Error ShikumiError :> es, Prim :> es) => - Dataset i o -> - FeedbackMetric o -> - Program i o -> - Eff es (FeedbackLog, [Double]) + Dataset i o -> FeedbackMetric o -> Program i o -> Eff es (FeedbackLog, [Double]) captureFeedback ds fm prog = do - let paths = programNodePaths prog - (scores, fblog) <- - runFeedback $ - forM (datasetExamples ds) $ \(Example inp expd) -> do - out <- runProgram prog inp - let (sc, crit) = fm expd (prediction out) - when (not (T.null crit)) (forM_ paths (\p -> attachFeedback p crit)) - pure (unScore sc) - pure (fblog, scores) + captured <- captureEvidence defaultFeedbackConfig ds (legacyFeedback fm) prog + (_, logbook) <- runFeedback $ forM_ captured $ \(_, fb, _) -> + forM_ (programCritique fb) $ \(source, t) -> + when (not (T.null t)) (attachFeedback (NodePath []) ("program (" <> tshow source <> "): " <> t)) + pure (logbook, [unScore (overallScore fb) | (_, fb, _) <- captured]) -- --------------------------------------------------------------------------- -- M2 — reflective mutation @@ -202,10 +182,32 @@ Budget -> Optimizer i o gepa proposer fbMetric budget = Optimizer $ \train metric student -> do + -- Preserve the legacy all-or-nothing predicted seed gate. + if datasetSize train == 0 || maxLmCalls budget < scoringCost train student || maxCandidates budget <= 0 + then pure (freezeProgram student) + else do + let controls = X.defaultRunConfig {X.runLimits = X.RunLimits (max 0 (maxLmCalls budget)) (max 0 (maxCandidates budget)) 1 1} + cfg = + (defaultGEPAConfig (FeedbackCallback (legacyFeedback fbMetric))) + { feedbackConfig = defaultFeedbackConfig {includeProgramCritique = True}, + minibatchSize = datasetSize train + } + (result, _) <- X.runSearchSession controls $ \session -> runConfiguredOptimizer (gepaWith cfg proposer) session train metric student + either throwError pure result + +-- | An effectful callback portable across the optimizer's existing effect row. +newtype FeedbackCallback o = FeedbackCallback + { runFeedbackCallback :: + forall es. + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + EvidenceMetric es o + } + +gepaWithFeedback :: FeedbackConfig -> Program ReflectIn ReflectOut -> FeedbackCallback o -> Budget -> Optimizer i o +gepaWithFeedback cfg proposer callback budget = Optimizer $ \train metric student -> do + either throwError pure (validateFeedbackConfig cfg) meter <- newBudgetMeter budget let paths = programNodePaths student - fields = nodeFieldsIndexed student - nNodes = max 1 (length paths) progSummary = fallbackProgramSummary (length paths) dataSummary = fallbackDatasetSummary (datasetSize train) maxCands = maxCandidates budget @@ -216,8 +218,11 @@ if not seedFits then pure (freezeProgram student) else do - seedRpt <- evaluatePure train metric student - let seedCand = Candidate (foldParams student) (perEx seedRpt) (aggregateScore seedRpt) + seedRows <- captureEvidence cfg train (legacyFeedback (\e p -> (metric e p, ""))) student + let candidateFrom prog rows = + let scores = [unScore (overallScore fb) | (_, fb, _) <- rows] + in Candidate (foldParams prog) scores (if null scores then 0 else sum scores / fromIntegral (length scores)) + seedCand = candidateFrom student seedRows -- A full step costs: capture + child evaluation over the whole dataset, -- plus one reflective proposer call. @@ -235,27 +240,94 @@ 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) + captured <- captureEvidence cfg train (runFeedbackCallback callback) parentProg + child <- mutateFromEvidence cfg proposer progSummary dataSummary captured step parentProg + if foldParams child == foldParams parentProg + then loop (step + 1) cands seed' frontier + else do + rows <- captureEvidence cfg train (legacyFeedback (\e p -> (metric e p, ""))) child + let childCand = candidateFrom child rows frontier' = paretoFrontier (childCand : frontier) loop (step + 1) (childCand : cands) seed' frontier' best <- loop 0 [seedCand] 1 [seedCand] pure (freezeProgram (rebuild best)) --- | The per-example score vector from a report, in dataset order. -perEx :: Report -> [Double] -perEx rpt = [unScore s | ExampleResult {score = s} <- results rpt] +-- | Reflect only on executed nodes with attributed critiques (or explicitly +-- enabled program fallback). Redaction covers all evidence, errors and critiques +-- before it reaches the proposer. Rejected retries remain labeled evidence. +mutateFromEvidence :: + (LLM :> es, Error ShikumiError :> es) => + FeedbackConfig -> + Program ReflectIn ReflectOut -> + Text -> + Text -> + [(EvaluationEvidence o, FeedbackResult, a)] -> + Int -> + Program i o -> + Eff es (Program i o) +mutateFromEvidence cfg proposer progSummary dataSummary rows step prog = do + either throwError pure (validateFeedbackConfig cfg) + validated <- mapM (\(ev, fb, _) -> (ev,) <$> either throwError pure (validateFeedback cfg paths ev fb)) rows + let relevant path ev fb = + [ obs + | obs <- observations ev, + observationPath obs == path, + not (observationOpaque obs), + any (\f -> feedbackPath f == path && not (T.null (critique f)) && maybe True (== observationInvocation obs) (feedbackInvocation f)) (nodeCritiques fb) + || (includeProgramCritique cfg && maybe False (not . T.null . snd) (programCritique fb)) + ] + evidence path = [(ev, fb, obs) | (ev, fb) <- validated, obs <- relevant path ev fb] + eligible = [p | p <- paths, not (null (evidence p))] + case eligible of + [] -> pure prog + _ | reflectionExamples cfg == 0 || reflectionCharacters cfg == 0 -> pure prog + _ -> do + let path = eligible !! (max 0 step `mod` length eligible) + local = sortOn (\(_, _, obs) -> (isRight (observationStatus obs), null (observationRejectedBy obs))) (evidence path) + chosen = take (reflectionExamples cfg) local + render (ev, fb, obs) = + "example " + <> tshow (exampleIndex ev) + <> "; invocation " + <> tshow (observationInvocation obs) + <> "; status: " + <> tshow (observationStatus obs) + <> "; rejected scopes: " + <> tshow (observationRejectedBy obs) + <> "\nnode critiques: " + <> T.intercalate + "\n" + [ tshow (provenance f) <> ": " <> critique f + | f <- nodeCritiques fb, + feedbackPath f == path, + maybe True (== observationInvocation obs) (feedbackInvocation f) + ] + <> (if includeProgramCritique cfg then "\nprogram critique: " <> maybe "" tshow (programCritique fb) else "") + <> "\ninput: " + <> maybe (tshow (observationInputFields obs)) tshow (observationInput obs) + <> "\noutput: " + <> maybe (tshow (observationOutputFields obs)) tshow (observationOutput obs) + payload = + T.intercalate "\n\n" (map render chosen) + <> if length chosen < length local then "\n[examples truncated]" else "" + clean = boundText (reflectionCharacters cfg) . redactEvidence cfg + case findIndex (== path) paths of + Nothing -> pure prog + Just idx -> do + ReflectOut newInstruction <- + runProgram + proposer + ( ReflectIn + (clean (effectiveInstructionAt idx prog)) + (clean payload) + (clean progSummary) + (clean dataSummary) + (clean (renderFields (drop idx (nodeFieldsIndexed prog)))) + ) + pure (setNodeInstrIfNew idx newInstruction prog) + where + paths = programNodePaths prog -- | The frontier candidate with the highest aggregate (earliest on ties); falls back -- to the seed if the frontier is somehow empty. @@ -272,3 +344,129 @@ tshow :: (Show a) => a -> Text tshow = T.pack . show + +-- | Callbacks are trusted code. The framework sends only training evidence to +-- reflection; it is not a security sandbox around caller closures. +data GEPAConfig i o = GEPAConfig + { validationDataset :: !(Maybe (Dataset i o)), + feedbackConfig :: !FeedbackConfig, + feedbackCallback :: !(FeedbackCallback o), + objectivePolicy :: !R.ObjectivePolicy, + objectiveCallback :: !(Maybe (ObjectiveCallback o)), + minibatchSize :: !Int, + childrenPerGeneration :: !Int + } + +newtype ObjectiveCallback o = ObjectiveCallback + { runObjectiveCallback :: + forall es. + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + X.ObjectiveMetric es o + } + +defaultGEPAConfig :: FeedbackCallback o -> GEPAConfig i o +defaultGEPAConfig callback = GEPAConfig Nothing defaultFeedbackConfig callback R.qualityPolicy Nothing 4 1 + +-- | Configured reflective evolution, with full validation before frontier entry. +-- Children are proposed serially from a generation snapshot, then scored in +-- bounded batches. Width one retains adaptive single-child evolution. +gepaWith :: GEPAConfig i o -> Program ReflectIn ReflectOut -> ConfiguredOptimizer i o +gepaWith cfg proposer = ConfiguredOptimizer $ \session train metric student -> do + either throwError pure (validateFeedbackConfig (feedbackConfig cfg)) + either (throwError . ValidationFailure) pure (R.validateObjectives (objectivePolicy cfg)) + when (datasetSize train == 0 || maybe False ((== 0) . datasetSize) (validationDataset cfg)) $ + throwError (ValidationFailure "GEPA training and explicit validation datasets must be nonempty") + when (minibatchSize cfg <= 0 || childrenPerGeneration cfg <= 0) $ + throwError (ValidationFailure "GEPA minibatch and generation sizes must be positive") + let validation = maybe train id (validationDataset cfg) + mode = maybe "training-as-validation compatibility" (const "explicit validation") (validationDataset cfg) + policy = objectivePolicy cfg + minibatch = dataset (take (minibatchSize cfg) (datasetExamples train)) + objectivesFor expected measured = case objectiveCallback cfg of + Nothing -> X.scalarObjectives metric expected measured + Just callback -> runObjectiveCallback callback expected measured + evaluate ident prog = do + X.addPredictedWork session (scoringCost validation prog) + report <- + X.evaluateCandidate + session + ident + validation + (runProgramObserved prog) + (failureClassification (feedbackConfig cfg)) + metric + policy + objectivesFor + pure (report, prog) + best completed = case R.selectObjectiveWinner policy (map fst completed) of + Nothing -> student + Just winner -> maybe student snd (find (\(r, _) -> R.candidateId r == R.candidateId winner) completed) + propose step parent = do + X.addPredictedWork session (2 * scoringCost minibatch parent + 1) + captured <- captureEvidence (feedbackConfig cfg) minibatch (runFeedbackCallback (feedbackCallback cfg)) parent + child <- + mutateFromEvidence + (feedbackConfig cfg) + proposer + (fallbackProgramSummary (length (programNodePaths student))) + (fallbackDatasetSummary (datasetSize train)) + captured + step + parent + -- The screen verifies training execution/feedback before expensive full + -- validation, without rejecting a child solely for lower training quality. + screened <- captureEvidence (feedbackConfig cfg) minibatch (legacyFeedback (\e p -> (metric e p, ""))) child + when (all (\(ev, _, _) -> not (isRight (executionResult ev))) screened) $ + throwError (ValidationFailure "GEPA training screen has no successful executions") + pure child + loop step completed = do + halted <- X.sessionStopped session + remaining <- X.remainingCandidates session + if halted || remaining == 0 + then do + when (remaining == 0) (X.reserveCandidate session >> pure ()) + pure completed + else do + let front = R.objectiveFrontier policy (map fst completed) + parents = [p | (r, p) <- completed, R.candidateId r `elem` map R.candidateId front] + parent = if null parents then best completed else parents !! ((X.deterministicSeed (X.sessionLimits session) + step) `mod` length parents) + proposals <- tryShikumi $ forM [0 .. min remaining (childrenPerGeneration cfg) - 1] $ \offset -> do + ident <- X.reserveCandidate session + case ident of + Nothing -> pure Nothing + Just ix -> do + result <- tryShikumi (propose (step + offset) parent) + case result of + Right child -> pure (Just (ix, child)) + Left e -> do + -- Close a reserved proposal as an incomplete/failed candidate + -- through the same generic lifecycle boundary. + _ <- + X.evaluateCandidate + session + ix + validation + (\_ -> throwError e) + (failureClassification (feedbackConfig cfg)) + metric + policy + objectivesFor + pure Nothing + case proposals of + Left e -> do + haltedNow <- X.sessionStopped session + if haltedNow && e == BudgetExceeded "optimizer operation admission exhausted" then pure completed else throwError e + Right pending -> do + results <- X.evaluateCandidates session (uncurry evaluate) (catMaybes pending) + let completed' = completed ++ results + X.setSelection session mode policy + loop (step + childrenPerGeneration cfg) completed' + X.setSelection session mode policy + seedId <- X.reserveCandidate session + completed <- case seedId of + Nothing -> pure [] + Just ident -> do + seed <- evaluate ident student + loop 0 [seed] + X.setSelection session mode policy + pure (freezeProgram (best completed)) diff --git a/src/Shikumi/Optimize/MIPRO.hs b/src/Shikumi/Optimize/MIPRO.hs --- a/src/Shikumi/Optimize/MIPRO.hs +++ b/src/Shikumi/Optimize/MIPRO.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE GADTs #-} {-# LANGUAGE ScopedTypeVariables #-} -- | MIPROv2 (EP-20): the joint instruction-and-demonstration optimizer. Where @@ -45,11 +46,12 @@ import Control.Monad (forM) import Data.Aeson (ToJSON) import Data.Aeson.Text (encodeToLazyText) +import Data.Map.Strict qualified as Map import Data.Text (Text) import Data.Text.Lazy qualified as TL import Effectful (Eff, (:>)) import Effectful.Concurrent (Concurrent) -import Effectful.Error.Static (Error, catchError) +import Effectful.Error.Static (Error) import Effectful.Prim (Prim) import GHC.Generics (Generic) import Shikumi.Effect.Time (Time) @@ -61,11 +63,9 @@ dataset, datasetExamples, datasetSize, - prediction, - unScore, ) import Shikumi.LLM (LLM) -import Shikumi.Optimize.Bootstrap (recoverDemo) +import Shikumi.Optimize.Bootstrap (BootstrapConfig (BootstrapConfig), NodeBootstrapConfig (..), bootstrapDemosFor, defaultNodeBootstrapConfig, recoverDemo) import Shikumi.Optimize.Propose ( ProposeRequest (..), ProposeResult (..), @@ -76,11 +76,11 @@ import Shikumi.Program ( Demo (..), Params (..), - Program, + Program (..), foldParams, mapParamsAt, - runProgram, ) +import Shikumi.Trace.Node (programNodePaths) -- --------------------------------------------------------------------------- -- Configuration and presets @@ -147,8 +147,8 @@ -- | For each node (in @foldParams@ order), a list of candidate demo sets. Candidate -- 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. +-- (bootstrapped). Labelled outer pairs are candidates only for a bare single +-- prediction; composite nodes receive only their own recovered invocations. bootstrapDemoCandidates :: (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es, Prim :> es) => Miprov2Config -> @@ -175,32 +175,18 @@ 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 - let s = unScore (metric expd (prediction out)) - pure [recoverDemo inp out | s >= bootstrapThreshold cfg] - ) - `catchError` \_ (_ :: ShikumiError) -> pure [] - 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 - nodeSets ps = + bootCfg = defaultNodeBootstrapConfig {nodeBootstrapConfig = BootstrapConfig (bootstrapThreshold cfg) cap} + (pools, _) <- bootstrapDemosFor bootCfg meter teacher train metric student + let labeledSet = case student of + Predict {} -> take cap [recoverDemo i o | Example i o <- datasetExamples train] + PredictCaptured {} -> take cap [recoverDemo i o | Example i o <- datasetExamples train] + _ -> [] + nodeSets path ps = take (max 1 (numDemoCandidates cfg)) - (dedup (demos ps : [] : filter (not . null) [labeledSet, bootSet])) - pure (map nodeSets (foldParams student)) + (dedup (demos ps : [] : filter (not . null) [labeledSet, Map.findWithDefault [] path pools])) + pure (zipWith nodeSets (programNodePaths student) (foldParams student)) -- --------------------------------------------------------------------------- -- Phase 2 — propose instruction candidates diff --git a/src/Shikumi/Optimize/RandomSearch.hs b/src/Shikumi/Optimize/RandomSearch.hs --- a/src/Shikumi/Optimize/RandomSearch.hs +++ b/src/Shikumi/Optimize/RandomSearch.hs @@ -21,8 +21,7 @@ import Data.List (sortBy) import Data.Ord (comparing) import Shikumi.Eval (Dataset, dataset, datasetExamples) -import Shikumi.Optimize.Bootstrap (bootstrapKeptDemos, defaultBootstrapConfig) -import Shikumi.Optimize.LabeledFewShot (withDemos) +import Shikumi.Optimize.Bootstrap (NodeBootstrapConfig (..), bootstrapDemosFor, defaultBootstrapConfig, defaultNodeBootstrapConfig, withNodeDemos) import Shikumi.Optimize.Search (freezeProgram, meteredScore, newBudgetMeter, selectBestMetered) import Shikumi.Optimize.Types (Budget (..), Optimizer (..), Scored (..)) import Shikumi.Program (Program) @@ -84,8 +83,8 @@ let seeds = [1 .. max 1 numCandidates] candidateFor seed = do let cfg' = defaultBootstrapConfig & #maxBootstrappedDemos .~ sizeFor cfg seed - demos <- bootstrapKeptDemos cfg' meter teacher (shuffle seed train) metric - pure (withDemos demos student) + (demos, _) <- bootstrapDemosFor (defaultNodeBootstrapConfig {nodeBootstrapConfig = cfg', nodeSeed = Just seed}) meter teacher (shuffle seed train) metric student + pure (withNodeDemos demos student) seeded <- mapM candidateFor seeds let cands = student : seeded -- zero-shot baseline first best <- selectBestMetered meter (\p -> meteredScore meter train metric p) cands diff --git a/src/Shikumi/Optimize/Report.hs b/src/Shikumi/Optimize/Report.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Report.hs @@ -0,0 +1,202 @@ +-- | Versioned diagnostic metadata. No prompts, examples or executable closures. +module Shikumi.Optimize.Report + ( Direction (..), + Aggregation (..), + MissingPolicy (..), + ObjectiveSpec (..), + ObjectivePolicy (..), + ObjectiveValues, + qualityPolicy, + validateObjectives, + aggregateObjectives, + objectiveFrontier, + selectObjectiveWinner, + RunStatus (..), + CandidateStatus (..), + CandidateReport (..), + EventKind (..), + OptimizationEvent (..), + OptimizationReport (..), + ) +where + +import Control.Monad (forM, unless) +import Data.Aeson (FromJSON (..), ToJSON, withObject, (.!=), (.:), (.:?)) +import Data.List (minimumBy, nub, sortOn) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Maybe (mapMaybe) +import Data.Ord (comparing) +import Data.Text (Text) +import GHC.Generics (Generic) + +data Direction = Maximize | Minimize + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +data Aggregation = Mean | Total | Worst + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +data MissingPolicy = Required | Substitute Double + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +data ObjectiveSpec = ObjectiveSpec + { objectiveId :: !Text, + unit :: !Text, + direction :: !Direction, + aggregation :: !Aggregation, + missingPolicy :: !MissingPolicy, + lowerBound :: !(Maybe Double), + upperBound :: !(Maybe Double) + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +data ObjectivePolicy = ObjectivePolicy + { objectives :: ![ObjectiveSpec], + primaryObjective :: !Text, + tieBreakObjectives :: ![Text] + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +type ObjectiveValues = Map Text Double + +qualityPolicy :: ObjectivePolicy +qualityPolicy = ObjectivePolicy [ObjectiveSpec "quality" "score" Maximize Mean Required (Just 0) (Just 1)] "quality" [] + +finite :: Double -> Bool +finite x = not (isNaN x || isInfinite x) + +validateObjectives :: ObjectivePolicy -> Either Text () +validateObjectives p = do + let ids = map objectiveId (objectives p) + ordered = primaryObjective p : tieBreakObjectives p + unless (not (null ids) && all (/= "") ids && nub ids == ids) (Left "objectives must have distinct nonempty IDs") + unless (all (`elem` ids) ordered && nub ordered == ordered) (Left "selection objectives must be declared and distinct") + mapM_ check (objectives p) + where + check s = do + unless (all finite (mapMaybe id [lowerBound s, upperBound s]) && case missingPolicy s of Required -> True; Substitute v -> finite v) (Left "objective configuration must be finite") + unless (case (lowerBound s, upperBound s) of (Just l, Just u) -> l <= u; _ -> True) (Left "objective bounds are reversed") + +aggregateObjectives :: ObjectivePolicy -> [ObjectiveValues] -> Either Text ObjectiveValues +aggregateObjectives p rows = do + validateObjectives p + unless (not (null rows)) (Left "no complete evaluation rows") + unless (all (all finite . Map.elems) rows) (Left "non-finite objective value") + Map.fromList + <$> forM + (objectives p) + ( \s -> do + xs <- forM rows $ \row -> case Map.lookup (objectiveId s) row of + Just v -> Right v + Nothing -> case missingPolicy s of Required -> Left "required objective missing"; Substitute v -> Right v + let v = case aggregation s of + Mean -> sum xs / fromIntegral (length xs) + Total -> sum xs + Worst -> (if direction s == Maximize then minimum else maximum) xs + unless (finite v) (Left "non-finite objective aggregate") + pure (objectiveId s, v) + ) + +eligible :: ObjectivePolicy -> CandidateReport -> Bool +eligible p c = candidateStatus c == CandidateCompleted && all valid (objectives p) + where + valid s = case Map.lookup (objectiveId s) (objectiveValues c) of + Nothing -> False + Just v -> finite v && maybe True (v >=) (lowerBound s) && maybe True (v <=) (upperBound s) + +objectiveFrontier :: ObjectivePolicy -> [CandidateReport] -> [CandidateReport] +objectiveFrontier p cs = filter (\c -> not (any (`dominates` c) valid)) valid + where + valid = sortOn candidateId (filter (eligible p) cs) + oriented s c = (if direction s == Maximize then negate else id) (Map.findWithDefault 0 (objectiveId s) (objectiveValues c)) + dominates a b = + let pairs = [(oriented s a, oriented s b) | s <- objectives p] + in all (uncurry (<=)) pairs && any (uncurry (<)) pairs + +selectObjectiveWinner :: ObjectivePolicy -> [CandidateReport] -> Maybe CandidateReport +selectObjectiveWinner p cs = case objectiveFrontier p cs of + [] -> Nothing + xs -> Just (minimumBy (comparing key) xs) + where + key c = ([value ident c | ident <- primaryObjective p : tieBreakObjectives p], candidateId c) + value ident c = + let v = Map.findWithDefault 0 ident (objectiveValues c) + in if any (\s -> objectiveId s == ident && direction s == Maximize) (objectives p) then negate v else v + +data RunStatus = Completed | BudgetStopped | Failed | Cancelled + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +data CandidateStatus = Unscored | CandidateCompleted | CandidateFailed | CandidateIncomplete + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +data CandidateReport = CandidateReport + { candidateId :: !Int, + candidateStatus :: !CandidateStatus, + requiredExamples :: !Int, + completedExamples :: !Int, + exampleScores :: ![(Int, Double)], + objectiveValues :: !ObjectiveValues, + candidateOperations :: !Int, + candidateReason :: !(Maybe Text) + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +data EventKind = RunStarted | CandidateMetadata Int (Map Text Text) | CandidateStarted Int | CandidateEnded Int CandidateStatus | BudgetStop | RunFinished RunStatus + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +data OptimizationEvent = OptimizationEvent {eventId :: !Int, eventKind :: !EventKind} + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +data OptimizationReport = OptimizationReport + { reportVersion :: !Int, + runStatus :: !RunStatus, + admittedOperations :: !Int, + predictedWork :: !Int, + runControls :: !(Map Text Int), + reportedPolicy :: !(Maybe ObjectivePolicy), + candidates :: ![CandidateReport], + candidateMetadata :: !(Map Int (Map Text Text)), + unexecutedReservations :: ![Int], + frontier :: ![Int], + selectedCandidate :: !(Maybe Int), + resultStatus :: !(Maybe CandidateStatus), + selectionReason :: !Text, + candidateDetailAvailable :: !Bool, + validationMode :: !Text, + observerFailures :: !Int, + events :: ![OptimizationEvent] + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +instance FromJSON OptimizationReport where + parseJSON = withObject "OptimizationReport" $ \o -> do + version <- o .: "reportVersion" + unless (version == (1 :: Int)) (fail "unsupported optimization report version") + OptimizationReport version + <$> o .: "runStatus" + <*> o .: "admittedOperations" + <*> o .: "predictedWork" + <*> o .: "runControls" + <*> o .: "reportedPolicy" + <*> o .: "candidates" + <*> (o .:? "candidateMetadata" .!= Map.empty) + <*> o .: "unexecutedReservations" + <*> o .: "frontier" + <*> o .: "selectedCandidate" + <*> o .: "resultStatus" + <*> o .: "selectionReason" + <*> o .: "candidateDetailAvailable" + <*> o .: "validationMode" + <*> o .: "observerFailures" + <*> o .: "events" diff --git a/src/Shikumi/Optimize/Structure.hs b/src/Shikumi/Optimize/Structure.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Structure.hs @@ -0,0 +1,89 @@ +-- | Experimental finite structure selection using shared operation admission. +module Shikumi.Optimize.Structure (StructureSearchResult (..), structureSearchWith) where + +import Control.Monad (forM, unless) +import Data.List.NonEmpty qualified as NE +import Data.Map.Strict qualified as Map +import Data.Maybe (catMaybes) +import Data.Text qualified as T +import Effectful (Eff, (:>)) +import Effectful.Concurrent (Concurrent) +import Effectful.Error.Static (Error, throwError) +import Effectful.Prim (Prim) +import Shikumi.Compile.Structure +import Shikumi.Compile.Types (CompiledProgram (..)) +import Shikumi.Effect.Time (Time) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.Eval (Dataset, Metric, datasetSize) +import Shikumi.Eval.Report (FailurePolicy) +import Shikumi.LLM (LLM) +import Shikumi.Optimize.Execution +import Shikumi.Optimize.Report +import Shikumi.Optimize.Search (scoringCost) +import Shikumi.Trace.Observation (runProgramObserved) + +data StructureSearchResult i o = StructureSearchResult + { selectedRecipeId :: !RecipeId, + selectedRecipeRevision :: !RecipeRevision, + selectedStructure :: !(CompiledProgram i o), + structureReport :: !OptimizationReport + } + +-- | Training is validated but never evaluated or used for ranking in this finite +-- enumeration. Every supplied recipe is scored on validation only. Callbacks are +-- trusted application code. Width one gives reproducible dispatch scheduling; +-- ties always retain registry order, including with concurrent completion. +structureSearchWith :: + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + RunConfig -> + Dataset i o -> + Dataset i o -> + Metric o -> + (ShikumiError -> FailurePolicy) -> + ObjectivePolicy -> + ObjectiveMetric es o -> + StructureRegistry i o -> + Eff es (StructureSearchResult i o) +structureSearchWith cfg training validation metric classify policy objective registry = do + either throwError pure (validateRunConfig cfg) + unless + (datasetSize training > 0 && datasetSize validation > 0) + (throwError (ValidationFailure "structure search requires nonempty training and validation datasets")) + either (throwError . ValidationFailure) pure (validateObjectives policy) + (outcome, report) <- runSearchSession cfg $ \session -> do + jobs <- + catMaybes + <$> forM + (NE.toList (registryRecipes registry)) + ( \recipe -> do + reservation <- reserveCandidate session + forM reservation $ \ident -> do + annotateCandidate session ident (Map.fromList [("recipeId", recipeIdText (recipeId recipe)), ("recipeRevision", T.pack (show (recipeRevisionNumber (recipeRevision recipe)))), ("registryId", registryIdText (registryId registry))]) + pure (ident, recipe) + ) + completedRows <- + evaluateCandidates + session + ( \(ident, recipe) -> do + available <- canStartCandidate session + if not available + then pure Nothing + else do + addPredictedWork session (scoringCost validation (recipeProgram recipe)) + row <- evaluateCandidate session ident validation (runProgramObserved (recipeProgram recipe)) classify metric policy objective + pure (Just (recipe, row)) + ) + jobs + setSelection session "explicit validation; experimental structure search" policy + let rows = catMaybes completedRows + winner = selectObjectiveWinner policy (map snd rows) + pure $ case winner of + Nothing -> NE.head (registryRecipes registry) + Just row -> maybe (NE.head (registryRecipes registry)) fst (findRow (candidateId row) rows) + recipe <- either throwError pure outcome + pure (StructureSearchResult (recipeId recipe) (recipeRevision recipe) (CompiledProgram (recipeProgram recipe)) report) + where + findRow _ [] = Nothing + findRow ident (x@(_, row) : xs) + | candidateId row == ident = Just x + | otherwise = findRow ident xs diff --git a/src/Shikumi/Optimize/Types.hs b/src/Shikumi/Optimize/Types.hs --- a/src/Shikumi/Optimize/Types.hs +++ b/src/Shikumi/Optimize/Types.hs @@ -7,10 +7,12 @@ -- and a starting 'Program', it proposes new node parameters (instructions and -- few-shot demonstrations), scores each candidate by running the program over the -- dataset, and returns the best-scoring 'CompiledProgram' it found. An optimizer --- never changes a program's structure or types — only its parameters — so the --- optimized program is the same typed function, merely better-behaved. +-- normally changes parameters while preserving boundary types. -- --- Two optimizers are explicit structure-changing exceptions. 'Shikumi.Optimize.KNN.knnFewShot' +-- The additive 'Shikumi.Optimize.Structure.structureSearchWith' API selects among +-- caller-registered typed implementations using shared validation and admission; +-- its experimental structure artifacts restore through that same registry. +-- Two legacy optimizers are also 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' @@ -41,6 +43,8 @@ -- the plan's Decision Log. module Shikumi.Optimize.Types ( Optimizer (..), + ConfiguredOptimizer (..), + fromLegacyOptimizer, Budget (..), defaultBudget, Scored (..), @@ -51,11 +55,12 @@ import Effectful.Concurrent (Concurrent) import Effectful.Error.Static (Error) import Effectful.Prim (Prim) -import Shikumi.Compile.Types (CompiledProgram) +import Shikumi.Compile.Types (CompiledProgram (..)) import Shikumi.Effect.Time (Time) import Shikumi.Error (ShikumiError) import Shikumi.Eval (Dataset, Metric) import Shikumi.LLM (LLM) +import Shikumi.Optimize.Execution (SearchSession, markLegacy, sessionStopped) import Shikumi.Program (Program) -- | A search strategy that, given a training dataset, a metric, and a starting @@ -109,3 +114,20 @@ score :: !Double } deriving stock (Eq, Show) + +-- | A strategy receives an explicit, shared execution session. +newtype ConfiguredOptimizer i o = ConfiguredOptimizer + { runConfiguredOptimizer :: + forall es. + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + SearchSession es -> Dataset i o -> Metric o -> Program i o -> Eff es (CompiledProgram i o) + } + +fromLegacyOptimizer :: Optimizer i o -> ConfiguredOptimizer i o +fromLegacyOptimizer opt = ConfiguredOptimizer $ \session ds metric prog -> do + halted <- sessionStopped session + if halted + then pure (CompiledProgram prog) + else do + markLegacy session + runOptimizer opt ds metric prog diff --git a/test/BootstrapSpec.hs b/test/BootstrapSpec.hs --- a/test/BootstrapSpec.hs +++ b/test/BootstrapSpec.hs @@ -14,6 +14,7 @@ import Effectful.Concurrent (Concurrent, runConcurrent) import Effectful.Error.Static (Error, runErrorNoCallStack) import Effectful.Prim (Prim, runPrim) +import NodeBootstrapSpec qualified import Shikumi.Compile.Types (CompiledProgram, Compiler (runCompiler), compiledProgram) import Shikumi.Compile.ZeroShot (zeroShot) import Shikumi.Effect.Time (Time, runTime) @@ -64,7 +65,8 @@ tests = testGroup "M2 bootstrap" - [ testCase "recoverDemo produces a round-trippable typed demo" $ do + [ NodeBootstrapSpec.tests, + testCase "recoverDemo produces a round-trippable typed demo" $ do let d = recoverDemo (Sentence "good film") (Label "positive") (fromModel (input d) :: Either ShikumiError Sentence) @?= Right (Sentence "good film") (fromModel (output d) :: Either ShikumiError Label) @?= Right (Label "positive"), @@ -84,7 +86,7 @@ let budget = Budget {maxLmCalls = 6, maxCandidates = 32} res <- runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ - runStubLMCounting ref (optimize (bootstrapFewShot sentimentPipeline budget) budgetTrainset exactMatch sentimentProg) + runStubLMCounting ref (optimize (bootstrapFewShot sentimentPipeline budget) budgetTrainset exactMatch sentimentPipeline) case res of Left e -> assertFailure ("unexpected error: " <> show e) Right _ -> pure () diff --git a/test/ExecutionSpec.hs b/test/ExecutionSpec.hs new file mode 100644 --- /dev/null +++ b/test/ExecutionSpec.hs @@ -0,0 +1,295 @@ +{-# LANGUAGE TypeApplications #-} + +module ExecutionSpec (tests) where + +import Control.Monad (forM, replicateM_) +import Data.Aeson (eitherDecode, encode) +import Data.Either (isLeft) +import Effectful (Eff, IOE, runEff, (:>)) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Concurrent.Async (cancel, mapConcurrently, waitCatch, withAsync) +import Effectful.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar) +import Effectful.Dispatch.Dynamic (interpose) +import Effectful.Error.Static (Error, catchError, runErrorNoCallStack, throwError) +import Effectful.Exception qualified as E +import Effectful.Prim (Prim, runPrim) +import Effectful.Prim.IORef qualified as Ref +import Shikumi.Effect.Time (Time, runTime) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.Eval (Dataset, FailurePolicy, Metric, dataset, exactMatch, example, scoreZero) +import Shikumi.LLM (LLM (..), complete, stream) +import Shikumi.Optimize (Optimizer (..), freezeProgram, fromLegacyOptimizer, optimizeWith) +import Shikumi.Optimize.Execution +import Shikumi.Optimize.Feedback (candidateFailurePolicy) +import Shikumi.Optimize.Report +import Shikumi.Program (embed, runProgram) +import Shikumi.Trace.Observation (NodeObservation, runProgramObserved) +import StubLM (Label (..), Sentence (..), runGepaStubLM, sentimentProg) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +run :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +run act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runGepaStubLM act + +cfg :: Int -> Int -> RunConfig +cfg cap width = defaultRunConfig {runLimits = RunLimits cap 8 width 1} + +tests :: TestTree +tests = + testGroup + "Execution" + [ testCase "invalid controls reject without dispatch" $ do + r <- run (runSearchSession (cfg (-1) 1) (const (pure ()))) + assertBool "invalid" (isLeft r), + testCase "caught budget error cannot admit cap+1" $ do + r <- run $ runSearchSession (cfg 2 1) $ \_ -> + replicateM_ 8 $ + (runProgram sentimentProg (Sentence "good") >> pure ()) `catchError` \_ (_ :: ShikumiError) -> pure () + check r $ \(_, report) -> do + admittedOperations report @?= 2 + runStatus report @?= BudgetStopped + last (map eventKind (events report)) @?= RunFinished BudgetStopped + eitherDecode (encode report) @?= Right report + assertBool "future report version rejected" (isLeft (eitherDecode (encode report {reportVersion = 2}) :: Either String OptimizationReport)), + testCase "opaque optimizer and zero cap" $ do + r <- run $ optimizeWith (cfg 0 1) (fromLegacyOptimizer (Optimizer $ \_ _ p -> runProgram p (Sentence "good") >> pure (freezeProgram p))) ds exactMatch sentimentProg + check r $ \(_, report) -> do + admittedOperations report @?= 0 + resultStatus report @?= Just Unscored, + testCase "opaque optimizer interrupted mid-run returns explicitly unscored baseline" $ do + let legacy = Optimizer $ \_ _ p -> do + _ <- runProgram p (Sentence "good") + _ <- runProgram p (Sentence "bad") + pure (freezeProgram p) + r <- run $ optimizeWith (cfg 1 1) (fromLegacyOptimizer legacy) ds exactMatch sentimentProg + check r $ \(_, report) -> do + admittedOperations report @?= 1 + candidateDetailAvailable report @?= False + resultStatus report @?= Just Unscored, + testCase "mid candidate stop is incomplete" $ do + r <- run $ runSearchSession (cfg 1 1) $ \s -> + evaluateFresh + s + ds + (runProgramObserved sentimentProg) + (candidateFailurePolicy scoreZero) + exactMatch + qualityPolicy + (scalarObjectives exactMatch) + check r $ \(_, report) -> do + map candidateStatus (candidates report) @?= [CandidateIncomplete] + map completedExamples (candidates report) @?= [1] + map candidateOperations (candidates report) @?= [1] + admittedOperations report @?= 1, + testCase "concurrent final-slot race obeys ceiling and ordered reports" $ do + r <- run $ runSearchSession (cfg 3 4) $ \s -> do + ids <- forM [0 .. 3 :: Int] (const (reserveCandidate s)) + evaluateCandidates + s + ( \ix -> + evaluateCandidate + s + ix + ds + (runProgramObserved sentimentProg) + (candidateFailurePolicy scoreZero) + exactMatch + qualityPolicy + (scalarObjectives exactMatch) + ) + [ix | Just ix <- ids] + check r $ \(_, report) -> do + admittedOperations report @?= 3 + sum (map candidateOperations (candidates report)) @?= 3 + map candidateId (candidates report) @?= [0, 1, 2, 3] + length [() | OptimizationEvent _ (CandidateEnded _ _) <- events report] @?= 4, + testCase "observer exception isolated" $ do + let observer = (cfg 8 1) {eventSink = \_ -> E.throwIO (userError "observer unavailable")} + r <- run $ runSearchSession observer (const (pure ())) + check r $ \(_, report) -> do + runStatus report @?= Completed + observerFailures report @?= 2, + testCase "barriers prove dispatch width and reversed completion" $ do + r <- run $ do + gate <- newEmptyMVar + active <- Ref.newIORef (0 :: Int, 0 :: Int, 0 :: Int) + let provider op = + E.bracket + (Ref.atomicModifyIORef' active (\(n, high, total) -> ((n + 1, max high (n + 1), total + 1), total))) + (\_ -> Ref.atomicModifyIORef' active (\(n, high, total) -> ((n - 1, high, total), ()))) + ( \ordinal -> do + if even ordinal then takeMVar gate else putMVar gate () + op + ) + result <- interpose + ( \_ -> \case + Complete m c o -> provider (complete m c o) + Stream m c o -> stream m c o + ) + $ runSearchSession (cfg 8 2) + $ \s -> do + ids <- forM [0 .. 3 :: Int] (const (reserveCandidate s)) + evaluateCandidates + s + ( \ix -> + evaluateCandidate + s + ix + (dataset [example (Sentence "good") (Label "positive")]) + (runProgramObserved sentimentProg) + (candidateFailurePolicy scoreZero) + exactMatch + qualityPolicy + (scalarObjectives exactMatch) + ) + [ix | Just ix <- ids] + counts <- Ref.readIORef active + pure (result, counts) + check r $ \((_, report), (active, high, total)) -> do + active @?= 0 + high @?= 2 + total @?= 4 + map candidateId (candidates report) @?= [0, 1, 2, 3] + map candidateStatus (candidates report) @?= replicate 4 CandidateCompleted, + testCase "cancellation closes candidate and run and propagates" $ do + r <- run $ do + gate <- newEmptyMVar + blocked <- newEmptyMVar + observed <- Ref.newIORef [] + let controls = (cfg 8 1) {eventSink = \event -> Ref.atomicModifyIORef' observed (\xs -> (xs ++ [event], ()))} + result <- withAsync + ( runSearchSession controls $ \s -> + evaluateFresh + s + ds + (\inp -> putMVar gate () >> takeMVar blocked >> runProgramObserved sentimentProg inp) + (candidateFailurePolicy scoreZero) + exactMatch + qualityPolicy + (scalarObjectives exactMatch) + ) + $ \worker -> do + takeMVar gate + cancel worker + waitCatch worker + evs <- Ref.readIORef observed + pure (isLeft result, evs) + check r $ \(cancelled, evs) -> do + assertBool "cancellation propagated" cancelled + length [() | OptimizationEvent _ (CandidateEnded 0 CandidateIncomplete) <- evs] @?= 1 + length [() | OptimizationEvent _ (RunFinished Cancelled) <- evs] @?= 1, + testCase "stream and completion share admission" $ do + r <- run $ runSearchSession (cfg 1 1) $ \_ -> + interpose + ( \_ -> \case + Complete m c o -> stream m c o >> complete m c o + Stream m c o -> stream m c o + ) + (runProgram sentimentProg (Sentence "good")) + check r $ \(_, report) -> do + admittedOperations report @?= 1 + runStatus report @?= BudgetStopped, + testCase "failed admitted operations are never refunded" $ do + r <- run + $ interpose + ( \_ -> \case + Complete {} -> throwError (ProviderFailure "fixture transport failed") + Stream m c o -> stream m c o + ) + $ runSearchSession (cfg 2 1) + $ \_ -> + replicateM_ 5 $ + (runProgram sentimentProg (Sentence "good") >> pure ()) `catchError` \_ (_ :: ShikumiError) -> pure () + check r $ \(_, report) -> admittedOperations report @?= 2, + testCase "caught stop cannot turn a candidate into completed success" $ do + r <- run $ runSearchSession (cfg 1 1) $ \s -> + evaluateFresh + s + ds + ( \inp -> do + replicateM_ 3 $ (runProgram sentimentProg inp >> pure ()) `catchError` \_ (_ :: ShikumiError) -> pure () + pure (Right (Label "positive"), []) + ) + (candidateFailurePolicy scoreZero) + exactMatch + qualityPolicy + (scalarObjectives exactMatch) + check r $ \(_, report) -> do + map candidateStatus (candidates report) @?= [CandidateIncomplete] + map completedExamples (candidates report) @?= [0], + testCase "reserved IDs execute once and unused reservations are reported" $ do + r <- run $ runSearchSession (cfg 8 1) $ \s -> do + first <- reserveCandidate s + _ <- reserveCandidate s + case first of + Nothing -> throwError (ValidationFailure "missing test reservation") + Just ident -> do + let evaluate = + evaluateCandidate + s + ident + ds + (runProgramObserved sentimentProg) + (candidateFailurePolicy scoreZero) + exactMatch + qualityPolicy + (scalarObjectives exactMatch) + _ <- evaluate + evaluate + check r $ \(out, report) -> do + assertBool "reuse rejected" (isLeft out) + length (candidates report) @?= 1 + unexecutedReservations report @?= [1], + testCase "Embed bodies share the operation ceiling" $ do + r <- run $ runSearchSession (cfg 1 1) $ \_ -> + runProgram + (embed (\inp -> runProgram sentimentProg inp >> runProgram sentimentProg inp)) + (Sentence "good") + check r $ \(_, report) -> do + admittedOperations report @?= 1 + runStatus report @?= BudgetStopped, + testCase "nested concurrent runner counts only admitted operations" $ do + r <- run $ runSearchSession (cfg 2 2) $ \s -> + evaluateFresh + s + (dataset [example (Sentence "good") (Label "positive")]) + ( \inp -> do + rows <- mapConcurrently (const (runProgramObserved sentimentProg inp)) [1 .. 4 :: Int] + case rows of + row : _ -> pure row + [] -> throwError (ValidationFailure "missing nested test rows") + ) + (candidateFailurePolicy scoreZero) + exactMatch + qualityPolicy + (scalarObjectives exactMatch) + check r $ \(_, report) -> do + admittedOperations report @?= 2 + map candidateOperations (candidates report) @?= [2] + map candidateStatus (candidates report) @?= [CandidateIncomplete], + testCase "caller BudgetExceeded remains failure" $ do + r <- run $ runSearchSession (cfg 8 1) (\_ -> throwError (BudgetExceeded "caller") :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] ()) + check r $ \(out, report) -> do + out @?= Left (BudgetExceeded "caller") + runStatus report @?= Failed + ] + where + ds = dataset [example (Sentence "good") (Label "positive"), example (Sentence "bad") (Label "negative")] + check (Left e) _ = assertFailure (show e) + check (Right x) f = f x + +evaluateFresh :: + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + SearchSession es -> + Dataset i o -> + (i -> Eff es (Either ShikumiError o, [NodeObservation])) -> + (ShikumiError -> FailurePolicy) -> + Metric o -> + ObjectivePolicy -> + ObjectiveMetric es o -> + Eff es CandidateReport +evaluateFresh s ds runner classifier metric policy objective = do + ident <- reserveCandidate s + case ident of + Nothing -> throwError (ValidationFailure "test exhausted candidate reservations") + Just ix -> evaluateCandidate s ix ds runner classifier metric policy objective diff --git a/test/FeedbackSpec.hs b/test/FeedbackSpec.hs new file mode 100644 --- /dev/null +++ b/test/FeedbackSpec.hs @@ -0,0 +1,237 @@ +module FeedbackSpec (tests) where + +import Baikai (AssistantContent (..), emptyResponse, emptyTextContent) +import Baikai.Error (contentFiltered) +import Control.Exception qualified as E +import Control.Lens ((&), (.~)) +import Data.IORef (atomicModifyIORef', newIORef) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Vector qualified as V +import Effectful (Eff, IOE, liftIO, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Dispatch.Dynamic (interpret) +import Effectful.Error.Static (Error, runErrorNoCallStack, throwError) +import Effectful.Prim (Prim, runPrim) +import Shikumi.Compile.Types (compiledProgram) +import Shikumi.Effect.Time (Time, runTime) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.Eval (Dataset, boolScore, dataset, exactMatch, example, scoreZero, unScore) +import Shikumi.Eval.Report (FailurePolicy (..), FailureReason (..)) +import Shikumi.LLM (LLM (..)) +import Shikumi.Module (predict) +import Shikumi.Optimize (optimize) +import Shikumi.Optimize.Feedback +import Shikumi.Optimize.GEPA (FeedbackCallback (..), ReflectIn (..), ReflectOut (..), captureFeedback, gepaWithFeedback, mutateFromEvidence) +import Shikumi.Optimize.Types (Budget (..)) +import Shikumi.Program (Params (..), Program (..), programParams) +import Shikumi.Signature (mkSignature) +import Shikumi.Trace.Feedback (feedbackFor) +import Shikumi.Trace.Node (NodePath (..), NodeStep (..), programNodePaths) +import Shikumi.Trace.Observation (NodeObservation (..)) +import StubLM (Label (..), Sentence (..), sentimentProg) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +runScript :: [Either ShikumiError Text] -> Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runScript replies action = do + ref <- newIORef replies + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + interpret + ( \_ -> \case + Complete {} -> do + next <- liftIO $ atomicModifyIORef' ref $ \case + [] -> ([], Right "{\"sentiment\":\"positive\"}") + x : xs -> (xs, x) + body <- either throwError pure next + if body == "CANCEL" + then liftIO (E.throwIO E.ThreadKilled) + else + pure (emptyResponse & #message . #content .~ V.singleton (AssistantText (emptyTextContent & #text .~ body))) + Stream {} -> pure [] + ) + action + +ok :: Either ShikumiError Text +ok = Right "{\"sentiment\":\"positive\"}" + +batch :: Dataset Sentence Label +batch = dataset [example (Sentence (T.pack (show i))) (Label "positive") | i <- [0 .. 2 :: Int]] + +metric :: (Applicative m) => Label -> EvaluationEvidence Label -> m FeedbackResult +metric _ _ = pure (FeedbackResult (boolScore True) Nothing []) + +checkRight :: (Show e) => Either e a -> (a -> IO ()) -> IO () +checkRight value check = either (assertFailure . show) check value + +tests :: TestTree +tests = + testGroup + "failure-aware feedback" + [ testCase "middle decode failure retains positions, original error, and missing codecs" $ do + result <- runScript [ok, Right "not JSON", ok] (captureEvidence defaultFeedbackConfig batch metric sentimentProg) + checkRight result $ \rows -> do + map (exampleIndex . first) rows @?= [0, 1, 2] + map (unScore . overallScore . second) rows @?= [1, 0, 1] + let ev = first (rows !! 1) + case executionResult ev of Left (MissingField "sentiment") -> pure (); other -> assertFailure (show other) + length (observations ev) @?= 1 + map observationInput (observations ev) @?= [Nothing] + map observationPath (observations ev) @?= programNodePaths sentimentProg + assertBool "failed leaf retained" (all (not . observationEligible) (observations ev)), + testCase "abort returns exact output error; budget and infrastructure always escape by default" $ do + let errors = [InvalidJSON "original", BudgetExceeded "budget", ProviderFailure "offline", ProviderError (contentFiltered "refused"), Timeout "infra"] + mapM_ + ( \err -> do + let cfg = if err == InvalidJSON "original" then defaultFeedbackConfig {failureClassification = const FailAbort} else defaultFeedbackConfig + result <- runScript [ok, Left err, ok] (captureEvidence cfg batch metric sentimentProg) + result @?= Left err + ) + errors + forced <- runScript [Left (BudgetExceeded "hard")] (captureEvidence (defaultFeedbackConfig {failureClassification = const (FailScore scoreZero)}) batch metric sentimentProg) + forced @?= Left (BudgetExceeded "hard"), + testCase "cancellation is never a completed batch" $ do + result <- E.try @E.AsyncException (runScript [ok, Right "CANCEL", ok] (captureEvidence defaultFeedbackConfig batch metric sentimentProg)) + result @?= Left E.ThreadKilled, + testCase "metric failure is separately labeled" $ do + result <- runScript [ok] (captureEvidence defaultFeedbackConfig batch (\_ _ -> throwError (ValidationFailure "critic failed")) sentimentProg) + checkRight result $ \rows -> map third rows @?= replicate 3 (Just (MetricError "ValidationFailure \"critic failed\"")), + testCase "custom infrastructure classification and failure score" $ do + let cfg = defaultFeedbackConfig {failureClassification = const (FailScore (boolScore True))} + result <- runScript [Left (ProviderFailure "explicit")] (captureEvidence cfg batch metric sentimentProg) + checkRight result $ \rows -> unScore (overallScore (second (rows !! 0))) @?= 1, + testCase "Map attribution rejects an unexecuted ordinal and wrong example/path" $ do + let prog = Map 1 sentimentProg + ds = dataset [example [Sentence "one", Sentence "two"] [Label "positive", Label "positive"]] + result <- runScript [ok, ok] (captureEvidence defaultFeedbackConfig ds (\_ _ -> pure (FeedbackResult scoreZero Nothing [])) prog) + checkRight result $ \rows -> do + let ev = first (rows !! 0) + target = NodeFeedback 0 (NodePath [StepMap]) (Just 1) "fix" Caller + validate f = validateFeedback defaultFeedbackConfig (programNodePaths prog) ev (FeedbackResult scoreZero Nothing [f]) + assertBool "real invocation accepted" (either (const False) (const True) (validate target)) + mapM_ + (\f -> assertBool "invalid target rejected" (either (const True) (const False) (validate f))) + [target {feedbackInvocation = Just 2}, target {feedbackExample = 1}, target {feedbackPath = NodePath []}, target {provenance = LegacyProgram}], + testCase "zero and odd critique bounds, and negative preflight" $ do + boundText 0 "abcdef" @?= "" + boundText 3 "abcdef" @?= "ab…" + boundText 1 "abcdef" @?= "…" + result <- runScript [Right "CANCEL"] (captureEvidence (defaultFeedbackConfig {critiqueCharacters = -1}) batch metric sentimentProg) + case result of Left ValidationFailure {} -> pure (); _ -> assertFailure "expected preflight rejection", + testCase "second node alone is reflected on with local input and redaction" $ do + let prog = Compose sentimentProg (predict (mkSignature "classify") :: Program Label Label) + ds = dataset [example (Sentence "outer-secret") (Label "positive")] + target _ ev = pure (FeedbackResult scoreZero Nothing [NodeFeedback (exampleIndex ev) (NodePath [StepComposeR]) Nothing "fix secret" Caller]) + cfg = defaultFeedbackConfig {redactEvidence = T.replace "secret" "REDACTED"} + proposer = Embed $ \input -> + if "positive" `T.isInfixOf` feedback input + && "REDACTED" `T.isInfixOf` feedback input + && not ("secret" `T.isInfixOf` feedback input) + && not ("outer-" `T.isInfixOf` feedback input) + then pure (ReflectOut "correct classifier") + else throwError (ValidationFailure (feedback input)) + result <- runScript [ok, ok] $ do + rows <- captureEvidence cfg ds target prog + changed <- mutateFromEvidence cfg proposer "program" "dataset" rows 0 prog + pure (programParams changed) + checkRight result $ \params -> do + take 1 params @?= take 1 (programParams prog) + instructionOverride (params !! 1) @?= Just "correct classifier", + testCase "GEPA callback improves only classifier in the complete search" $ do + let prog = Compose sentimentProg (predict (mkSignature "classify") :: Program Label Label) + ds = dataset [example (Sentence "outer") (Label "positive")] + target = FeedbackCallback $ \_ ev -> + pure + ( FeedbackResult + scoreZero + Nothing + [NodeFeedback (exampleIndex ev) (NodePath [StepComposeR]) Nothing "correct classifier" Caller] + ) + proposer = Embed $ \input -> + if "positive" `T.isInfixOf` feedback input && not ("outer" `T.isInfixOf` feedback input) + then pure (ReflectOut "correct classifier") + else throwError (ValidationFailure "wrong evidence") + wrong = Right "{\"sentiment\":\"negative\"}" + result <- + runScript [ok, wrong, ok, wrong, ok, ok] $ + optimize (gepaWithFeedback defaultFeedbackConfig proposer target (Budget 20 2)) ds exactMatch prog + checkRight result $ \cp -> do + let params = programParams (compiledProgram cp) + take 1 params @?= take 1 (programParams prog) + instructionOverride (params !! 1) @?= Just "correct classifier", + testCase "empty feedback and zero reflection limits never invoke proposer" $ do + let ds = dataset [example (Sentence "input") (Label "positive")] + proposer = Embed (const (throwError (ProviderFailure "must not reflect"))) + target _ ev = pure (FeedbackResult scoreZero Nothing [NodeFeedback (exampleIndex ev) (NodePath []) Nothing "fix" Caller]) + mapM_ + ( \cfg -> do + result <- runScript [ok] $ do + rows <- captureEvidence cfg ds target sentimentProg + programParams <$> mutateFromEvidence cfg proposer "" "" rows 0 sentimentProg + result @?= Right (programParams sentimentProg) + ) + [defaultFeedbackConfig {critiqueCharacters = 0}, defaultFeedbackConfig {reflectionExamples = 0}, defaultFeedbackConfig {reflectionCharacters = 0}], + testCase "reflection excludes sibling critiques" $ do + let prog = Compose sentimentProg (predict (mkSignature "classify") :: Program Label Label) + ds = dataset [example (Sentence "input") (Label "positive")] + target _ ev = + pure + ( FeedbackResult + scoreZero + Nothing + [ NodeFeedback (exampleIndex ev) (NodePath [StepComposeL]) Nothing "sibling-only" Caller, + NodeFeedback (exampleIndex ev) (NodePath [StepComposeR]) Nothing "local-only" Model + ] + ) + proposer = Embed $ \input -> + if "local-only" `T.isInfixOf` feedback input && not ("sibling-only" `T.isInfixOf` feedback input) + then pure (ReflectOut "changed") + else throwError (ValidationFailure "sibling contamination") + result <- runScript [ok, ok] $ do + rows <- captureEvidence defaultFeedbackConfig ds target prog + programParams <$> mutateFromEvidence defaultFeedbackConfig proposer "" "" rows 1 prog + checkRight result $ \params -> instructionOverride (params !! 1) @?= Just "changed", + testCase "odd reflection bound retains prioritized failure and marks truncation" $ do + let prog = Map 1 sentimentProg + ds = dataset [example [Sentence "good", Sentence "bad"] [Label "positive", Label "positive"]] + cfg = defaultFeedbackConfig {reflectionExamples = 1, reflectionCharacters = 101} + target _ ev = + pure + ( FeedbackResult + scoreZero + Nothing + [NodeFeedback (exampleIndex ev) (NodePath [StepMap]) Nothing "fix the failed invocation" Caller] + ) + proposer = Embed $ \input -> + if T.length (feedback input) <= 101 + && "MissingField" `T.isInfixOf` feedback input + && "…" `T.isSuffixOf` feedback input + && "invocation 1" `T.isInfixOf` feedback input + then pure (ReflectOut "changed") + else throwError (ValidationFailure (feedback input)) + result <- runScript [ok, Right "bad"] $ do + rows <- captureEvidence cfg ds target prog + programParams <$> mutateFromEvidence cfg proposer "" "" rows 0 prog + checkRight result $ \params -> instructionOverride (params !! 0) @?= Just "changed", + testCase "legacy multi-node critique is program-scoped and never broadcast" $ do + let prog = Compose sentimentProg (predict (mkSignature "classify") :: Program Label Label) + ds = dataset [example (Sentence "input") (Label "positive")] + result <- runScript [ok, ok] (captureFeedback ds (\_ _ -> (scoreZero, "fix")) prog) + checkRight result $ \(logbook, scores) -> do + scores @?= [0] + map (`feedbackFor` logbook) (programNodePaths prog) @?= [[], []] + feedbackFor (NodePath []) logbook @?= ["program (LegacyProgram): fix"], + testCase "retry evidence keeps rejected and accepted invocation lineage" $ do + let prog = Retry 2 sentimentProg + ds = dataset [example (Sentence "input") (Label "positive")] + result <- runScript [Right "bad JSON", ok] (captureEvidence defaultFeedbackConfig ds metric prog) + checkRight result $ \rows -> do + let obs = observations (first (rows !! 0)) + map observationInvocation obs @?= [0, 1] + assertBool "first rejected" (not (null (observationRejectedBy (obs !! 0)))) + observationEligible (obs !! 1) @?= True + ] + where + first (a, _, _) = a + second (_, b, _) = b + third (_, _, c) = c diff --git a/test/GepaSpec.hs b/test/GepaSpec.hs --- a/test/GepaSpec.hs +++ b/test/GepaSpec.hs @@ -5,6 +5,9 @@ -- mutation, and the end-to-end held-out lift + serialization round-trip. module GepaSpec (tests) where +import Data.IORef (newIORef, readIORef) +import Data.Map.Strict qualified as Map +import Data.Text qualified as T import Effectful (Eff, IOE, runEff) import Effectful.Concurrent (Concurrent, runConcurrent) import Effectful.Error.Static (Error, runErrorNoCallStack) @@ -18,23 +21,31 @@ import Shikumi.Optimize ( Budget (..), Candidate (..), + FeedbackCallback (..), FeedbackMetric, + GEPAConfig (..), captureFeedback, defaultBudget, + defaultGEPAConfig, dominates, gepa, + gepaWith, mutateNode, optimize, + optimizeWith, paretoFrontier, reflectiveProposer, sampleParent, scoreOn, withLmCallCount, ) +import Shikumi.Optimize.Execution qualified as X +import Shikumi.Optimize.Feedback qualified as F +import Shikumi.Optimize.Report qualified as R import Shikumi.Program (Params (..), foldParams, nodeFieldsIndexed, programParams) import Shikumi.Trace.Feedback (feedbackFor) import Shikumi.Trace.Node (programNodePaths) -import StubLM (Label (..), Sentence (..), ruleInstruction, runGepaStubLM, sentimentProg) +import StubLM (Label (..), Sentence (..), ruleInstruction, runGepaStubLM, runGepaStubLMCapturing, sentimentProg) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) @@ -67,7 +78,7 @@ runGepa act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runGepaStubLM act tests :: TestTree -tests = testGroup "Gepa" [paretoPure, feedbackCapture, reflectiveMutation, heldoutLift, budgetGate, roundTrips] +tests = testGroup "Gepa" [paretoPure, feedbackCapture, reflectiveMutation, heldoutLift, budgetGate, roundTrips, validatedSplit, validatedBudgetStop, emptyValidation] -- --------------------------------------------------------------------------- -- Pure Pareto-frontier tests @@ -107,7 +118,7 @@ (p0 : _) -> do let crits = feedbackFor p0 fblog assertBool "node 0 accumulated a critique" (not (null crits)) - assertBool "the critique is the expected one" (any (== "be more specific") crits) + assertBool "the critique is the expected one" (any (== "program (LegacyProgram): be more specific") crits) [] -> assertFailure "expected at least one node" -- --------------------------------------------------------------------------- @@ -176,3 +187,67 @@ Right cp -> case decodeCompiledOnto sentimentProg (encodeCompiled cp) of Left err -> assertFailure ("decode failed: " <> err) Right cp' -> programParams (compiledProgram cp') @?= programParams (compiledProgram cp) + +validatedSplit :: TestTree +validatedSplit = testCase "validation reverses training ranking without leaking its sentinel" $ do + ref <- newIORef [] + let train = dataset [example (Sentence "good TRAIN") (Label "neutral")] + validation = dataset [example (Sentence "good VALIDATION_SENTINEL") (Label "positive")] + callback = FeedbackCallback (F.legacyFeedback (\_ _ -> (boolScore True, "be more specific"))) + cfg = + (defaultGEPAConfig callback) + { validationDataset = Just validation, + feedbackConfig = F.defaultFeedbackConfig {F.includeProgramCritique = True}, + minibatchSize = 1 + } + controls = X.defaultRunConfig {X.runLimits = X.RunLimits 30 2 2 1} + result <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runGepaStubLMCapturing ref (optimizeWith controls (gepaWith cfg reflectiveProposer) train exactMatch sentimentProg) + case result of + Left e -> assertFailure (show e) + Right (_, report) -> do + R.selectedCandidate report @?= Just 1 + map R.objectiveValues (R.candidates report) @?= map (Map.singleton "quality") [0, 1] + requests <- readIORef ref + let reflections = filter (T.isInfixOf "## proposedInstruction ##") requests + assertBool "reflection occurred" (not (null reflections)) + assertBool "validation absent from reflection" (all (not . T.isInfixOf "VALIDATION_SENTINEL") reflections) + +validatedBudgetStop :: TestTree +validatedBudgetStop = testCase "budget stop during later proposal retains completed validation winner" $ do + let callback = FeedbackCallback (F.legacyFeedback (\_ _ -> (boolScore True, "be more specific"))) + cfg = + (defaultGEPAConfig callback) + { validationDataset = Just (dataset [example (Sentence "good validation") (Label "positive")]), + feedbackConfig = F.defaultFeedbackConfig {F.includeProgramCritique = True}, + minibatchSize = 1 + } + controls = X.defaultRunConfig {X.runLimits = X.RunLimits 6 3 1 1} + result <- + runGepa $ + optimizeWith + controls + (gepaWith cfg reflectiveProposer) + (dataset [example (Sentence "good training") (Label "neutral")]) + exactMatch + sentimentProg + case result of + Left e -> assertFailure (show e) + Right (compiled, report) -> do + R.selectedCandidate report @?= Just 1 + R.runStatus report @?= R.BudgetStopped + R.admittedOperations report @?= 6 + map R.candidateStatus (R.candidates report) @?= [R.CandidateCompleted, R.CandidateCompleted, R.CandidateIncomplete] + case foldParams (compiledProgram compiled) of + p : _ -> instructionOverride p @?= Just ruleInstruction + [] -> assertFailure "missing predictor" + +emptyValidation :: TestTree +emptyValidation = testCase "explicit empty validation fails even at zero budget" $ do + let cfg = (defaultGEPAConfig (FeedbackCallback (F.legacyFeedback fbMetric))) {validationDataset = Just (dataset [])} + controls = X.defaultRunConfig {X.runLimits = X.RunLimits 0 0 1 1} + result <- runGepa (optimizeWith controls (gepaWith cfg reflectiveProposer) trainset exactMatch sentimentProg) + case result of + Left _ -> pure () + Right _ -> assertFailure "empty validation was silently accepted" diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -9,16 +9,20 @@ import BootstrapSpec qualified import CoproSpec qualified import EnsembleSpec qualified +import ExecutionSpec qualified +import FeedbackSpec qualified import GepaSpec qualified import InstructionSpec qualified import KNNSpec qualified import LabeledFewShotSpec qualified import Miprov2Spec qualified +import ObjectiveSpec qualified import OptimizeSpec qualified import ProposeSpec qualified import RandomSearchSpec qualified import SearchSpec qualified import SeedingSpec qualified +import StructureSpec qualified import Test.Tasty (defaultMain, testGroup) main :: IO () @@ -26,7 +30,8 @@ defaultMain $ testGroup "shikumi-optimize" - [ OptimizeSpec.tests, + [ StructureSpec.tests, + OptimizeSpec.tests, LabeledFewShotSpec.tests, BootstrapSpec.tests, InstructionSpec.tests, @@ -34,7 +39,10 @@ ProposeSpec.tests, Miprov2Spec.tests, CoproSpec.tests, + ExecutionSpec.tests, + ObjectiveSpec.tests, GepaSpec.tests, + FeedbackSpec.tests, KNNSpec.tests, RandomSearchSpec.tests, SearchSpec.tests, diff --git a/test/NodeBootstrapSpec.hs b/test/NodeBootstrapSpec.hs new file mode 100644 --- /dev/null +++ b/test/NodeBootstrapSpec.hs @@ -0,0 +1,245 @@ +module NodeBootstrapSpec (tests) where + +import Baikai (AssistantContent (..), emptyResponse, emptyTextContent) +import Control.Lens ((&), (.~), (^.)) +import Data.Aeson (ToJSON, object, toJSON, (.=)) +import Data.Generics.Labels () +import Data.IORef (IORef, modifyIORef', newIORef, readIORef) +import Data.Map.Strict qualified as Map +import Data.Maybe (fromMaybe) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Vector qualified as V +import Effectful (Eff, IOE, liftIO, runEff, (:>)) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Dispatch.Dynamic (interpret) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +import GHC.Generics (Generic) +import Shikumi.Adapter (ToPrompt) +import Shikumi.Compile.Serialize (decodeCompiledOnto, encodeCompiled) +import Shikumi.Compile.Types (compiledProgram) +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, predictCaptured) +import Shikumi.Optimize +import Shikumi.Program (CaptureCodec (..), Demo (..), Params (..), Program (..), emptyParams, mapParamsAt, programParams, runProgram) +import Shikumi.Schema (FromModel, ToSchema, Validatable, deriveSchema) +import Shikumi.Signature (Signature, mkSignature) +import Shikumi.Trace.Node (NodePath (..), NodeStep (..)) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +newtype Question = Question {question :: Text} deriving stock (Eq, Show, Generic) + +instance ToJSON Question + +instance ToSchema Question + +instance FromModel Question + +instance ToPrompt Question + +instance Validatable Question + +newtype City = City {city :: Text} deriving stock (Eq, Show, Generic) + +instance ToJSON City + +instance ToSchema City + +instance FromModel City + +instance ToPrompt City + +instance Validatable City + +newtype Country = Country {country :: Text} deriving stock (Eq, Show, Generic) + +instance ToJSON Country + +instance ToSchema Country + +instance FromModel Country + +instance ToPrompt Country + +instance Validatable Country + +citySig :: Signature Question City +citySig = mkSignature "Extract city" + +countrySig :: Signature City Country +countrySig = mkSignature "Resolve country" + +program :: Program Question Country +program = Compose (predictCaptured citySig) (predictCaptured countrySig) + +train :: Dataset Question Country +train = dataset [example (Question "Where is Paris?") (Country "France")] + +expected :: [[Demo]] +expected = [[Demo (toJSON (Question "Where is Paris?")) (toJSON (City "Paris"))], [Demo (toJSON (City "Paris")) (toJSON (Country "France"))]] + +runCity :: IORef Int -> Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runCity calls = + runEff + . runPrim + . runTime + . runConcurrent + . runErrorNoCallStack @ShikumiError + . interpret + ( \_ -> \case + Complete _ ctx _ -> do + liftIO (modifyIORef' calls (+ 1)) + callNumber <- liftIO (readIORef calls) + let prompt = fromMaybe "" (ctx ^. #systemPrompt) + canAnswer = not ("Student" `T.isInfixOf` prompt) || V.length (ctx ^. #messages) > 1 + body + | "Retry country" `T.isInfixOf` prompt && callNumber == 2 = "{\"country\":\"Spain\"}" + | "Extract city" `T.isInfixOf` prompt = if canAnswer then "{\"city\":\"Paris\"}" else "{\"city\":\"Lyon\"}" + | otherwise = if canAnswer then "{\"country\":\"France\"}" else "{\"country\":\"Spain\"}" + pure (emptyResponse & #message . #content .~ V.singleton (AssistantText (emptyTextContent & #text .~ body))) + Stream {} -> pure [] + ) + +recover :: (LLM :> es, Error ShikumiError :> es, Prim :> es) => NodeBootstrapConfig -> Program Question Country -> Program i o -> Eff es (Map.Map NodePath [Demo], BootstrapReport) +recover cfg teacher student = do + meter <- newBudgetMeter defaultBudget + bootstrapNodeDemos cfg meter teacher train exactMatch student + +tests :: TestTree +tests = + testGroup + "node-local bootstrap" + [ testCase "bootstrap excludes successful leaves from a rejected retry attempt" $ do + calls <- newIORef 0 + let retryTeacher = + Retry + 2 + ( Validate + (\out -> if out == Country "France" then Right out else Left "retry") + (mapParamsAt 1 (\ps -> ps {instructionOverride = Just "Retry country"}) program) + ) + cfg = + defaultNodeBootstrapConfig + { nodeMapping = + Just + [ (NodePath [StepRetry, StepValidate, StepComposeL], NodePath [StepComposeL]), + (NodePath [StepRetry, StepValidate, StepComposeR], NodePath [StepComposeR]) + ] + } + result <- runCity calls (recover cfg retryTeacher program) + case result of Right (pools, _) -> Map.elems pools @?= expected; Left e -> assertFailure (show e) + readIORef calls >>= (@?= 4), + testCase "explicit mapping handles a differently wrapped teacher" $ do + calls <- newIORef 0 + let wrapped = Validate Right program + pairs = [(NodePath [StepValidate, StepComposeL], NodePath [StepComposeL]), (NodePath [StepValidate, StepComposeR], NodePath [StepComposeR])] + automatic <- runCity calls (recover defaultNodeBootstrapConfig wrapped program) + case automatic of Left (ValidationFailure _) -> pure (); _ -> assertFailure "expected explicit mapping requirement" + readIORef calls >>= (@?= 0) + mapped <- runCity calls (recover (defaultNodeBootstrapConfig {nodeMapping = Just pairs}) wrapped program) + case mapped of Right (pools, _) -> Map.elems pools @?= expected; Left e -> assertFailure (show e), + testCase "compatible predictors use independent seeded subsets" $ do + calls <- newIORef 0 + let parallel = Parallel (predictCaptured citySig) (predictCaptured citySig) + examples = dataset [example (Question (T.pack (show n))) (City "Paris", City "Paris") | n <- [1 .. 12 :: Int]] + cfg = defaultNodeBootstrapConfig {nodeSeed = Just 17, nodeBootstrapConfig = BootstrapConfig 1 3} + collect = do + meter <- newBudgetMeter defaultBudget + bootstrapNodeDemos cfg meter parallel examples exactMatch parallel + first <- runCity calls collect + again <- runCity calls collect + first @?= again + case first of + Right (pools, _) -> case Map.elems pools of + [a, b] -> do + map length [a, b] @?= [3, 3] + assertBool "independent subsets" (a /= b) + _ -> assertFailure "expected two nodes" + Left e -> assertFailure (show e), + testCase "duplicate targets require explicit merge" $ do + calls <- newIORef 0 + let parallel = Parallel (predictCaptured citySig) (predictCaptured citySig) + examples = dataset [example (Question "Paris?") (City "Paris", City "Paris")] + cfg = defaultNodeBootstrapConfig {nodeMapping = Just [(NodePath [StepParallelL], NodePath []), (NodePath [StepParallelR], NodePath [])]} + collect c = do + meter <- newBudgetMeter defaultBudget + bootstrapNodeDemos c meter parallel examples exactMatch (predictCaptured citySig) + rejected <- runCity calls (collect cfg) + case rejected of Left (ValidationFailure _) -> pure (); _ -> assertFailure "expected duplicate rejection" + readIORef calls >>= (@?= 0) + merged <- runCity calls (collect (cfg {mergeTargetMappings = True})) + case merged of Right (pools, _) -> map length (Map.elems pools) @?= [2]; Left e -> assertFailure (show e), + testCase "heterogeneous demos run and round-trip onto the capture template" $ do + calls <- newIORef 0 + result <- runCity calls $ do + cp <- optimize (bootstrapFewShot program defaultBudget) train exactMatch program + out <- runProgram (compiledProgram cp) (Question "Where is Paris?") + pure (cp, out) + case result of + Left e -> assertFailure (show e) + Right (cp, out) -> do + out @?= Country "France" + map demos (programParams (compiledProgram cp)) @?= expected + case decodeCompiledOnto program (encodeCompiled cp) of + Left e -> assertFailure e + Right restored -> do + map demos (programParams (compiledProgram restored)) @?= expected + replay <- runCity calls (runProgram (compiledProgram restored) (Question "Where is Paris?")) + replay @?= Right (Country "France"), + testCase "incompatible mappings fail before calls" $ do + calls <- newIORef 0 + let cfg = defaultNodeBootstrapConfig {nodeMapping = Just [(NodePath [StepComposeL], NodePath [StepComposeR])]} + result <- runCity calls (recover cfg program program) + case result of Left (ValidationFailure _) -> pure (); _ -> assertFailure "expected preflight rejection" + readIORef calls >>= (@?= 0), + testCase "uncaptured composite fails before calls" $ do + calls <- newIORef 0 + result <- runCity calls (recover defaultNodeBootstrapConfig (Compose (predict citySig) (predict countrySig)) program) + case result of Left (ValidationFailure msg) -> assertBool "actionable" ("predictCaptured" `T.isInfixOf` msg); _ -> assertFailure "expected missing codec" + readIORef calls >>= (@?= 0), + testCase "failed root contributes no demos" $ do + calls <- newIORef 0 + let failed = Validate (const (Left "rejected")) program + cfg = defaultNodeBootstrapConfig {nodeMapping = Just [(NodePath [StepValidate, StepComposeL], NodePath [StepComposeL]), (NodePath [StepValidate, StepComposeR], NodePath [StepComposeR])]} + result <- runCity calls (recover cfg failed program) + case result of Right (pools, _) -> Map.elems pools @?= [[], []]; Left e -> assertFailure (show e), + testCase "a dishonest codec is rejected by target decoding" $ do + calls <- newIORef 0 + let codec = CaptureCodec toJSON (const (object ["wrong" .= ("Paris" :: Text)])) (deriveSchema @Question) (deriveSchema @City) + bad = Compose (PredictCaptured codec citySig emptyParams) (predictCaptured countrySig) + result <- runCity calls (recover defaultNodeBootstrapConfig bad program) + case result of + Right (pools, BootstrapReport notes) -> do + Map.lookup (NodePath [StepComposeL]) pools @?= Just [] + assertBool "diagnostic" (not (null notes)) + Left e -> assertFailure (show e), + testCase "seeded pools reproduce and obey per-node caps" $ do + calls <- newIORef 0 + let cfg = defaultNodeBootstrapConfig {nodeSeed = Just 17, nodeBootstrapConfig = BootstrapConfig 1 1} + a <- runCity calls (recover cfg program program) + b <- runCity calls (recover cfg program program) + a @?= b + case a of Right (pools, _) -> map length (Map.elems pools) @?= [1, 1]; Left e -> assertFailure (show e), + testCase "random search accepts a heterogeneous capture program" $ do + calls <- newIORef 0 + let student = + mapParamsAt 0 (\ps -> ps {instructionOverride = Just "Student Extract city"}) $ + mapParamsAt 1 (\ps -> ps {instructionOverride = Just "Student Resolve country"}) program + result <- runCity calls (optimize (bootstrapRandomSearch program 2 defaultBudget) train exactMatch student) + case result of + Right cp -> do + map demos (programParams (compiledProgram cp)) @?= expected + runCity calls (runProgram (compiledProgram cp) (Question "Paris?")) >>= (@?= Right (Country "France")) + Left e -> assertFailure (show e), + testCase "MIPRO builds separate valid node candidate pools" $ do + calls <- newIORef 0 + result <- runCity calls (bootstrapDemoCandidates ((miprov2Auto Miprov2Light) {numDemoCandidates = 4}) program train exactMatch program) + case result of + Right sets -> map (filter (not . null)) sets @?= map (: []) expected + Left e -> assertFailure (show e) + ] diff --git a/test/ObjectiveSpec.hs b/test/ObjectiveSpec.hs new file mode 100644 --- /dev/null +++ b/test/ObjectiveSpec.hs @@ -0,0 +1,38 @@ +module ObjectiveSpec (tests) where + +import Data.Aeson (eitherDecode, encode) +import Data.Map.Strict qualified as Map +import Shikumi.Optimize.Report +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Named objectives" + [ testCase "quality/cost frontier and hard ceiling" $ do + map candidateId (objectiveFrontier policy rows) @?= [0, 1] + fmap candidateId (selectObjectiveWinner policy rows) @?= Just 0 + let bounded = policy {objectives = [quality, cost {upperBound = Just 2}]} + fmap candidateId (selectObjectiveWinner bounded rows) @?= Just 1, + testCase "directions and creation-order ties" $ do + fmap candidateId (selectObjectiveWinner (policy {primaryObjective = "cost"}) rows) @?= Just 1 + fmap candidateId (selectObjectiveWinner policy [a {candidateId = 9}, a]) @?= Just 0, + testCase "required values, nonfinite values, duplicates fail" $ do + assertBool "missing" (isLeft (aggregateObjectives policy [Map.singleton "quality" 1])) + assertBool "NaN" (isLeft (aggregateObjectives policy [Map.fromList [("quality", 0 / 0), ("cost", 1)]])) + assertBool "duplicates" (isLeft (validateObjectives (policy {objectives = [quality, quality]}))), + testCase "substitution and aggregation" $ do + let p = policy {objectives = [quality, cost {missingPolicy = Substitute 2, aggregation = Total}]} + aggregateObjectives p [Map.singleton "quality" 1, Map.fromList [("quality", 0), ("cost", 3)]] @?= Right (Map.fromList [("quality", 0.5), ("cost", 5)]), + testCase "configuration JSON" $ eitherDecode (encode policy) @?= Right policy + ] + where + quality = ObjectiveSpec "quality" "score" Maximize Mean Required (Just 0) (Just 1) + cost = ObjectiveSpec "cost" "operations" Minimize Mean Required (Just 0) Nothing + policy = ObjectivePolicy [quality, cost] "quality" [] + a = row 0 1 3 + rows = [a, row 1 0.9 1, row 2 0.8 4] + row ix q c = CandidateReport ix CandidateCompleted 1 1 [(0, q)] (Map.fromList [("quality", q), ("cost", c)]) 0 Nothing + isLeft (Left _) = True + isLeft _ = False diff --git a/test/StructureSpec.hs b/test/StructureSpec.hs new file mode 100644 --- /dev/null +++ b/test/StructureSpec.hs @@ -0,0 +1,239 @@ +module StructureSpec (tests) where + +import Baikai (AssistantContent (..), Context, Response, emptyResponse, emptyTextContent) +import Control.Lens ((&), (.~), (^.)) +import Control.Monad (replicateM_) +import Data.Aeson (eitherDecode, encode) +import Data.Either (isLeft) +import Data.Generics.Labels () +import Data.List.NonEmpty qualified as NE +import Data.Map.Strict qualified as Map +import Data.Text (Text) +import Data.Text qualified as T +import Data.Vector qualified as V +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar) +import Effectful.Dispatch.Dynamic (interpose, interpret) +import Effectful.Error.Static (Error, catchError, runErrorNoCallStack, throwError) +import Effectful.Prim (Prim, runPrim) +import Effectful.Prim.IORef qualified as Ref +import GHC.Generics (Generic) +import Shikumi.Adapter (ToPrompt) +import Shikumi.Compile +import Shikumi.Effect.Time (Time, runTime) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.Eval (Dataset, dataset, exactMatch, example, scoreZero) +import Shikumi.LLM (LLM (..), complete, stream) +import Shikumi.Module (predict) +import Shikumi.Optimize.Execution +import Shikumi.Optimize.Feedback (candidateFailurePolicy) +import Shikumi.Optimize.Report +import Shikumi.Optimize.Structure +import Shikumi.Program (Program (..), embed, runProgram) +import Shikumi.Schema (FromModel, ToSchema, Validatable) +import Shikumi.Signature (mkSignature) +import Test.Tasty +import Test.Tasty.HUnit + +newtype Input = Input {question :: Text} + deriving stock (Eq, Show, Generic) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +newtype Output = Output {answer :: Text} + deriving stock (Eq, Show, Generic) + deriving anyclass (ToSchema, FromModel, ToPrompt, Validatable) + +base :: Program Input Output +base = predict (mkSignature "Answer") + +response :: Bool -> Text -> Response +response cot directAnswer = emptyResponse & #message . #content .~ V.singleton (AssistantText (emptyTextContent & #text .~ body)) + where + body = + if cot + then "[[ ## reasoning ## ]]\nreason\n[[ ## value ## ]]\n{\"answer\":\"cot\"}\n[[ ## completed ## ]]" + else "[[ ## answer ## ]]\n" <> directAnswer <> "\n[[ ## completed ## ]]" + +respond :: Context -> Response +respond context = response (maybe False (T.isInfixOf "step by step") (context ^. #systemPrompt)) (if "held-out 1" `T.isInfixOf` T.pack (show context) then "cot" else "direct") + +-- The provider returns different typed answers for direct and reasoning requests. +run :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +run action = + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + interpret (\_ -> \case Complete _ c _ -> pure (respond c); Stream {} -> pure []) action + +right :: (Show e) => Either e a -> IO a +right = either (fail . show) pure + +cfg :: Int -> Int -> Int -> RunConfig +cfg cap candidatesCount width = defaultRunConfig {runLimits = RunLimits cap candidatesCount width 1} + +train :: Dataset Input Output +train = dataset [example (Input "training") (Output "direct")] + +validation :: Dataset Input Output +validation = dataset [example (Input "held-out 1") (Output "cot"), example (Input "held-out 2") (Output "cot")] + +search :: RunConfig -> StructureRegistry Input Output -> Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] (StructureSearchResult Input Output) +search controls = structureSearchWith controls train validation exactMatch (candidateFailurePolicy scoreZero) qualityPolicy (scalarObjectives exactMatch) + +tests :: TestTree +tests = + testGroup + "structure search" + [ testCase "validation selects CoT, report identity and artifact behavior agree" $ do + registry <- right (directCotRegistry "fixture" base) + result <- run (search (cfg 8 2 1) registry) >>= right + recipeIdText (selectedRecipeId result) @?= "cot" + let report = structureReport result + selectedCandidate report @?= Just 1 + map objectiveValues (candidates report) @?= [Map.singleton "quality" 0.5, Map.singleton "quality" 1] + admittedOperations report @?= 4 + Map.lookup 1 (candidateMetadata report) @?= Just (Map.fromList [("recipeId", "cot"), ("recipeRevision", "1"), ("registryId", "fixture")]) + eitherDecode (encode report) @?= Right report + bytes <- right (encodeStructureArtifact registry (selectedRecipeId result) (selectedStructure result)) + restored <- right (decodeStructureArtifact registry bytes) + (before, restoredOutput, requests) <- + run + ( do + ref <- Ref.newIORef [] + let capture = + interpose + ( \_ -> \case + Complete m c o -> Ref.atomicModifyIORef' ref (\xs -> (xs ++ [c], ())) >> complete m c o + Stream m c o -> stream m c o + ) + before <- capture (runCompiled (selectedStructure result) (Input "same")) + restoredOutput <- capture (runCompiled restored (Input "same")) + requests <- Ref.readIORef ref + pure (before, restoredOutput, requests) + ) + >>= right + before @?= Output "cot" + restoredOutput @?= before + length requests @?= 2 + case requests of + [originalRequest, restoredRequest] -> originalRequest @?= restoredRequest + _ -> assertFailure "expected two captured requests", + testCase "opposite validation ranking selects direct despite CoT training labels" $ do + registry <- right (directCotRegistry "fixture" base) + result <- run (structureSearchWith (cfg 8 2 1) validation train exactMatch (candidateFailurePolicy scoreZero) qualityPolicy (scalarObjectives exactMatch) registry) >>= right + recipeIdText (selectedRecipeId result) @?= "direct", + testCase "zero candidate or operation budgets return unscored baseline" $ do + registry <- right (directCotRegistry "fixture" base) + mapM_ + ( \controls -> do + result <- run (search controls registry) >>= right + recipeIdText (selectedRecipeId result) @?= "direct" + resultStatus (structureReport result) @?= Just Unscored + admittedOperations (structureReport result) @?= 0 + ) + [cfg 0 2 1, cfg 8 0 1], + testCase "partial CoT cannot displace completed direct" $ do + registry <- right (directCotRegistry "fixture" base) + result <- run (search (cfg 3 2 1) registry) >>= right + recipeIdText (selectedRecipeId result) @?= "direct" + map candidateStatus (candidates (structureReport result)) @?= [CandidateCompleted, CandidateIncomplete] + admittedOperations (structureReport result) @?= 3, + testCase "candidate ceiling only evaluates the first recipe" $ do + registry <- right (directCotRegistry "fixture" base) + result <- run (search (cfg 8 1 1) registry) >>= right + recipeIdText (selectedRecipeId result) @?= "direct" + length (candidates (structureReport result)) @?= 1 + runStatus (structureReport result) @?= BudgetStopped, + testCase "opaque caught denials and nested retries cannot overspend" $ do + let opaque = + embed + ( \inp -> do + replicateM_ 9 ((runProgram base inp >> pure ()) `catchError` \_ (_ :: ShikumiError) -> pure ()) + pure (Output "cot") + ) + recipe <- right (structureRecipe "opaque" 1 "" (Retry 3 opaque)) + registry <- right (structureRegistry "fixture" [recipe]) + (result, actual) <- + run + ( do + counter <- Ref.newIORef (0 :: Int) + result <- + interpose + ( \_ -> \case + Complete m c o -> Ref.atomicModifyIORef' counter (\n -> (n + 1, ())) >> complete m c o + Stream m c o -> Ref.atomicModifyIORef' counter (\n -> (n + 1, ())) >> stream m c o + ) + (search (cfg 2 1 1) registry) + actual <- Ref.readIORef counter + pure (result, actual) + ) + >>= right + actual @?= 2 + admittedOperations (structureReport result) @?= actual + map candidateStatus (candidates (structureReport result)) @?= [CandidateIncomplete], + testCase "queued candidates do not start after exact exhaustion" $ do + registry <- right (directCotRegistry "fixture" base) + result <- run (search (cfg 2 2 1) registry) >>= right + map candidateStatus (candidates (structureReport result)) @?= [CandidateCompleted] + unexecutedReservations (structureReport result) @?= [1] + selectedCandidate (structureReport result) @?= Just 0, + testCase "nested ensemble stream and completion share the cap" $ do + recipe <- right (structureRecipe "ensemble" 1 "" (Ensemble [base, base] (const (Output "cot")))) + registry <- right (structureRegistry "fixture" [recipe]) + result <- run (searchWithStreams registry) >>= right + admittedOperations (structureReport result) @?= 3 + map candidateStatus (candidates (structureReport result)) @?= [CandidateIncomplete], + testCase "ties retain registry order under reversed completion" $ do + registry <- right (directCotRegistry "fixture" base) + result <- + run + ( do + gate <- newEmptyMVar + let controls = + (cfg 8 2 2) + { eventSink = \event -> case eventKind event of + CandidateEnded 1 CandidateCompleted -> putMVar gate () + _ -> pure () + } + interpose + ( \_ -> \case + Complete m c o -> + if maybe False (T.isInfixOf "step by step") (c ^. #systemPrompt) + then complete m c o + else takeMVar gate >> complete m c o + Stream m c o -> stream m c o + ) + (structureSearchWith controls train train exactMatch (candidateFailurePolicy scoreZero) qualityPolicy (\_ _ -> pure (Map.singleton "quality" 1)) registry) + ) + >>= right + recipeIdText (selectedRecipeId result) @?= "direct" + selectedCandidate (structureReport result) @?= Just 0, + testCase "empty validation and non-finite objectives reject" $ do + registry <- right (directCotRegistry "fixture" base) + invalid <- run (structureSearchWith (cfg 0 0 1) train (dataset []) exactMatch (candidateFailurePolicy scoreZero) qualityPolicy (scalarObjectives exactMatch) registry) + assertBool "invalid dataset" (isLeft invalid) + result <- run (structureSearchWith (cfg 8 2 1) train validation exactMatch (candidateFailurePolicy scoreZero) qualityPolicy (\_ _ -> pure (Map.singleton "quality" (0 / 0))) registry) >>= right + resultStatus (structureReport result) @?= Just Unscored + map candidateStatus (candidates (structureReport result)) @?= [CandidateFailed, CandidateFailed], + testCase "infrastructure failures propagate" $ do + recipe <- right (structureRecipe "broken" 1 "" (embed (\_ -> throwError (ProviderFailure "offline failure")) :: Program Input Output)) + registry <- right (structureRegistry "fixture" [recipe]) + result <- run (search (cfg 8 1 1) registry) + assertBool "propagated" (isLeft result) + ] + +-- Each predictor dispatches a stream followed by completion inside the session. +searchWithStreams :: StructureRegistry Input Output -> Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] (StructureSearchResult Input Output) +searchWithStreams registry = do + let withStreams = + embed + ( \inp -> + interpose + ( \_ -> \case + Complete m c o -> stream m c o >> complete m c o + Stream m c o -> stream m c o + ) + (runProgram (recipeProgram (NE.head (registryRecipes registry))) inp) + ) + case structureRecipe "stream-ensemble" 1 "" withStreams >>= \recipe -> structureRegistry "streams" [recipe] of + Left err -> throwError (ValidationFailure (T.pack (show err))) + Right registered -> search (cfg 3 1 1) registered diff --git a/test/StubLM.hs b/test/StubLM.hs --- a/test/StubLM.hs +++ b/test/StubLM.hs @@ -88,7 +88,7 @@ import Shikumi.Adapter (ToPrompt) import Shikumi.Combinator ((>>>)) import Shikumi.LLM (LLM (..)) -import Shikumi.Module (predict) +import Shikumi.Module (predict, predictCaptured) import Shikumi.Program (Program) import Shikumi.Schema (FromModel, ToSchema, Validatable) import Shikumi.Signature (Signature, mkSignature) @@ -148,7 +148,7 @@ -- | A two-node sentiment pipeline: classify a sentence, then echo the label. sentimentPipeline :: Program Sentence Label -sentimentPipeline = sentimentProg >>> predict echoSig +sentimentPipeline = predictCaptured sentimentSig >>> predictCaptured echoSig -- --------------------------------------------------------------------------- -- Ground truth and helpers