diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,24 @@ ## Unreleased +## 0.3.1.0 — 2026-11-02 + +### Added + +- Add ReAnchor calibration for decision thresholds, score cuts and choice weights. + Candidate trials use run-local response replay above admission, retain only strict + improvements that pass held-out folds, and preserve absent settings on rejection. + Instructions and demos stay unchanged. Reports include retained settings, score + summaries, evidence ranges and fold counters; optional validation is report-only. +- Expose accepted-candidate selection for strategies with an additional selection + gate. Unaccepted trials remain diagnostic and cannot become the reported winner. + +### Changed + +- Upgrade `mori://shinzui/baikai/packages/baikai` to `>=0.10.0.0 && <0.11`. + +- Update internal dependency bounds for the breaking `shikumi 0.5.0.0` release. + ## 0.3.0.2 — 2026-10-16 ### 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.3.0.2 +version: 0.3.1.0 synopsis: The optimizer framework for shikumi LM programs (EP-10) category: AI description: @@ -64,6 +64,9 @@ Shikumi.Optimize.Propose.Tips Shikumi.Optimize.Propose.Types Shikumi.Optimize.RandomSearch + Shikumi.Optimize.ReAnchor + Shikumi.Optimize.ReAnchor.Fit + Shikumi.Optimize.ReAnchor.Replay Shikumi.Optimize.Report Shikumi.Optimize.Search Shikumi.Optimize.Structure @@ -71,15 +74,17 @@ build-depends: aeson >=2.2 && <2.3, + baikai >=0.10.0.0 && <0.11, base >=4.20 && <5, containers >=0.6 && <0.9, effectful >=2.6 && <2.8, generic-lens >=2.2 && <2.4, lens ^>=5.3, - shikumi ^>=0.4.1.1, - shikumi-compile ^>=0.2.1.2, - shikumi-eval ^>=0.3.0.2, - shikumi-trace ^>=0.3.0.2, + shikumi ^>=0.5.0.0, + shikumi-cache ^>=0.2.1.0, + shikumi-compile ^>=0.2.2.0, + shikumi-eval ^>=0.3.1.0, + shikumi-trace ^>=0.3.1.0, text ^>=2.1, vector >=0.13 && <0.14, @@ -109,6 +114,7 @@ OptimizeSpec ProposeSpec RandomSearchSpec + ReAnchorSpec SearchSpec SeedingSpec StructureSpec @@ -116,17 +122,18 @@ build-depends: aeson, - baikai >=0.8.0.0 && <0.9, + baikai >=0.10.0.0 && <0.11, base, containers, effectful, generic-lens, lens, - shikumi ^>=0.4.1.1, - shikumi-compile ^>=0.2.1.2, - shikumi-eval ^>=0.3.0.2, - shikumi-optimize ^>=0.3.0.2, - shikumi-trace ^>=0.3.0.2, + shikumi ^>=0.5.0.0, + shikumi-compile ^>=0.2.2.0, + shikumi-eval ^>=0.3.1.0, + shikumi-optimize ^>=0.3.1.0, + shikumi-testing, + shikumi-trace ^>=0.3.1.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 @@ -29,6 +29,7 @@ module Shikumi.Optimize.Pareto, module Shikumi.Optimize.KNN, module Shikumi.Optimize.RandomSearch, + module Shikumi.Optimize.ReAnchor, module Shikumi.Optimize.Ensemble, ) where @@ -53,6 +54,7 @@ import Shikumi.Optimize.MIPRO import Shikumi.Optimize.Pareto import Shikumi.Optimize.RandomSearch +import Shikumi.Optimize.ReAnchor import Shikumi.Optimize.Report (CandidateStatus (Unscored), OptimizationReport (..), RunStatus (BudgetStopped)) import Shikumi.Optimize.Search import Shikumi.Optimize.Structure diff --git a/src/Shikumi/Optimize/Execution.hs b/src/Shikumi/Optimize/Execution.hs --- a/src/Shikumi/Optimize/Execution.hs +++ b/src/Shikumi/Optimize/Execution.hs @@ -13,6 +13,7 @@ runSearchSession, markLegacy, setSelection, + setSelectedCandidate, sessionStopped, canStartCandidate, addPredictedWork, @@ -194,6 +195,20 @@ reportedPolicy = Just policy, validationMode = mode } + +-- | Pin the strategy's accepted candidate after an additional selection gate +-- such as a held-out fold check. Trial scores alone must not bypass that gate. +setSelectedCandidate :: (Prim :> es, Error ShikumiError :> es) => SearchSession es -> Maybe CandidateId -> Text -> Eff es () +setSelectedCandidate s chosen reason = do + current <- readIORef (state s) + ident <- case chosen of + Nothing -> pure Nothing + Just (CandidateId owner ix) -> do + unless + (owner == reserved s && any (\c -> candidateId c == ix && candidateStatus c == CandidateCompleted) (candidates current)) + (throwError (ValidationFailure "selected candidate must be completed in this session")) + pure (Just ix) + modifyReport s (\r -> r {selectedCandidate = ident, frontier = maybe [] (: []) ident, resultStatus = Just (maybe Unscored (const CandidateCompleted) ident), selectionReason = reason}) addPredictedWork :: (Prim :> es) => SearchSession es -> Int -> Eff es () addPredictedWork s n = modifyReport s (\r -> r {predictedWork = predictedWork r + max 0 n}) diff --git a/src/Shikumi/Optimize/ReAnchor.hs b/src/Shikumi/Optimize/ReAnchor.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/ReAnchor.hs @@ -0,0 +1,230 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Fit local thresholds, cuts and weights. All trials share run-local response +-- replay inside SearchSession admission; only new requests consume operations. +module Shikumi.Optimize.ReAnchor + ( ReAnchorConfig (..), + defaultReAnchorConfig, + reAnchor, + reAnchorWith, + ReAnchorReport (..), + FieldCalibration (..), + ObservationSummary (..), + ) +where + +import Baikai (Response) +import Control.Monad (foldM, when) +import Data.Aeson (ToJSON, Value (..), object, (.=)) +import Data.Aeson.Key qualified as K +import Data.Aeson.KeyMap qualified as KM +import Data.List (nub) +import Data.Map.Strict qualified as Map +import Data.Maybe (fromMaybe) +import Data.Text (Text) +import Data.Text qualified as T +import Effectful (Eff, (:>)) +import Effectful.Concurrent (Concurrent) +import Effectful.Error.Static (Error, throwError) +import Effectful.Prim.IORef (Prim, newIORef, readIORef) +import GHC.Generics (Generic) +import Shikumi.Adapter (evidenceSchema, responseValue) +import Shikumi.Compile.Types (CompiledProgram (..)) +import Shikumi.Decision +import Shikumi.Decision.Runtime (decodeDecisionEvidence, evidenceOf) +import Shikumi.Decision.Setting +import Shikumi.Effect.Time (Time) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.Eval (Dataset, FailurePolicy (..), Metric, datasetSize) +import Shikumi.LLM (LLM) +import Shikumi.Optimize.Execution +import Shikumi.Optimize.ReAnchor.Fit +import Shikumi.Optimize.ReAnchor.Replay +import Shikumi.Optimize.Report qualified as R +import Shikumi.Optimize.Types (ConfiguredOptimizer (..)) +import Shikumi.Program +import Shikumi.Trace.Node (NodePath, programNodePaths) + +-- | Validation is report-only. For a full search, use RunLimits with +-- candidateLimit = 2000; replayed candidates consume no operation budget. +data ReAnchorConfig i o = ReAnchorConfig + {validation :: !(Maybe (Dataset i o)), maxCandidatesPerStep :: !Int, foldCount :: !Int} + +defaultReAnchorConfig :: ReAnchorConfig i o +defaultReAnchorConfig = ReAnchorConfig Nothing 40 5 + +data ObservationSummary = ObservationSummary + {calls :: !Int, distinctValues :: !Int, minimumValue :: !(Maybe Double), maximumValue :: !(Maybe Double), candidatesTried :: !Int} + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +data FieldCalibration = FieldCalibration + { nodeIndex :: !Int, + nodePath :: !Text, + field :: !Text, + parameter :: !Text, + fittedValue :: !(Maybe DecisionSetting), + skippedReason :: !(Maybe Text), + trainScoreOriginal :: !Double, + trainScoreAtStart :: !Double, + trainScoreRetained :: !Double, + observations :: !ObservationSummary, + foldChecksPassed :: !Int, + foldChecksFailed :: !Int + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +data ReAnchorReport = ReAnchorReport + {trainScoreBefore :: !Double, trainScore :: !Double, valScoreBefore :: !(Maybe Double), valScore :: !(Maybe Double), fitted :: ![FieldCalibration]} + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +reAnchor :: ReAnchorConfig i o -> ConfiguredOptimizer i o +reAnchor cfg = ConfiguredOptimizer $ \session ds metric program -> fst <$> reAnchorWith cfg session ds metric program + +mean :: [Double] -> Double +mean [] = 0 +mean xs = sum xs / fromIntegral (length xs) + +reAnchorWith :: + forall es i o. + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + ReAnchorConfig i o -> SearchSession es -> Dataset i o -> Metric o -> Program i o -> Eff es (CompiledProgram i o, ReAnchorReport) +reAnchorWith cfg session train metric original = do + when (datasetSize train == 0) (throwError (ValidationFailure "ReAnchor: empty training set")) + when (maxCandidatesPerStep cfg < 2 || foldCount cfg < 2) (throwError (ValidationFailure "ReAnchor: candidate and fold counts must be at least two")) + let targets = [(ix, path, j) | (ix, (path, js)) <- zip [0 ..] (zip (programNodePaths original) (nodeDecisionFieldsIndexed original)), j <- js] + when (null targets) (throwError (ValidationFailure "ReAnchor: program has no rich decision outputs")) + store <- newReplayStore + let score ds metadata program = do + reservation <- reserveCandidate session + case reservation of + Nothing -> pure Nothing + Just ident -> do + annotateCandidate session ident metadata + seen <- newIORef [] + report <- + evaluateCandidate + session + ident + ds + ( \input -> do + result <- runResponses store seen program input + pure (result, []) + ) + (const FailAbort) + metric + R.qualityPolicy + (scalarObjectives metric) + responses <- readIORef seen + pure $ + if R.candidateStatus report == R.CandidateCompleted + then Just (ident, map snd (R.exampleScores report), responses) + else Nothing + scoreTrain = score train + optionalValidation program = case validation cfg of + Nothing -> pure Nothing + Just ds + | datasetSize ds == 0 -> pure Nothing + | otherwise -> fmap (\(_, scores, _) -> mean scores) <$> score ds (Map.singleton "role" "validation-report-only") program + initial <- scoreTrain (Map.singleton "role" "baseline") original + valBefore <- optionalValidation original + case initial of + Nothing -> do + setSelectedCandidate session Nothing "ReAnchor: unscored baseline; budget exhausted" + pure (CompiledProgram original, ReAnchorReport 0 0 valBefore Nothing []) + Just (initialId, initialScores, _) -> do + let step (program, retainedId, retainedScores, rows) (ix, path, j) = do + before <- scoreTrain (metadata ix path j "original" Nothing) program + case before of + Nothing -> pure (program, retainedId, retainedScores, rows) + Just (beforeId, beforeScores, _) -> do + let saved = decisionSettingFor (judgmentField j) (foldParams program !! ix) + start = fromMaybe (defaultSetting j) saved + set setting = mapParamsAt ix (setDecisionSetting (judgmentField j) setting) + base <- scoreTrain (metadata ix path j "start" (Just start)) (set start program) + case base of + Nothing -> pure (program, retainedId, retainedScores, rows) + Just (baseId, baseScores, responses) -> do + evidence <- either throwError pure (traverse (readEvidence j . snd) (filter ((== path) . fst) responses)) + let trials current scores ident candidates = do + tried <- + traverse + ( \(setting, tie) -> do + measured <- scoreTrain (metadata ix path j "trial" (Just setting)) (set setting program) + pure (fmap (\(cid, sc, _) -> (sc, tie, (setting, cid))) measured) + ) + candidates + let available = [x | Just x <- tried] + (kept, refused) = selectKept (foldCount cfg) scores available + pure $ case kept of + Just (sc, _, (setting, cid)) -> (setting, sc, cid, 1, 0, length available) + Nothing -> (current, scores, ident, 0, if refused then 1 else 0, length available) + fitNoul = case start of + NoulSetting t -> trials start baseScores baseId [(NoulSetting c, tie) | (c, tie) <- thresholdCandidates (maxCandidatesPerStep cfg) [p | FieldNoul (NoulEvidence p) <- evidence] t] + _ -> pure (start, baseScores, baseId, 0, 0, 0) + fitCuts = case start of + ScoreSetting cs -> + foldM + ( \(current, sc, cid, passed, failed, count) k -> case current of + ScoreSetting values -> do + let lo = if k == 0 then 0 else values !! (k - 1) + hi = if k + 1 == length values then fromIntegral (length values) else values !! (k + 1) + means = [v | FieldScore e <- evidence, Right (v, _) <- [deriveScore values e]] + candidates = [(ScoreSetting (take k values ++ [c] ++ drop (k + 1) values), tie) | (c, tie) <- cutCandidates (maxCandidatesPerStep cfg) means lo hi (values !! k)] + (next, sc', cid', p, f, n) <- trials current sc cid candidates + pure (next, sc', cid', passed + p, failed + f, count + n) + _ -> pure (current, sc, cid, passed, failed, count) + ) + (start, baseScores, baseId, 0, 0, 0) + [0 .. length cs - 1] + _ -> pure (start, baseScores, baseId, 0, 0, 0) + fitWeights labels = + foldM + ( \(current, sc, cid, passed, failed, count) label -> case current of + ChoiceSetting ws -> do + let weight l = Map.findWithDefault 1 l ws + flips = [rival / p | FieldChoice (ChoiceEvidence ps _) <- evidence, let p = Map.findWithDefault 0 label ps, let rival = maximum (0 : [weight l * Map.findWithDefault 0 l ps | l <- labels, l /= label]), p > 0, rival > 0] + candidates = [(ChoiceSetting (Map.insert label c ws), tie) | (c, tie) <- weightCandidates (maxCandidatesPerStep cfg) flips (weight label)] + (next, sc', cid', p, f, n) <- trials current sc cid candidates + pure (next, sc', cid', passed + p, failed + f, count + n) + _ -> pure (current, sc, cid, passed, failed, count) + ) + (start, baseScores, baseId, 0, 0, 0) + labels + (setting, scores, ident, passed, failed, count) <- case judgmentKind j of + JudgeBoolean {} -> fitNoul + JudgeOrdered {} -> fitCuts + JudgeChoice labels -> fitWeights (map fst labels) + let (accepted, refused) = selectKept (foldCount cfg) beforeScores [(scores, [], ())] + keep = case accepted of Just _ -> True; _ -> False + chosen = if keep then set setting program else program + values = concatMap observedValues evidence + summary = ObservationSummary (length evidence) (length (nub values)) (safeMinimum values) (safeMaximum values) count + row = FieldCalibration ix (T.pack (show path)) (judgmentField j) (kindName j) (if keep then Just setting else Nothing) (if keep then Nothing else Just "fitted behaviour did not beat the original") (mean beforeScores) (mean baseScores) (mean (if keep then scores else beforeScores)) summary (passed + if keep then 1 else 0) (failed + if refused then 1 else 0) + pure (chosen, if keep then ident else beforeId, if keep then scores else beforeScores, rows ++ [row]) + (final, selected, scores, rows) <- foldM step (original, initialId, initialScores, []) targets + valAfter <- optionalValidation final + setSelection session "ReAnchor: strict improvement with held-out fold check" R.qualityPolicy + setSelectedCandidate session (Just selected) "ReAnchor: retained settings passed strict improvement and held-out folds" + pure (CompiledProgram final, ReAnchorReport (mean initialScores) (mean scores) valBefore valAfter rows) + where + metadata :: Int -> NodePath -> Judgment -> Text -> Maybe DecisionSetting -> Map.Map Text Text + metadata ix path j role setting = Map.fromList [("node", T.pack (show ix)), ("path", T.pack (show path)), ("field", judgmentField j), ("parameter", kindName j), ("role", role), ("value", T.pack (show setting))] + kindName j = case judgmentKind j of JudgeBoolean {} -> "threshold"; JudgeOrdered {} -> "cuts"; JudgeChoice {} -> "weights" + observedValues = \case FieldNoul (NoulEvidence p) -> [p]; FieldChoice (ChoiceEvidence ps _) -> Map.elems ps; FieldScore e@(ScoreEvidence _ _) -> case deriveScore [] e of Right (v, _) -> [v]; _ -> [] + safeMinimum [] = Nothing + safeMinimum xs = Just (minimum xs) + safeMaximum [] = Nothing + safeMaximum xs = Just (maximum xs) + +readEvidence :: Judgment -> Response -> Either ShikumiError FieldEvidence +readEvidence j response = case evidenceOf [j] response of + Just result -> result >>= \fields -> maybe (Left (MissingField (judgmentField j))) Right (Map.lookup (judgmentField j) fields) + Nothing -> do + value <- responseValue (evidenceSchema (object ["properties" .= object []]) [j]) response + raw <- case value of + Object fields -> maybe (Left (MissingField (judgmentField j))) Right (KM.lookup (K.fromText (judgmentField j)) fields) + _ -> Left (SchemaMismatch "ReAnchor: expected evidence object") + decodeDecisionEvidence j raw diff --git a/src/Shikumi/Optimize/ReAnchor/Fit.hs b/src/Shikumi/Optimize/ReAnchor/Fit.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/ReAnchor/Fit.hs @@ -0,0 +1,76 @@ +-- | Bounded candidates and deterministic held-out selection, without model calls. +module Shikumi.Optimize.ReAnchor.Fit (gaps, tidy, folds, pickBest, selectKept, thresholdCandidates, cutCandidates, weightCandidates) where + +import Data.List (maximumBy, nub, sort, sortOn) +import Data.Ord (comparing) +import Data.Word (Word64) +import GHC.Float (castDoubleToWord64, castWord64ToDouble) + +tidy :: Double -> Double -> Double -> Double +tidy lo hi midpoint = case [x | digits <- [1 .. 11 :: Int], let scale = 10 ^^ digits, let x = fromInteger (round (midpoint * scale)) / scale, x > lo, x < hi] of + x : _ -> x + [] -> midpoint + +gaps :: Int -> Bool -> [Double] -> Double -> Double -> [(Double, Double)] +gaps limit geometric values lo hi + | lo >= hi || (geometric && lo <= 0) = [] + | otherwise = [(tidy a b mid, width a b) | (a, b) <- zip points (drop 1 points), let mid = if geometric then sqrt (a * b) else (a + b) / 2, mid > a, mid < b] + where + inner = nub (sort [x | x <- values, x > lo, x < hi]) + count = max 1 (limit - 1) + thinned + | length inner <= count = inner + | count == 1 = [inner !! (length inner `div` 2)] + | otherwise = nub [inner !! round (fromIntegral i * fromIntegral (length inner - 1) / fromIntegral (count - 1) :: Double) | i <- [0 .. count - 1]] + points = lo : thinned ++ [hi] + width a b = if geometric then log (b / a) else b - a + +-- | A fixed-seed shuffle followed by round-robin fold assignment. It is stable +-- across runs and independent of candidate order or mutable random state. +folds :: Int -> Int -> [[Int]] +folds requested n = [[x | (i, x) <- zip [0 :: Int ..] shuffled, i `mod` k == j] | j <- [0 .. k - 1]] + where + k = min (max 1 requested) n + shuffled = shuffle (1 :: Word64) [0 .. n - 1] + shuffle _ [] = [] + shuffle seed xs = + let seed' = seed * 6364136223846793005 + 1442695040888963407 + ix = fromIntegral (seed' `mod` fromIntegral (length xs)) + (a, b) = splitAt ix xs + in case b of y : ys -> y : shuffle seed' (a ++ ys); [] -> [] + +-- | Strict improvement; on equal score and tie key the earliest trial wins. +pickBest :: [Int] -> [Double] -> [([Double], [Double], a)] -> Maybe ([Double], [Double], a) +pickBest rows start tried = case filter (\(scores, _, _) -> total scores > total start) tried of + [] -> Nothing + x : xs -> Just (foldl choose x xs) + where + total scores = sum [scores !! i | i <- rows] + key (scores, tie, _) = (total scores, tie) + choose a b = if key b > key a then b else a + +selectKept :: Int -> [Double] -> [([Double], [Double], a)] -> (Maybe ([Double], [Double], a), Bool) +selectKept k start tried = case pickBest rows start tried of + Nothing -> (Nothing, False) + Just best + | held > sum start -> (Just best, False) + | otherwise -> (Nothing, True) + where + rows = [0 .. length start - 1] + held = sum [sum [scores !! i | i <- fold] | fold <- folds k (length start), let outside = filter (`notElem` fold) rows, let scores = maybe start (\(s, _, _) -> s) (pickBest outside start tried)] + +thresholdCandidates :: Int -> [Double] -> Double -> [(Double, [Double])] +thresholdCandidates limit observed start = [(t, [w, negate (abs (t - start)), t]) | (t, w) <- candidates, t /= start] + where + sorted = nub (sort observed) + adjacent = [(b, b - a) | (a, b) <- zip sorted (drop 1 sorted), a >= 0, castWord64ToDouble (castDoubleToWord64 a + 1) == b] + candidates = gaps limit False observed 0 1 ++ adjacent ++ [(0, 0) | 0 `elem` observed] + +cutCandidates :: Int -> [Double] -> Double -> Double -> Double -> [(Double, [Double])] +cutCandidates limit means lo hi current = [(c, [w]) | (c, w) <- gaps limit False means lo hi, c /= current] + +weightCandidates :: Int -> [Double] -> Double -> [(Double, [Double])] +weightCandidates limit flips current = [(c, [w]) | (c, w) <- gaps limit True flips low high, c /= current] + where + low = case flips of [] -> 1e-3; _ -> max 1e-3 (minimum flips / 4) + high = case flips of [] -> 1e3; _ -> min 1e3 (maximum flips * 4) diff --git a/src/Shikumi/Optimize/ReAnchor/Replay.hs b/src/Shikumi/Optimize/ReAnchor/Replay.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/ReAnchor/Replay.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE RankNTypes #-} + +-- | Run-local response replay, placed inside candidate admission. Node-path +-- capture works for ordinary Predict leaves without requiring typed codecs. +module Shikumi.Optimize.ReAnchor.Replay (ReplayStore, newReplayStore, withReplay, runResponses) where + +import Baikai (Response) +import Data.Map.Strict qualified as Map +import Effectful (Eff, (:>)) +import Effectful.Dispatch.Dynamic (interpose) +import Effectful.Error.Static (Error, catchError) +import Effectful.Prim.IORef (IORef, Prim, modifyIORef', newIORef, readIORef) +import Shikumi.Cache.Key (CacheKey, cacheKey) +import Shikumi.Error (ShikumiError) +import Shikumi.LLM (LLM (..), complete, stream) +import Shikumi.Program (Program, runProgram) +import Shikumi.Trace.Node (NodePath) +import Shikumi.Trace.Program (walkProgramWith) + +newtype ReplayStore = ReplayStore (IORef (Map.Map CacheKey Response)) + +newReplayStore :: (Prim :> es) => Eff es ReplayStore +newReplayStore = ReplayStore <$> newIORef Map.empty + +withReplay :: (LLM :> es, Prim :> es) => ReplayStore -> Eff es a -> Eff es a +withReplay (ReplayStore store) = interpose $ \_ -> \case + Complete m c o -> do + let key = cacheKey m c o + recorded <- Map.lookup key <$> readIORef store + case recorded of + Just response -> pure response + Nothing -> do + response <- complete m c o + modifyIORef' store (Map.insert key response) + pure response + Stream m c o -> stream m c o + +runResponses :: forall es i o. (LLM :> es, Prim :> es, Error ShikumiError :> es) => ReplayStore -> IORef [(NodePath, Response)] -> Program i o -> i -> Eff es (Either ShikumiError o) +runResponses store seen program input = + (Right <$> walkProgramWith (\_ _ action -> action) leaf (pure ()) program input) `catchError` \_ err -> pure (Left err) + where + leaf :: forall x y. NodePath -> Program x y -> x -> Eff es y + leaf path node x = + interpose + ( \_ -> \case + Complete m c o -> do + response <- withReplay store (complete m c o) + modifyIORef' seen (++ [(path, response)]) + pure response + Stream m c o -> stream m c o + ) + (runProgram node x) diff --git a/test/ExecutionSpec.hs b/test/ExecutionSpec.hs --- a/test/ExecutionSpec.hs +++ b/test/ExecutionSpec.hs @@ -16,7 +16,7 @@ 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.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 @@ -38,7 +38,19 @@ tests = testGroup "Execution" - [ testCase "invalid controls reject without dispatch" $ do + [ testCase "strategy selection preserves its accepted reservation" $ do + r <- run $ runSearchSession (cfg 8 1) $ \s -> do + reserved <- reserveCandidate s + case reserved of + Nothing -> throwError (ValidationFailure "missing reservation") + Just ix -> do + _ <- evaluateCandidate s ix ds (runProgramObserved sentimentProg) (const FailAbort) exactMatch qualityPolicy (scalarObjectives exactMatch) + setSelectedCandidate s (Just ix) "held-out gate accepted baseline" + check r $ \(_, report) -> do + selectedCandidate report @?= Just 0 + frontier report @?= [0] + selectionReason report @?= "held-out gate accepted baseline", + 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 diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -20,6 +20,7 @@ import OptimizeSpec qualified import ProposeSpec qualified import RandomSearchSpec qualified +import ReAnchorSpec qualified import SearchSpec qualified import SeedingSpec qualified import StructureSpec qualified @@ -44,6 +45,7 @@ GepaSpec.tests, FeedbackSpec.tests, KNNSpec.tests, + ReAnchorSpec.tests, RandomSearchSpec.tests, SearchSpec.tests, SeedingSpec.tests, diff --git a/test/ReAnchorSpec.hs b/test/ReAnchorSpec.hs new file mode 100644 --- /dev/null +++ b/test/ReAnchorSpec.hs @@ -0,0 +1,210 @@ +{-# LANGUAGE DataKinds #-} + +module ReAnchorSpec (tests) where + +import Baikai qualified +import Baikai.Models.Generated (openai_gpt_4o_mini) +import Control.Lens ((^.)) +import Data.Aeson (Value (..), eitherDecode, encode, toJSON) +import Data.Aeson.KeyMap qualified as KM +import Data.Generics.Labels () +import Data.IORef (modifyIORef', newIORef, readIORef) +import Data.List (sort) +import Data.Map.Strict qualified as Map +import Data.Text qualified as T +import Data.Vector qualified as V +import Effectful (liftIO, runEff) +import Effectful.Concurrent (runConcurrent) +import Effectful.Dispatch.Dynamic (interpose) +import Effectful.Error.Static (runErrorNoCallStack) +import Effectful.Prim (runPrim) +import GHC.Float (castDoubleToWord64, castWord64ToDouble) +import Shikumi.Compile.Types (CompiledProgram (..)) +import Shikumi.Decision +import Shikumi.Decision.Setting +import Shikumi.Effect.Time (runTime) +import Shikumi.Error +import Shikumi.Eval qualified as E +import Shikumi.LLM (LLM (..), complete, stream) +import Shikumi.Optimize hiding (field) +import Shikumi.Optimize.Execution +import Shikumi.Optimize.ReAnchor.Fit +import Shikumi.Optimize.Report qualified as R +import Shikumi.Program +import Shikumi.Routing +import Shikumi.Schema (field, unField) +import Shikumi.Signature (mkSignature) +import Shikumi.Testing.DecisionFixtures +import Shikumi.Testing.Response (evidenceJSONResponse, evidenceMarkerResponse) +import Shikumi.Testing.StubLLM +import Shikumi.Testing.SystemOne +import Test.Tasty +import Test.Tasty.HUnit + +number :: Value -> Double +number (Object o) | Just (Object inputs) <- KM.lookup "inputs" o, Just (String t) <- KM.lookup "ticket" inputs = read (T.unpack t) +number _ = 0 + +assessment :: Double -> Double -> Double -> Assessment +assessment p mean categoryP = Assessment (field (Noul (p >= 0.72) p 0)) (field (Score mean (if mean >= 1.6 then 2 else if mean >= 0.5 then 1 else 0) (probs mean) Nothing)) (field (Choice (if categoryP >= 0.7 then Technical else Billing) (Map.fromList [("technical", categoryP), ("billing", 1 - categoryP)]) Nothing)) + +probs :: Double -> V.Vector Double +probs mean + | mean <= 1 = V.fromList [1 - mean, mean, 0] + | otherwise = V.fromList [0, 2 - mean, mean - 1] + +response :: Double -> [(T.Text, FieldEvidence)] +response p = [("urgent", FieldNoul (NoulEvidence p)), ("severity", FieldScore (ScoreEvidence (probs (2 * p)) Nothing)), ("category", FieldChoice (ChoiceEvidence (Map.fromList [("technical", p), ("billing", 1 - p)]) Nothing))] + +metric :: E.Metric Assessment +metric expected actual = E.boolScore (flag expected == flag (E.predictionPrimary actual)) + where + flag (Assessment u _ _) = case unField u of Noul v _ _ -> v + +train :: E.Dataset Ticket Assessment +train = E.dataset [E.example (Ticket (field (T.pack (show p)))) (assessment p (2 * p) p) | p <- [0.1, 0.2, 0.3, 0.4, 0.55, 0.6, 0.65, 0.75, 0.8, 0.9]] + +cfg :: RunConfig +cfg = defaultRunConfig {runLimits = RunLimits 100 1000 1 1} + +tests :: TestTree +tests = + testGroup + "ReAnchor" + [ testCase "gaps tidy, thin, and separate adjacent floats" $ do + map fst (gaps 40 False [0.98, 1] 0 1) @?= [0.5, 0.99] + assertBool "bounded" (length (gaps 40 False [fromIntegral n / 100 | n <- [1 .. 50 :: Int]] 0 1) <= 40) + let lo = 0.5; hi = castWord64ToDouble (castDoubleToWord64 lo + 1) + assertBool "adjacent upper" (hi `elem` map fst (thresholdCandidates 40 [lo, hi] 0.2)) + assertBool "zero threshold" (0 `elem` map fst (thresholdCandidates 40 [0, 0.5] 0.2)) + assertBool "ordered cuts" (all (\(c, _) -> c > 0.5 && c < 1.8) (cutCandidates 40 [0.8, 1.2, 1.7] 0.5 1.8 1.5)), + testCase "fold selection refuses isolated gain and accepts spread gain" $ do + let (a, refused) = selectKept 5 [0, 0, 0, 0, 0] [([1, 0, 0, 0, 0], [], ())] + fmap (const ()) a @?= Nothing + refused @?= True + let (b, _) = selectKept 5 [0, 0, 0, 0, 0] [([1, 1, 1, 1, 1], [], ())] + fmap (const ()) b @?= Just () + sort (concat (folds 5 13)) @?= [0 .. 12], + testCase "fits threshold, replays every candidate, and persists only settings" $ do + calls <- newIORef 0 + let program = Predict assessmentSig emptyParams + result <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError . runRouting (systemOneModel "jev-latest") . runStubLLM (systemOneResponder (response . number)) . interpose (\_ -> \case Complete m c o -> liftIO (modifyIORef' calls (+ 1)) >> complete m c o; Stream m c o -> stream m c o) . routeLLM $ + runSearchSession cfg (\s -> reAnchorWith defaultReAnchorConfig s train metric program) + case result of + Right (Right (CompiledProgram tuned, report), diagnostics) -> do + trainScoreBefore report @?= 0.7 + trainScore report @?= 1 + case decisionSettingFor "urgent" (head (foldParams tuned)) of + Just (NoulSetting t) -> assertBool "observed gap" (t > 0.65 && t <= 0.75) + other -> assertFailure (show other) + let ps = foldParams tuned + eitherDecode (encode ps) @?= Right ps + map instructionOverride ps @?= [Nothing] + map demos ps @?= [[]] + R.admittedOperations diagnostics @?= 10 + assertBool "accepted candidate recorded" (R.selectedCandidate diagnostics /= Nothing) + Left err -> assertFailure (show err) + Right (Left err, _) -> assertFailure (show err) + readIORef calls >>= (@?= 10), + testCase "fold-refused fit restores absent entries and budgets keep accepted state" $ do + let ds = E.dataset [E.example (Ticket (field (T.pack (show p)))) (assessment p (2 * p) p) | p <- [0.6, 0.1, 0.2, 0.3, 0.8]] + result <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError . runRouting (systemOneModel "jev-latest") . runStubLLM (systemOneResponder (response . number)) . routeLLM $ + optimizeWith cfg (reAnchor defaultReAnchorConfig) ds metric (Predict assessmentSig emptyParams) + case result of + Right (CompiledProgram tuned, _) -> map decisionSettings (foldParams tuned) @?= [Map.empty] + Left err -> assertFailure (show err), + testCase "score cuts and choice weights improve distributed errors" $ do + let scoreMetric expected actual = E.boolScore (lev expected == lev (E.predictionPrimary actual)) + lev (Assessment _ s _) = case unField s of Score _ l _ _ -> l + choiceMetric expected actual = E.boolScore (cat expected == cat (E.predictionPrimary actual)) + cat (Assessment _ _ c) = case unField c of Choice v _ _ -> v + scoreData = E.dataset [E.example (Ticket (field (T.pack (show p)))) (assessment p (2 * p) p) | p <- [0.1, 0.2, 0.3, 0.4, 0.755, 0.77, 0.785, 0.825, 0.9, 0.95]] + run ds met = + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError . runRouting (systemOneModel "jev-latest") . runStubLLM (systemOneResponder (response . number)) . routeLLM $ + runSearchSession cfg (\s -> reAnchorWith defaultReAnchorConfig s ds met (Predict assessmentSig emptyParams)) + fittedScore <- run scoreData scoreMetric + case fittedScore of + Right (Right (CompiledProgram tuned, report), _) -> do + trainScore report @?= 1 + case decisionSettingFor "severity" (head (foldParams tuned)) of + Just (ScoreSetting [a, b]) -> assertBool "ordered fitted boundary" (a < b && b > 1.57 && b <= 1.65) + _ -> assertFailure "missing fitted cuts" + Left err -> assertFailure (show err) + Right (Left err, _) -> assertFailure (show err) + fittedChoice <- run train choiceMetric + case fittedChoice of + Right (Right (CompiledProgram tuned, report), _) -> do + trainScore report @?= 1 + case decisionSettingFor "category" (head (foldParams tuned)) of + Just (ChoiceSetting ws) -> assertBool "technical has lower relative weight" (Map.findWithDefault 1 "technical" ws < Map.findWithDefault 1 "billing" ws) + _ -> assertFailure "missing fitted weights" + Left err -> assertFailure (show err) + Right (Left err, _) -> assertFailure (show err), + testCase "fits a three-label choice" $ do + let ds = E.dataset [E.example (Ticket (field (T.pack (show p)))) (ChoiceAssessment (Choice (if p >= 0.7 then FirstLabel else SecondLabel) Map.empty Nothing)) | p <- [0.1, 0.2, 0.3, 0.4, 0.55, 0.6, 0.65, 0.75, 0.8, 0.9]] + met (ChoiceAssessment (Choice expected _ _)) predicted = case E.predictionPrimary predicted of ChoiceAssessment (Choice actual _ _) -> E.boolScore (expected == actual) + resp state = [("selection", FieldChoice (ChoiceEvidence (Map.fromList [("FirstLabel", p), ("SecondLabel", 0.8 * (1 - p)), ("ThirdLabel", 0.2 * (1 - p))]) Nothing))] where p = number state + result <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError . runRouting (systemOneModel "jev-latest") . runStubLLM (systemOneResponder resp) . routeLLM $ + runSearchSession cfg (\s -> reAnchorWith defaultReAnchorConfig s ds met (Predict (mkSignature "Choose a label") emptyParams)) + case result of + Right (Right (CompiledProgram tuned, report), _) -> do + trainScore report @?= 1 + case decisionSettingFor "selection" (head (foldParams tuned)) of + Just (ChoiceSetting ws) -> assertBool "over-picked label downweighted" (Map.findWithDefault 1 "FirstLabel" ws < Map.findWithDefault 1 "SecondLabel" ws) + _ -> assertFailure "missing fitted weights" + Left err -> assertFailure (show err) + Right (Left err, _) -> assertFailure (show err), + testCase "changed upstream answers admit new downstream requests; budget stops safely" $ do + let program = Compose (Predict assessmentSig emptyParams) (Predict (mkSignature "Downstream decision") emptyParams) + respond state = response $ case state of + Object o | Just (Object inputs) <- KM.lookup "inputs" o, Just (Bool flag) <- KM.lookup "urgent" inputs -> if flag then 0.9 else 0.1 + _ -> number state + result <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError . runRouting (systemOneModel "jev-latest") . runStubLLM (systemOneResponder respond) . routeLLM $ + optimizeWith cfg (reAnchor defaultReAnchorConfig) train metric program + case result of + Right (_, diagnostics) -> do + assertBool "new downstream states admitted" (R.admittedOperations diagnostics > 12) + assertBool "new requests respect admission ceiling" (R.admittedOperations diagnostics <= 100) + Left err -> assertFailure (show err) + stopped <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError . runRouting (systemOneModel "jev-latest") . runStubLLM (systemOneResponder (response . number)) . routeLLM $ + optimizeWith (cfg {runLimits = RunLimits 100 1 1 1}) (reAnchor defaultReAnchorConfig) train metric (Predict assessmentSig emptyParams) + case stopped of + Right (CompiledProgram tuned, diagnostics) -> do + map decisionSettings (foldParams tuned) @?= [Map.empty] + R.selectedCandidate diagnostics @?= Just 0 + Left err -> assertFailure (show err), + testCase "generative JSON and marker evidence fit the same threshold" $ do + let payload = \case FieldNoul e -> toJSON e; FieldChoice (ChoiceEvidence ps _) -> toJSON (ChoiceEvidence ps (Just 0.8)); FieldScore (ScoreEvidence ps _) -> toJSON (ScoreEvidence ps (Just 0.8)) + mapM_ + ( \builder -> do + result <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError . runRouting openai_gpt_4o_mini . runStubLLM (\ctx -> builder [(k, payload e) | (k, e) <- response (numberFromContext ctx)]) . routeLLM $ + optimizeWith cfg (reAnchor defaultReAnchorConfig) train metric (Predict assessmentSig emptyParams) + case result of + Right (CompiledProgram tuned, diagnostics) -> do + case decisionSettingFor "urgent" (head (foldParams tuned)) of + Just (NoulSetting t) -> assertBool "fitted gap" (t > 0.65 && t <= 0.75) + _ -> assertFailure "missing threshold" + R.admittedOperations diagnostics @?= 10 + Left err -> assertFailure (show err) + ) + [evidenceJSONResponse, evidenceMarkerResponse], + testCase "no decision output fails before dispatch" $ do + let program = embed (\input -> pure input) :: Program Ticket Ticket + result <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError . runStubLLM (const (systemOneResponse [])) $ + optimizeWith cfg (reAnchor defaultReAnchorConfig) (E.dataset [E.example (Ticket (field "x")) (Ticket (field "x"))]) E.exactMatch program + assertBool "typed refusal" (case result of Left (ValidationFailure _) -> True; _ -> False) + ] + +-- Native contexts retain the ordinary user turn. Its typed Field Show spelling +-- contains the ticket number; use the shared prompt text projection for fixtures. +numberFromContext :: Baikai.Context -> Double +numberFromContext ctx = case [t | Baikai.UserMessage p <- V.toList (ctx ^. #messages), Baikai.UserText (Baikai.TextContent t) <- V.toList (p ^. #content)] of + t : _ -> case reads (T.unpack (T.filter (\c -> c `elem` ("0123456789." :: String)) t)) of (p, _) : _ -> p; _ -> 0 + _ -> 0