diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,10 @@ +# Changelog + +## Unreleased + +## 0.1.0.0 - 2026-06-13 + +### Added + +- Initial Hackage release of optimizer APIs for shikumi programs. +- Labeled few-shot, bootstrap few-shot, instruction search, random search, GEPA, KNN, COPRO, MIPRO, ensemble, Pareto, and proposal helpers. diff --git a/shikumi-optimize.cabal b/shikumi-optimize.cabal new file mode 100644 --- /dev/null +++ b/shikumi-optimize.cabal @@ -0,0 +1,111 @@ +cabal-version: 3.4 +name: shikumi-optimize +version: 0.1.0.0 +synopsis: The optimizer framework for shikumi LM programs (EP-10) +category: AI +description: + The optimizer framework for shikumi: search procedures that automatically + improve a 'Shikumi.Program.Program' by rewriting its per-node optimizable + parameters (instruction strings and few-shot demonstrations) and scoring each + candidate with @evaluate@ (EP-8) over a typed dataset and metric. Four + strategies ship — labeled few-shot demo selection, bootstrap few-shot (recover + demonstrations from a teacher's successful runs), instruction search + (MIPRO/COPRO-style coordinate ascent under an explicit call budget), and + ensemble search (combine complementary candidates by majority vote). Each + optimizer returns a @CompiledProgram@ (EP-9, MasterPlan integration point #6). + Every acceptance assertion runs fully offline against a deterministic stub LM, + so the whole package tests with no network and no API key. + +license: BSD-3-Clause +author: Nadeem Bitar +maintainer: nadeem@gmail.com +build-type: Simple +extra-doc-files: CHANGELOG.md + +common common-options + ghc-options: + -Wall -Wcompat -Widentities -Wincomplete-uni-patterns + -Wincomplete-record-updates -Wredundant-constraints + -fhide-source-paths -Wmissing-export-lists -Wpartial-fields + -Wmissing-deriving-strategies + + default-language: GHC2024 + default-extensions: + DeriveAnyClass + DuplicateRecordFields + OverloadedLabels + OverloadedStrings + +library + import: common-options + hs-source-dirs: src + exposed-modules: + Shikumi.Optimize + Shikumi.Optimize.Bootstrap + Shikumi.Optimize.COPRO + Shikumi.Optimize.Ensemble + Shikumi.Optimize.GEPA + Shikumi.Optimize.Instruction + Shikumi.Optimize.KNN + Shikumi.Optimize.LabeledFewShot + Shikumi.Optimize.MIPRO + Shikumi.Optimize.Pareto + Shikumi.Optimize.Propose + Shikumi.Optimize.Propose.Grounded + Shikumi.Optimize.Propose.Summarize + Shikumi.Optimize.Propose.Tips + Shikumi.Optimize.Propose.Types + Shikumi.Optimize.RandomSearch + Shikumi.Optimize.Search + Shikumi.Optimize.Types + + build-depends: + , aeson + , base >=4.20 && <5 + , containers + , effectful + , shikumi ^>=0.1.0.0 + , shikumi-compile ^>=0.1.0.0 + , shikumi-eval ^>=0.1.0.0 + , shikumi-trace ^>=0.1.0.0 + , text ^>=2.1 + , vector + +test-suite shikumi-optimize-test + import: common-options + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Main.hs + ghc-options: -threaded -with-rtsopts=-N + other-modules: + AcceptanceSpec + BootstrapSpec + CoproSpec + EnsembleSpec + GepaSpec + InstructionSpec + KNNSpec + LabeledFewShotSpec + Miprov2Spec + OptimizeSpec + ProposeSpec + RandomSearchSpec + StubLM + + build-depends: + , aeson + , baikai + , base + , containers + , effectful + , generic-lens + , lens + , shikumi ^>=0.1.0.0 + , shikumi-compile ^>=0.1.0.0 + , shikumi-eval ^>=0.1.0.0 + , shikumi-optimize ^>=0.1.0.0 + , shikumi-trace ^>=0.1.0.0 + , tasty + , tasty-hunit + , text + , vector diff --git a/src/Shikumi/Optimize.hs b/src/Shikumi/Optimize.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize.hs @@ -0,0 +1,66 @@ +-- | The public surface of the optimizer framework (EP-10). +-- +-- 'optimize' is the one stable entry point EP-12's CLI calls: it applies an +-- 'Optimizer' strategy to a starting program and returns a 'CompiledProgram'. The +-- shared search-state plumbing ('selectBest', 'scoreOn', 'freezeProgram') lives in +-- "Shikumi.Optimize.Search" and is re-exported here; the four strategies are +-- re-exported from their own modules. +-- +-- __No global mutable state.__ Candidate programs are threaded as ordinary values +-- through pure folds; the /only/ effectful step is scoring a candidate (which runs +-- it over the dataset via 'scoreOn'). This is the "thread candidates explicitly" +-- discipline the MasterPlan mandates. +module Shikumi.Optimize + ( -- * The driver + optimize, + + -- * Re-exports + module Shikumi.Optimize.Types, + module Shikumi.Optimize.Search, + module Shikumi.Optimize.LabeledFewShot, + module Shikumi.Optimize.Bootstrap, + module Shikumi.Optimize.Instruction, + module Shikumi.Optimize.COPRO, + module Shikumi.Optimize.MIPRO, + module Shikumi.Optimize.GEPA, + module Shikumi.Optimize.Pareto, + module Shikumi.Optimize.KNN, + module Shikumi.Optimize.RandomSearch, + module Shikumi.Optimize.Ensemble, + ) +where + +import Effectful (Eff, (:>)) +import Effectful.Concurrent (Concurrent) +import Effectful.Error.Static (Error) +import Effectful.Prim (Prim) +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.Bootstrap +import Shikumi.Optimize.COPRO +import Shikumi.Optimize.Ensemble +import Shikumi.Optimize.GEPA +import Shikumi.Optimize.Instruction +import Shikumi.Optimize.KNN +import Shikumi.Optimize.LabeledFewShot +import Shikumi.Optimize.MIPRO +import Shikumi.Optimize.Pareto +import Shikumi.Optimize.RandomSearch +import Shikumi.Optimize.Search +import Shikumi.Optimize.Types +import Shikumi.Program (Program) + +-- | Apply an optimizer to a starting program. A thin wrapper around the strategy +-- so the public API is one stable name (and so EP-12's CLI dispatches on one +-- function regardless of which optimizer the user picked). +optimize :: + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + Optimizer i o -> + Dataset i o -> + Metric o -> + Program i o -> + Eff es (CompiledProgram i o) +optimize opt train metric prog = runOptimizer opt train metric prog diff --git a/src/Shikumi/Optimize/Bootstrap.hs b/src/Shikumi/Optimize/Bootstrap.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Bootstrap.hs @@ -0,0 +1,88 @@ +{-# 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. +module Shikumi.Optimize.Bootstrap + ( bootstrapFewShot, + bootstrapFewShotWith, + BootstrapConfig (..), + defaultBootstrapConfig, + recoverDemo, + ) +where + +import Data.Aeson (ToJSON, toJSON) +import Effectful.Error.Static (catchError) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval (Example (..), datasetExamples, prediction, unScore) +import Shikumi.Optimize.LabeledFewShot (withDemos) +import Shikumi.Optimize.Search (freezeProgram) +import Shikumi.Optimize.Types (Budget (..), Optimizer (..)) +import Shikumi.Program (Demo (..), Program, runProgram) + +-- | Tunables for a bootstrap search. +data BootstrapConfig = BootstrapConfig + { -- | minimum metric score for a teacher run to contribute a demo (default + -- @1.0@: keep only exactly-correct runs) + passThreshold :: !Double, + -- | cap on the demos attached, so prompts stay bounded (default @4@) + maxBootstrappedDemos :: !Int + } + deriving stock (Eq, Show) + +-- | Keep only perfectly-correct teacher runs; attach at most four demos. +defaultBootstrapConfig :: BootstrapConfig +defaultBootstrapConfig = BootstrapConfig {passThreshold = 1.0, maxBootstrappedDemos = 4} + +-- | Recover a demonstration from one teacher run: pair the typed input with the +-- teacher's produced output, serialized to the JSON 'Demo' the run-time adapter +-- decodes back into the node's typed demo channel. (The JSON keys are the record +-- field names, so @fromModel@ round-trips them — see the unit test.) +recoverDemo :: (ToJSON i, ToJSON o) => i -> o -> Demo +recoverDemo i o = Demo {input = toJSON i, output = toJSON o} + +-- | Bootstrap few-shot with the default configuration. +bootstrapFewShot :: (ToJSON i, ToJSON o) => Program i o -> Budget -> Optimizer i o +bootstrapFewShot = bootstrapFewShotWith defaultBootstrapConfig + +-- | Bootstrap few-shot with an explicit configuration. The @teacher@ may be a +-- stronger or chain-of-thought variant of the student, or the student itself; it +-- must share the student's input/output types. +bootstrapFewShotWith :: + (ToJSON i, ToJSON o) => + BootstrapConfig -> + -- | teacher program whose successful runs supply demos + Program i o -> + Budget -> + Optimizer i o +bootstrapFewShotWith cfg teacher budget = Optimizer $ \train metric student -> do + -- Bound the number of teacher runs by the budget (each run is >= 1 LM call). + let exs = take (max 0 (maxLmCalls budget)) (datasetExamples train) + -- Run the teacher on one example; if it succeeds and the metric passes, + -- emit its recovered demo, else emit nothing. Recovery is total: a teacher + -- error yields no demo rather than aborting the search. + 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 [] + kept <- concat <$> mapM keepIfPassing exs + pure (freezeProgram (withDemos (take (maxBootstrappedDemos cfg) kept) student)) diff --git a/src/Shikumi/Optimize/COPRO.hs b/src/Shikumi/Optimize/COPRO.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/COPRO.hs @@ -0,0 +1,146 @@ +-- | COPRO (EP-21): coordinate-ascent prompt optimization. Where +-- 'Shikumi.Optimize.Instruction.instructionSearch' is one-shot per node, COPRO +-- improves each node's instruction over several /rounds/ (depth), proposing several +-- candidates per round (breadth) and feeding the /scored attempt history/ forward so +-- later rounds learn from what scored well. It is the principled generalization of +-- @instructionSearch@ (depth-1, no-history COPRO ≈ @instructionSearch@), kept +-- alongside it rather than replacing it. +-- +-- COPRO consumes EP-19's grounded proposer ('Shikumi.Optimize.Propose.proposeInstructions') +-- directly: each round's call passes the node's current instruction and its scored +-- 'PastInstruction' history, and the proposer returns ranked candidates with the +-- current instruction always retained (the safety property — a node never degrades). +-- +-- Output is V1's 'Shikumi.Compile.Types.CompiledProgram' via 'freezeProgram', invoked +-- through 'Shikumi.Optimize.optimize' and serialized unchanged (integration point #4). +module Shikumi.Optimize.COPRO + ( CoproConfig (..), + defaultCoproConfig, + copro, + ) +where + +import Control.Monad (foldM) +import Data.Aeson (ToJSON) +import Data.Maybe (fromMaybe) +import Effectful (Eff, (:>)) +import Effectful.Concurrent (Concurrent) +import Effectful.Error.Static (Error) +import Effectful.Prim (Prim) +import Shikumi.Effect.Time (Time) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval (Dataset, Metric, datasetSize) +import Shikumi.LLM (LLM) +import Shikumi.Optimize.Propose + ( PastInstruction (..), + ProposeRequest (..), + ProposeResult (..), + proposeInstructions, + ) +import Shikumi.Optimize.Search (freezeProgram, instructionAt, scoreOn, setNodeInstr) +import Shikumi.Optimize.Types (Budget (..), Optimizer (..), defaultBudget) +import Shikumi.Program (Program, foldParams) + +-- | COPRO's two knobs plus the shared 'Budget'. +data CoproConfig = CoproConfig + { -- | candidate instructions generated per node per round (clamped to @>= 2@) + breadth :: !Int, + -- | number of coordinate-ascent rounds (clamped to @>= 1@) + depth :: !Int, + -- | LM-call / candidate ceilings + budget :: !Budget + } + deriving stock (Eq, Show) + +-- | Breadth 4, depth 3, the default budget. +defaultCoproConfig :: CoproConfig +defaultCoproConfig = CoproConfig {breadth = 4, depth = 3, budget = defaultBudget} + +-- | Coordinate-ascent instruction optimization. Visits each node in @foldParams@ +-- order, optimizing it over @depth@ rounds against the already-improved earlier +-- nodes, threading one running LM-call count so the 'Budget' bounds the whole search. +copro :: (ToJSON i, ToJSON o) => CoproConfig -> Optimizer i o +copro cfg = Optimizer $ \train metric student -> do + let nNodes = length (foldParams student) + (final, _calls) <- + foldM + (\acc idx -> optimizeNode cfg train metric idx acc) + (student, 0) + [0 .. nNodes - 1] + pure (freezeProgram final) + +-- | Optimize node @idx@ over @depth@ rounds. Returns the program with node @idx@ set +-- to its best-found instruction, plus the updated LM-call count. Each round proposes +-- @breadth - 1@ fresh candidates (plus the retained current instruction) via the +-- grounded proposer fed the scored attempt history, scores the not-yet-seen ones on +-- the whole training set, records @(instruction, best-score)@, and sets the node to +-- the best so far. Every spend is gated against the 'Budget'. +optimizeNode :: + (ToJSON i, ToJSON o, LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + CoproConfig -> + Dataset i o -> + Metric o -> + Int -> + (Program i o, Int) -> + Eff es (Program i o, Int) +optimizeNode cfg train metric idx (prog0, calls0) = goRound 1 prog0 calls0 [] + where + dsSize = datasetSize train + bdth = max 2 (breadth cfg) + dpth = max 1 (depth cfg) + proposerCost = 4 + (bdth - 1) + maxCalls = maxLmCalls (budget cfg) + maxCands = maxCandidates (budget cfg) + + goRound r prog calls evald + | r > dpth = pure (setBest prog evald, calls) + | otherwise = do + let cur = fromMaybe "" (instructionAt idx prog) + hist = [PastInstruction i s | (i, s) <- evald] + (cands, calls1) <- + if calls + proposerCost <= maxCalls + then do + ProposeResult cs <- + proposeInstructions + train + ProposeRequest + { program = prog, + targetNode = idx, + currentInstruction = cur, + history = hist, + bootstrappedDemos = [], + numCandidates = bdth - 1, + tipIndex = 1, + viewBatch = 2 + } + pure (cs, calls + proposerCost) + else pure ([cur], calls) + (evald', calls2) <- scoreNew prog calls1 evald (dedupNew evald cands) + -- set the node to the best instruction found so far, then continue + let prog' = setBest prog evald' + goRound (r + 1) prog' calls2 evald' + + -- Score the not-yet-evaluated candidates, threading the call count and stopping + -- before either Budget ceiling would be exceeded. + scoreNew _ calls evald [] = pure (evald, calls) + scoreNew prog calls evald (c : cs) + | calls + dsSize > maxCalls = pure (evald, calls) + | length evald >= maxCands = pure (evald, calls) + | otherwise = do + s <- scoreOn train metric (setNodeInstr idx c prog) + scoreNew prog (calls + dsSize) (evald ++ [(c, s)]) cs + + -- Candidates not already scored, de-duplicated against each other (order-preserving). + dedupNew evald = go (map fst evald) + where + go _ [] = [] + go seen (x : xs) + | x `elem` seen = go seen xs + | otherwise = x : go (x : seen) xs + + -- Set node idx to the highest-scoring instruction recorded (earliest on ties). + setBest prog evald = case evald of + [] -> prog + (e : es) -> setNodeInstr idx (fst (foldl' pick e es)) prog + where + pick best c = if snd c > snd best then c else best diff --git a/src/Shikumi/Optimize/Ensemble.hs b/src/Shikumi/Optimize/Ensemble.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Ensemble.hs @@ -0,0 +1,68 @@ +-- | M4 — ensemble search. Run an inner optimizer several times on different +-- /bootstrap resamples/ of the training set (sampling with replacement, the +-- classic bagging trick that makes the resulting candidates complementary), +-- collect the candidate programs, and combine them into one program with EP-5's +-- 'Shikumi.Combinator.ensemble' combinator under a majority-vote reducer. The +-- ensemble runs all members on the input and returns their modal answer, which is +-- correct more often than any single member when the members err on different +-- inputs. +-- +-- Resampling is deterministic (a fixed linear-congruential stream seeded by the +-- member index), so the search is reproducible run to run. +module Shikumi.Optimize.Ensemble + ( ensembleSearch, + majorityReducer, + ) +where + +import Data.Vector qualified as V +import Shikumi.Combinator (ensemble) +import Shikumi.Compile.Types (compiledProgram) +import Shikumi.Eval (Dataset, dataset, datasetExamples) +import Shikumi.Optimize.Search (freezeProgram) +import Shikumi.Optimize.Types (Optimizer (..)) + +-- | Build an @size@-member ensemble: run @inner@ on @size@ bootstrap resamples of +-- the training set and combine the resulting programs by majority vote. +ensembleSearch :: (Eq o) => Int -> Optimizer i o -> Optimizer i o +ensembleSearch size inner = Optimizer $ \train metric student -> do + let seeds = [1 .. max 1 size] + members <- + mapM + (\seed -> compiledProgram <$> runOptimizer inner (resample seed train) metric student) + seeds + pure (freezeProgram (ensemble members majorityReducer)) + +-- | Sample a dataset with replacement, deterministically, seeded by @seed@. An +-- empty dataset is returned unchanged. +resample :: Int -> Dataset i o -> Dataset i o +resample seed ds = + let exs = V.fromList (datasetExamples ds) + m = V.length exs + in if m == 0 + then ds + else dataset [exs V.! (i `mod` m) | i <- take m (lcg seed)] + +-- | A simple linear-congruential pseudo-random stream (glibc constants). Pure and +-- deterministic — there is no @Math.random@ in this environment, and reproducible +-- sampling is what the test relies on. +lcg :: Int -> [Int] +lcg s0 = drop 1 (iterate step s0) + where + step s = (1103515245 * s + 12345) `mod` 2147483648 + +-- | The modal value of a non-empty list under 'Eq': most frequent, ties broken by +-- first appearance. This is the ensemble's default vote. (EP-4's @modal@ is not +-- exported, so it is reproduced here.) +majorityReducer :: (Eq o) => [o] -> o +majorityReducer items = case counts of + [] -> error "Shikumi.Optimize.Ensemble.majorityReducer: empty member list" + (c : cs) -> fst (foldl' (\best cur -> if snd cur > snd best then cur else best) c cs) + where + counts = [(v, length (filter (== v) items)) | v <- distinct items] + distinct = go [] + where + go _ [] = [] + go seen (y : ys) + | any (== y) seen = go seen ys + | otherwise = y : go (seen ++ [y]) ys diff --git a/src/Shikumi/Optimize/GEPA.hs b/src/Shikumi/Optimize/GEPA.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/GEPA.hs @@ -0,0 +1,273 @@ +{-# 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). +module Shikumi.Optimize.GEPA + ( FeedbackMetric, + ReflectIn (..), + ReflectOut (..), + reflectiveProposer, + captureFeedback, + mutateNode, + gepa, + ) +where + +import Control.Monad (forM, forM_, when) +import Data.Maybe (fromMaybe) +import Data.Text (Text) +import Data.Text qualified as T +import Effectful (Eff, (:>)) +import Effectful.Error.Static (Error) +import Effectful.Prim (Prim) +import GHC.Generics (Generic) +import Shikumi.Adapter (ToPrompt) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval + ( Dataset, + Example (..), + ExampleResult (..), + Prediction, + Report (..), + Score, + datasetExamples, + datasetSize, + evaluatePure, + prediction, + unScore, + ) +import Shikumi.LLM (LLM) +import Shikumi.Module (predict) +import Shikumi.Optimize.Pareto (Candidate (..), paretoFrontier, sampleParent) +import Shikumi.Optimize.Search (freezeProgram) +import Shikumi.Optimize.Types (Budget (..), Optimizer (..)) +import Shikumi.Program + ( NodeFields (..), + Params (..), + Program, + foldParams, + mapParamsAt, + nodeFieldsIndexed, + runProgram, + setProgramParams, + ) +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) + +-- --------------------------------------------------------------------------- +-- The reflective proposer +-- --------------------------------------------------------------------------- + +-- | The reflective proposer's input: the node's current instruction, its accumulated +-- critiques, program/dataset summaries, and the node's field names. +data ReflectIn = ReflectIn + { currentInstruction :: Text, + feedback :: Text, + programSummary :: Text, + datasetSummary :: Text, + fieldSummary :: Text + } + deriving stock (Generic, Show) + +instance FromModel ReflectIn + +instance ToPrompt ReflectIn + +newtype ReflectOut = ReflectOut {proposedInstruction :: Text} + deriving stock (Generic, Show) + +instance ToSchema ReflectOut + +instance FromModel ReflectOut + +instance ToPrompt ReflectOut + +instance Validatable ReflectOut + +-- | The default reflective proposer: a single predict node that addresses the +-- feedback specifically. +reflectiveProposer :: Program ReflectIn ReflectOut +reflectiveProposer = + predict + ( mkSignature + "You are improving the instruction for one node of a language-model pipeline. You are \ + \given the node's current instruction, textual feedback describing how it failed on \ + \several examples, a summary of the whole program, a summary of the dataset, and the \ + \node's input/output field names. Write a single improved instruction that addresses the \ + \feedback specifically, in the `proposedInstruction` field." + ) + +-- --------------------------------------------------------------------------- +-- 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). +captureFeedback :: + (LLM :> es, Error ShikumiError :> es, Prim :> es) => + 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) + +-- --------------------------------------------------------------------------- +-- M2 — reflective mutation +-- --------------------------------------------------------------------------- + +-- | Reflect on node @idx@'s accumulated feedback and overwrite its instruction with +-- the proposal. A node with no feedback is left unchanged (nothing to reflect on). +mutateNode :: + (LLM :> es, Error ShikumiError :> es) => + Program ReflectIn ReflectOut -> + -- | program summary + Text -> + -- | dataset summary + Text -> + [NodeFields] -> + FeedbackLog -> + [NodePath] -> + Int -> + Program i o -> + Eff es (Program i o) +mutateNode proposer progSummary dataSummary fields fblog paths idx prog = + case drop idx paths of + [] -> pure prog + (path : _) -> + let crits = feedbackFor path fblog + in if null crits + then pure prog + else do + let cur = fromMaybe "" (instructionOverride (paramsAt idx prog)) + fldSummary = renderFields (drop idx fields) + fb = T.intercalate "\n" crits + ReflectOut newInstr <- + runProgram proposer (ReflectIn cur fb progSummary dataSummary fldSummary) + pure (mapParamsAt idx (\ps -> ps {instructionOverride = Just newInstr}) prog) + +-- | The 'Params' at a node index (empty if out of range). +paramsAt :: Int -> Program i o -> Params +paramsAt idx prog = case drop idx (foldParams prog) of + (ps : _) -> ps + [] -> Params Nothing [] + +-- | Render a node's field names for the proposer prompt. +renderFields :: [NodeFields] -> Text +renderFields [] = "inputs: ?; outputs: ?" +renderFields (NodeFields ins outs : _) = + "inputs: " <> commas ins <> "; outputs: " <> commas outs + where + commas xs = if null xs then "?" else T.intercalate ", " xs + +-- --------------------------------------------------------------------------- +-- M3 — the evolution loop +-- --------------------------------------------------------------------------- + +-- | The reflective evolutionary optimizer. Takes its reflective proposer and feedback +-- metric explicitly (so it is testable under a stub LM) and returns V1's +-- 'Optimizer'. +gepa :: + Program ReflectIn ReflectOut -> + FeedbackMetric o -> + Budget -> + Optimizer i o +gepa proposer fbMetric budget = Optimizer $ \train metric student -> do + let paths = programNodePaths student + fields = nodeFieldsIndexed student + nNodes = max 1 (length paths) + n = max 1 (datasetSize train) + progSummary = fallbackProgramSummary (length paths) + dataSummary = fallbackDatasetSummary (datasetSize train) + maxCalls = maxLmCalls budget + maxCands = maxCandidates budget + rebuild cand = either (const student) id (setProgramParams (candParams cand) student) + + seedRpt <- evaluatePure train metric student + let seedCand = Candidate (foldParams student) (perEx seedRpt) (aggregateScore seedRpt) + + -- A full step costs: capture (n) + mutate (1) + child eval (n) = 2n + 1. + stepCost = 2 * n + 1 + stepCap = maxCands + 4 + + loop step calls cands seed frontier + | step >= stepCap = pure (bestOf seedCand frontier) + | calls + stepCost > maxCalls = pure (bestOf seedCand frontier) + | length cands >= maxCands = pure (bestOf seedCand frontier) + | otherwise = case sampleParent seed (paretoFrontier frontier) of + Nothing -> pure (bestOf seedCand frontier) + Just (parent, seed') -> do + let parentProg = rebuild parent + idx = step `mod` nNodes + (fblog, _) <- captureFeedback train fbMetric parentProg + case drop idx paths of + (path : _) + | null (feedbackFor path fblog) -> + -- nothing to reflect on at this node; advance (capture cost only) + loop (step + 1) (calls + n) cands seed' frontier + _ -> do + child <- mutateNode proposer progSummary dataSummary fields fblog paths idx parentProg + rpt <- evaluatePure train metric child + let childCand = Candidate (foldParams child) (perEx rpt) (aggregateScore rpt) + frontier' = paretoFrontier (childCand : frontier) + loop (step + 1) (calls + stepCost) (childCand : cands) seed' frontier' + + best <- loop 0 n [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] + +-- | The frontier candidate with the highest aggregate (earliest on ties); falls back +-- to the seed if the frontier is somehow empty. +bestOf :: Candidate -> [Candidate] -> Candidate +bestOf seedCand cs = case cs of + [] -> seedCand + (x : xs) -> foldl' (\b c -> if candAggregate c > candAggregate b then c else b) x xs + +-- | A minimal program summary (EP-19's program describer is the richer source). +fallbackProgramSummary :: Int -> Text +fallbackProgramSummary k = "A language-model program with " <> tshow k <> " predict node(s)." + +-- | A minimal dataset summary (EP-19's dataset summarizer is the richer source). +fallbackDatasetSummary :: Int -> Text +fallbackDatasetSummary k = "A dataset of " <> tshow k <> " example(s)." + +tshow :: (Show a) => a -> Text +tshow = T.pack . show diff --git a/src/Shikumi/Optimize/Instruction.hs b/src/Shikumi/Optimize/Instruction.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Instruction.hs @@ -0,0 +1,92 @@ +-- | M3 — instruction search (a MIPRO/COPRO-style optimizer). For each node, the +-- /grounded/ proposer (EP-19, 'Shikumi.Optimize.Propose.proposeInstructions') suggests +-- several candidate instruction strings — fed a dataset summary, the program's +-- pseudo-code and description, the node's real input/output field names, the +-- instruction history, and a stylistic tip — and each candidate is scored; the best is +-- kept. Optimization is greedy coordinate ascent — one node at a time, holding the +-- others fixed — so the candidate count is linear in @nodes × proposals@. +-- +-- The proposer and its signal-gatherers are themselves ordinary shikumi 'Program's, so +-- they are typed, cached, traced, and testable with the same stub-LM machinery as +-- everything else — the optimizer is written in the framework it optimizes. The +-- /current/ instruction is always retained as a candidate (the proposer guarantees it), +-- so a node can never end up worse than where it started. +-- +-- __Budget.__ The grounded proposer makes @4 + proposalsPerNode@ LM calls per node +-- (dataset summary, program describe, module describe, and one generation per proposal); +-- scoring one candidate costs one LM call per dataset example. The search threads a +-- running raw-call count and stops — returning the best found /so far/ — before either +-- bound in the 'Budget' would be exceeded, so the recorded LM-call count never exceeds +-- @maxLmCalls@. When the remaining budget cannot cover a node's full proposal, that node +-- keeps its current instruction (no proposer call) rather than partially proposing. +-- +-- This module re-points V1's blind proposer at EP-19's grounded surface; the old +-- @ProposeIn@/@ProposeOut@/@proposeInstruction@ predictor is removed (the grounded +-- @GenerateInstructionIn@/@GenerateInstructionOut@ replaces it, still emitting a +-- @proposedInstruction@ output field). +module Shikumi.Optimize.Instruction + ( instructionSearch, + ) +where + +import Control.Monad (foldM) +import Data.Aeson (ToJSON) +import Data.Maybe (fromMaybe) +import Shikumi.Eval (datasetSize) +import Shikumi.Optimize.Propose + ( ProposeRequest (..), + ProposeResult (..), + proposeInstructions, + ) +import Shikumi.Optimize.Search (freezeProgram, instructionAt, scoreOn, setNodeInstr) +import Shikumi.Optimize.Types (Budget (..), Optimizer (..)) +import Shikumi.Program (foldParams) + +-- | Search for a better instruction at every node by greedy coordinate ascent under +-- an explicit LM-call budget, using the grounded proposer to generate candidates. +instructionSearch :: (ToJSON i, ToJSON o) => Int -> Budget -> Optimizer i o +instructionSearch proposalsPerNode budget = Optimizer $ \train metric student -> do + let dsSize = datasetSize train + nNodes = length (foldParams student) + -- The grounded proposer's fixed per-node LM-call cost (see module header). + proposerCost = 4 + max 0 proposalsPerNode + + -- Score candidate instructions for node @idx@, threading the call count and + -- the best-so-far; stop before the next scoring would exceed the budget. + scoreCands calls best idx prog cs = case cs of + [] -> pure (best, calls) + (c : rest) + | calls + dsSize > maxLmCalls budget -> pure (best, calls) + | otherwise -> do + s <- scoreOn train metric (setNodeInstr idx c prog) + let best' = case best of + Nothing -> Just (c, s) + Just (_, bs) -> if s > bs then Just (c, s) else best + scoreCands (calls + dsSize) best' idx prog rest + + -- Optimize one node, holding the others fixed. + stepNode (prog, calls) idx = do + let curInstr = fromMaybe "" (instructionAt idx prog) + (cands, calls1) <- + if calls + proposerCost <= maxLmCalls budget + then do + ProposeResult cs <- + proposeInstructions + train + ProposeRequest + { program = prog, + targetNode = idx, + currentInstruction = curInstr, + history = [], + bootstrappedDemos = [], + numCandidates = proposalsPerNode, + tipIndex = 0, + viewBatch = 2 + } + pure (cs, calls + proposerCost) + else pure ([curInstr], calls) + (best, calls2) <- scoreCands calls1 Nothing idx prog cands + pure (setNodeInstr idx (maybe curInstr fst best) prog, calls2) + + (final, _) <- foldM stepNode (student, 0) [0 .. nNodes - 1] + pure (freezeProgram final) diff --git a/src/Shikumi/Optimize/KNN.hs b/src/Shikumi/Optimize/KNN.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/KNN.hs @@ -0,0 +1,124 @@ +-- | KNN few-shot demo selection (EP-23): pick, for each input, the training +-- examples whose inputs are most /semantically similar/ to it, and show those as the +-- demos. Two forms: +-- +-- * 'knnFewShot' — the faithful run-time form: a single 'Embed' node that, per +-- input, embeds the input, ranks the training examples by cosine similarity, and +-- runs the student under the @k@ nearest as demos. Demos depend on the input. +-- * 'knnFewShotCentroid' — a compile-time fallback that bakes the @k@ examples +-- nearest the training centroid as a fixed demo set (a plain @Params@ artifact), +-- for callers who cannot run an embedder at execution time. +-- +-- The embedder is injected as a /pure/ closure @Text -> Vector Double@ (the shape of +-- EP-15's pure @runEmbedding@ argument), not the @Embedding@ effect: an 'Embed' +-- body's row is fixed to @(LLM, Error ShikumiError)@, so it cannot call @embedText@; +-- all embedding-effect work happens at the caller, outside the node. The run-time +-- form carries no @Params@ (it serializes as the empty vector, exactly like @react@). +module Shikumi.Optimize.KNN + ( -- * Run-time form (the faithful analog) + knnFewShot, + knnDemos, + + -- * Compile-time fallback + knnFewShotCentroid, + + -- * Pure helpers (exposed so callers/tests can reproduce the ranking) + cosine, + nearestDemos, + centroid, + ) +where + +import Data.Aeson (ToJSON, toJSON) +import Data.List (sortBy) +import Data.Ord (Down (..), comparing) +import Data.Text (Text) +import Data.Vector (Vector) +import Data.Vector qualified as V +import Shikumi.Adapter (ToPrompt, toPrompt) +import Shikumi.Eval (Dataset, Example (..), datasetExamples) +import Shikumi.Optimize.LabeledFewShot (withDemos) +import Shikumi.Optimize.Search (freezeProgram) +import Shikumi.Optimize.Types (Optimizer (..)) +import Shikumi.Program (Demo (..), Program, embed, runProgram) + +-- | Cosine similarity of two vectors: @dot a b / (‖a‖·‖b‖)@, in @[-1, 1]@. A +-- zero-norm vector yields @0@ (no signal). Ranking by cosine puts the most similar +-- examples first. +cosine :: Vector Double -> Vector Double -> Double +cosine a b + | na == 0 || nb == 0 = 0 + | otherwise = dot a b / (na * nb) + where + dot x y = V.sum (V.zipWith (*) x y) + na = sqrt (dot a a) + nb = sqrt (dot b b) + +-- | The @k@ training examples whose input embeddings are nearest a query vector, +-- nearest first, as JSON demos. Ties break by original order (a stable sort). +nearestDemosByVec :: + (ToJSON i, ToJSON o, ToPrompt i) => + (Text -> Vector Double) -> + Int -> + [Example i o] -> + Vector Double -> + [Demo] +nearestDemosByVec embedder k exs qv = + [Demo (toJSON i) (toJSON o) | (_, _, Example i o) <- take (max 0 k) sorted] + where + scored = [(cosine (embedder (toPrompt i)) qv, idx, ex) | (idx, ex@(Example i _)) <- zip [0 :: Int ..] exs] + sorted = sortBy (comparing (\(s, idx, _) -> (Down s, idx))) scored + +-- | The @k@ training examples nearest a rendered query input, nearest first. +nearestDemos :: + (ToJSON i, ToJSON o, ToPrompt i) => + (Text -> Vector Double) -> + Int -> + [Example i o] -> + Text -> + [Demo] +nearestDemos embedder k exs query = nearestDemosByVec embedder k exs (embedder query) + +-- | The component-wise mean of a list of equal-length vectors (empty list → empty +-- vector). +centroid :: [Vector Double] -> Vector Double +centroid [] = V.empty +centroid vs = V.map (/ fromIntegral (length vs)) (foldl1 (V.zipWith (+)) vs) + +-- | The run-time KNN node: for each input, attach the @k@ nearest training examples +-- as demos and run the student under them. A plain @Program i o@ (an 'Embed' node) +-- usable anywhere a Program is; it carries no @Params@ (like @react@). +knnDemos :: + (ToJSON i, ToJSON o, ToPrompt i) => + (Text -> Vector Double) -> + Int -> + Dataset i o -> + Program i o -> + Program i o +knnDemos embedder k train student = + let exs = datasetExamples train + in embed $ \i -> runProgram (withDemos (nearestDemos embedder k exs (toPrompt i)) student) i + +-- | Run-time KNN as an 'Optimizer': selection is by embedding geometry, not by +-- score, so it consults neither the metric nor the LM at optimize time. +knnFewShot :: + (ToJSON i, ToJSON o, ToPrompt i) => + (Text -> Vector Double) -> + Int -> + Optimizer i o +knnFewShot embedder k = Optimizer $ \train _metric student -> + pure (freezeProgram (knnDemos embedder k train student)) + +-- | Compile-time fallback: bake the @k@ training examples nearest the training +-- centroid as a fixed demo set on every node, and freeze. A plain @Params@ artifact, +-- no run-time embedder needed. +knnFewShotCentroid :: + (ToJSON i, ToJSON o, ToPrompt i) => + (Text -> Vector Double) -> + Int -> + Optimizer i o +knnFewShotCentroid embedder k = Optimizer $ \train _metric student -> + let exs = datasetExamples train + c = centroid [embedder (toPrompt i) | Example i _ <- exs] + ds = nearestDemosByVec embedder k exs c + in pure (freezeProgram (withDemos ds student)) diff --git a/src/Shikumi/Optimize/LabeledFewShot.hs b/src/Shikumi/Optimize/LabeledFewShot.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/LabeledFewShot.hs @@ -0,0 +1,54 @@ +-- | M1 — labeled few-shot demo selection, the simplest real optimizer (DSPy's +-- @LabeledFewShot@). Each training example /is/ a candidate demonstration; the +-- search forms several size-@k@ demo sets, scores each by attaching it to every +-- node and evaluating over the training set, and keeps the highest-scoring set. +-- +-- No LM calls are made beyond scoring. Candidate sets are enumerated +-- deterministically (all size-@k@ combinations of the training demos, in a fixed +-- order, bounded by the budget) rather than randomly sampled, so the result is +-- reproducible run to run — the tests rely on this. +module Shikumi.Optimize.LabeledFewShot + ( labeledFewShot, + labeledCandidateSets, + withDemos, + ) +where + +import Data.Aeson (ToJSON, toJSON) +import Shikumi.Eval (Dataset, Example (..), datasetExamples) +import Shikumi.Optimize.Search (freezeProgram, scoreOn, selectBest) +import Shikumi.Optimize.Types (Optimizer (..), Scored (..), defaultBudget) +import Shikumi.Program (Demo (..), Params (..), Program, mapParams) + +-- | Select the best size-@k@ set of labelled demonstrations from the training set. +-- Multi-node programs receive the same demo set at every node (the DSPy default). +labeledFewShot :: (ToJSON i, ToJSON o) => Int -> Optimizer i o +labeledFewShot k = Optimizer $ \train metric prog -> do + let sets = labeledCandidateSets k train + best <- selectBest defaultBudget (\ds -> scoreOn train metric (withDemos ds prog)) sets + pure $ case best of + Nothing -> freezeProgram prog + Just sc -> freezeProgram (withDemos (candidate sc) prog) + +-- | The candidate demo sets a 'labeledFewShot' search considers: every size-@k@ +-- combination of the training examples (each turned into a JSON 'Demo'), in +-- deterministic enumeration order. Exposed so tests can reproduce the exact set of +-- candidates the optimizer scored. +labeledCandidateSets :: (ToJSON i, ToJSON o) => Int -> Dataset i o -> [[Demo]] +labeledCandidateSets k train = combinations k allDemos + where + allDemos = + [ Demo {input = toJSON i, output = toJSON o} + | Example i o <- datasetExamples train + ] + +-- | Attach a demo set to every node, leaving each node's instruction untouched. +withDemos :: [Demo] -> Program i o -> Program i o +withDemos ds = mapParams (\ps -> ps {demos = ds}) + +-- | All size-@k@ sub-lists of a list, preserving relative order +-- (@combinations 2 [a,b,c] = [[a,b],[a,c],[b,c]]@). +combinations :: Int -> [a] -> [[a]] +combinations k _ | k <= 0 = [[]] +combinations _ [] = [] +combinations k (x : xs) = map (x :) (combinations (k - 1) xs) ++ combinations k xs diff --git a/src/Shikumi/Optimize/MIPRO.hs b/src/Shikumi/Optimize/MIPRO.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/MIPRO.hs @@ -0,0 +1,314 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +-- | MIPROv2 (EP-20): the joint instruction-and-demonstration optimizer. Where +-- 'Shikumi.Optimize.Instruction.instructionSearch' improves one axis (instructions) +-- one node at a time, MIPROv2 searches /both/ axes — instruction × demo set — across +-- all nodes jointly, in three phases: +-- +-- 1. __bootstrap__ candidate demo sets (per node) from the teacher's passing runs +-- plus the labelled training pairs ('bootstrapDemoCandidates'); +-- 2. __propose__ candidate instructions (per node) through EP-19's grounded +-- proposer ('proposeInstructionCandidates'); and +-- 3. __search__ the joint @(instruction × demoset)@ grid with cheap minibatch +-- trials gated by periodic full evaluation ('searchJoint'). +-- +-- The result is V1's 'Shikumi.Compile.Types.CompiledProgram', produced via +-- 'freezeProgram', invoked through V1's 'Shikumi.Optimize.optimize' and persisted +-- with V1's serialization — no parallel optimizer or persistence surface (MasterPlan +-- integration point #4). +-- +-- __Search surrogate.__ DSPy drives phase 3 with Optuna's TPE sampler, which has no +-- Haskell equivalent in this workspace. We implement __greedy coordinate descent with +-- minibatch pruning__ over the joint grid: each trial screens the one-coordinate +-- neighbours of the running best on a seeded minibatch, then full-evaluates the +-- best-screened neighbour and accepts it only if it strictly improves the full score. +-- This keeps the three essential, testable behaviours — joint grid, minibatch-screen +-- + full-eval-confirm, and a hard 'Budget' — while staying fully deterministic. A +-- later EP can swap the neighbour-selection step for a TPE-lite without changing this +-- module's public surface. +module Shikumi.Optimize.MIPRO + ( -- * Public surface + Miprov2Auto (..), + Miprov2Config (..), + miprov2Auto, + miprov2, + miprov2With, + + -- * Phases (exported for tests) + bootstrapDemoCandidates, + proposeInstructionCandidates, + searchJoint, + ) +where + +import Control.Monad (forM) +import Data.Aeson (ToJSON) +import Data.Aeson.Text (encodeToLazyText) +import Data.Maybe (fromMaybe) +import Data.Text (Text) +import Data.Text.Lazy qualified as TL +import Effectful (Eff, (:>)) +import Effectful.Concurrent (Concurrent) +import Effectful.Error.Static (Error, catchError) +import Effectful.Prim (Prim) +import Shikumi.Effect.Time (Time) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval + ( Dataset, + Example (..), + Metric, + dataset, + datasetExamples, + datasetSize, + prediction, + unScore, + ) +import Shikumi.LLM (LLM) +import Shikumi.Optimize.Bootstrap (recoverDemo) +import Shikumi.Optimize.Propose + ( ProposeRequest (..), + ProposeResult (..), + proposeInstructions, + ) +import Shikumi.Optimize.Search (freezeProgram, scoreOn) +import Shikumi.Optimize.Types (Budget (..), Optimizer (..), defaultBudget) +import Shikumi.Program + ( Demo (..), + Params (..), + Program, + foldParams, + mapParamsAt, + runProgram, + ) + +-- --------------------------------------------------------------------------- +-- Configuration and presets +-- --------------------------------------------------------------------------- + +-- | How aggressively to search. Mirrors DSPy's light/medium/heavy "auto" modes. +data Miprov2Auto = Miprov2Light | Miprov2Medium | Miprov2Heavy + deriving stock (Eq, Show) + +-- | The full, explicit configuration. Presets fill it in; advanced users build one +-- directly. +data Miprov2Config = Miprov2Config + { -- | instruction candidates proposed per node (incl. the current at index 0) + numInstructCandidates :: !Int, + -- | demo-set candidates per node (incl. the empty set at index 0) + numDemoCandidates :: !Int, + -- | minibatch trials the search performs + numTrials :: !Int, + -- | examples per minibatch screening evaluation + minibatchSize :: !Int, + -- | run a full eval of a trial's proposal every k trials (also retained for + -- config compatibility; the implementation full-evaluates each trial's single + -- proposal to gate acceptance — see the module header) + fullEvalEvery :: !Int, + -- | cap on demos in one bootstrapped set (prompt-size guard) + maxBootstrappedDemos :: !Int, + -- | min metric score for a teacher run to contribute demos + bootstrapThreshold :: !Double, + -- | hard LM-call / candidate ceiling (V1's 'Budget') + budget :: !Budget + } + deriving stock (Eq, Show) + +-- | Turn a preset into a concrete config. Values are small enough for hermetic tests +-- yet keep DSPy's monotonic ordering (heavier => more candidates, more trials). +miprov2Auto :: Miprov2Auto -> Miprov2Config +miprov2Auto Miprov2Light = Miprov2Config 2 2 6 3 3 4 1.0 defaultBudget +miprov2Auto Miprov2Medium = Miprov2Config 3 3 12 5 4 4 1.0 defaultBudget +miprov2Auto Miprov2Heavy = Miprov2Config 4 4 18 8 5 4 1.0 defaultBudget + +-- --------------------------------------------------------------------------- +-- Public optimizers +-- --------------------------------------------------------------------------- + +-- | MIPROv2 with a preset. The teacher (whose successful runs seed bootstrapped +-- demos) defaults to the student itself, matching DSPy's default. +miprov2 :: (ToJSON i, ToJSON o) => Miprov2Auto -> Optimizer i o +miprov2 a = Optimizer $ \train metric student -> + runOptimizer (miprov2With (miprov2Auto a) student) train metric student + +-- | MIPROv2 with an explicit config and an explicit teacher program. +miprov2With :: (ToJSON i, ToJSON o) => Miprov2Config -> Program i o -> Optimizer i o +miprov2With cfg teacher = Optimizer $ \train metric student -> do + demoCands <- bootstrapDemoCandidates cfg teacher train metric student + instrCands <- proposeInstructionCandidates cfg student train demoCands + best <- searchJoint cfg train metric student instrCands demoCands + pure (freezeProgram best) + +-- --------------------------------------------------------------------------- +-- Phase 1 — bootstrap demo-set candidates +-- --------------------------------------------------------------------------- + +-- | For each node (in @foldParams@ order), a list of candidate demo sets. Candidate +-- 0 is always the empty set (so "no demos" is always reachable). The remaining sets +-- come from the teacher's metric-passing runs (bootstrapped) and from the labelled +-- training pairs (labelled demos, DSPy's @max_labeled_demos@) — recovered at the +-- program-I/O level and attached to every node (the documented degradation while +-- per-node trace recovery via EP-16 is deferred; for single-node programs the two +-- coincide). +bootstrapDemoCandidates :: + (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es) => + Miprov2Config -> + -- | teacher + Program i o -> + Dataset i o -> + Metric o -> + -- | student (for the node count) + Program i o -> + Eff es [[[Demo]]] +bootstrapDemoCandidates cfg teacher train metric student = do + let exs = datasetExamples train + 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 [] + bootstrapped <- concat <$> mapM keepIfPassing exs + let cap = max 1 (maxBootstrappedDemos cfg) + labeledSet = take cap (map (\(Example i o) -> recoverDemo i o) exs) + bootSet = take cap bootstrapped + sets = take (max 1 (numDemoCandidates cfg)) ([] : filter (not . null) [labeledSet, bootSet]) + nNodes = length (foldParams student) + pure (replicate nNodes sets) + +-- --------------------------------------------------------------------------- +-- Phase 2 — propose instruction candidates +-- --------------------------------------------------------------------------- + +-- | For each node (in @foldParams@ order), a list of candidate instructions with the +-- node's current instruction first (the grounded proposer guarantees retention). +-- Uses EP-19's 'proposeInstructions'; the per-candidate tip starts at the /creative/ +-- tip (@tipIndex = 1@) so a node always gets a strong proposal even at small N. +proposeInstructionCandidates :: + (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es) => + Miprov2Config -> + Program i o -> + Dataset i o -> + [[[Demo]]] -> + Eff es [[Text]] +proposeInstructionCandidates cfg student train demoCands = + forM (zip3 [0 ..] (foldParams student) demoCands) $ \(k, ps, nodeDemoSets) -> do + let curInstr = fromMaybe "" (instructionOverride ps) + demoTexts = map renderDemo (concat (take 1 (filter (not . null) nodeDemoSets))) + ProposeResult cs <- + proposeInstructions + train + ProposeRequest + { program = student, + targetNode = k, + currentInstruction = curInstr, + history = [], + bootstrappedDemos = demoTexts, + numCandidates = max 0 (numInstructCandidates cfg - 1), + tipIndex = 1, + viewBatch = 2 + } + pure cs + +-- | Render a recovered 'Demo' as @ => @ for the proposal +-- prompt's demo signal. +renderDemo :: Demo -> Text +renderDemo (Demo i o) = enc i <> " => " <> enc o + where + enc = TL.toStrict . encodeToLazyText + +-- --------------------------------------------------------------------------- +-- Phase 3 — the joint minibatch search +-- --------------------------------------------------------------------------- + +-- | A per-node selection: @(instructionIndex, demoSetIndex)@. +type JointVec = [(Int, Int)] + +-- | Search the joint per-node @(instruction × demoset)@ grid by greedy coordinate +-- descent with minibatch screening, returning the best program found within the +-- 'Budget'. Each trial screens the one-coordinate neighbours of the running best on a +-- seeded minibatch, then full-evaluates the best-screened neighbour and accepts it +-- only if it strictly improves the full score. +searchJoint :: + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + Miprov2Config -> + Dataset i o -> + Metric o -> + Program i o -> + -- | per-node instruction candidates + [[Text]] -> + -- | per-node demo-set candidates + [[[Demo]]] -> + Eff es (Program i o) +searchJoint cfg train metric student instrCands demoCands = do + let nNodes = length (foldParams student) + dsSize = datasetSize train + mbSize = max 1 (min dsSize (minibatchSize cfg)) + maxCalls = maxLmCalls (budget cfg) + apply = applyVec instrCands demoCands student + baseVec = replicate nNodes (0, 0) + if dsSize > maxCalls + then pure (apply baseVec) + else do + baseFull <- scoreOn train metric (apply baseVec) + let screen _t calls acc [] = pure (reverse acc, calls) + screen t calls acc (v : vs) + | calls + mbSize > maxCalls = pure (reverse acc, calls) + | otherwise = do + s <- scoreOn (minibatchAt mbSize (t * 7 + length acc) train) metric (apply v) + screen t (calls + mbSize) ((v, s) : acc) vs + go t calls bestVec bestFull + | t >= numTrials cfg = pure bestVec + | otherwise = do + (scored, calls1) <- screen t calls [] (oneCoordNeighbors instrCands demoCands bestVec) + case bestByScore scored of + Nothing -> pure bestVec + Just propVec + | calls1 + dsSize > maxCalls -> pure bestVec + | otherwise -> do + fs <- scoreOn train metric (apply propVec) + if fs > bestFull + then go (t + 1) (calls1 + dsSize) propVec fs + else go (t + 1) (calls1 + dsSize) bestVec bestFull + best <- go 0 dsSize baseVec baseFull + pure (apply best) + +-- | Apply a joint vector to the student: set each node's instruction and demos to the +-- selected candidates. +applyVec :: [[Text]] -> [[[Demo]]] -> Program i o -> JointVec -> Program i o +applyVec instrCands demoCands prog vec = foldl step prog (zip [0 ..] vec) + where + step p (k, (i, d)) = + mapParamsAt + k + ( \ps -> + ps + { instructionOverride = Just (instrCands !! k !! i), + demos = demoCands !! k !! d + } + ) + p + +-- | The one-coordinate neighbours of a vector: each move changes exactly one node's +-- instruction or demo index to a non-current alternative. +oneCoordNeighbors :: [[Text]] -> [[[Demo]]] -> JointVec -> [JointVec] +oneCoordNeighbors instrCands demoCands vec = + [setAt k (i, snd (vec !! k)) vec | k <- nodes, i <- idxs (instrCands !! k), i /= fst (vec !! k)] + ++ [setAt k (fst (vec !! k), d) vec | k <- nodes, d <- idxs (demoCands !! k), d /= snd (vec !! k)] + where + nodes = [0 .. length vec - 1] + idxs xs = [0 .. length xs - 1] + setAt k v xs = take k xs ++ [v] ++ drop (k + 1) xs + +-- | The vector with the strictly-highest score (earliest on ties); 'Nothing' if none. +bestByScore :: [(JointVec, Double)] -> Maybe JointVec +bestByScore [] = Nothing +bestByScore (x : xs) = Just (fst (foldl (\b c -> if snd c > snd b then c else b) x xs)) + +-- | A deterministic minibatch of @sz@ examples, rotated by @seed@ (a seeded sample is +-- as reproducible as the whole set, and far cheaper to screen on). +minibatchAt :: Int -> Int -> Dataset i o -> Dataset i o +minibatchAt sz seed ds = + let exs = datasetExamples ds + n = length exs + in if n == 0 then ds else dataset (take (max 1 sz) (drop (seed `mod` n) (cycle exs))) diff --git a/src/Shikumi/Optimize/Pareto.hs b/src/Shikumi/Optimize/Pareto.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Pareto.hs @@ -0,0 +1,60 @@ +-- | The Pareto-frontier bookkeeping for GEPA (EP-22): a pure, effect-free module so +-- the frontier logic is trivially testable and reproducible. A 'Candidate' is a +-- program identified by its node-parameter vector, carrying its per-example score +-- vector and aggregate. Keeping the /frontier/ (rather than one global best) +-- preserves candidates that win on some examples even if not best on average — the +-- room reflective evolution needs to escape local optima. +module Shikumi.Optimize.Pareto + ( Candidate (..), + dominates, + paretoFrontier, + sampleParent, + ) +where + +import Shikumi.Program (Params) + +-- | A candidate program, identified by its @foldParams@-order node parameters, with +-- its per-example scores (dataset order) and their mean. +data Candidate = Candidate + { candParams :: ![Params], + candPerExample :: ![Double], + candAggregate :: !Double + } + deriving stock (Eq, Show) + +-- | @a@ dominates @b@ when @a >= b@ on every example and @a > b@ on at least one. +-- Equal vectors are NOT domination, so equal candidates coexist on the frontier. +dominates :: Candidate -> Candidate -> Bool +dominates a b = + let pa = candPerExample a + pb = candPerExample b + in length pa == length pb && and (zipWith (>=) pa pb) && or (zipWith (>) pa pb) + +-- | The non-dominated subset, preserving input order (so the frontier is +-- reproducible). A candidate is kept iff no other candidate dominates it. +paretoFrontier :: [Candidate] -> [Candidate] +paretoFrontier cs = [c | c <- cs, not (any (`dominates` c) cs)] + +-- | Deterministically sample a parent from the frontier, biased toward candidates +-- that achieve the best score on more examples (the "pareto" selection strategy). +-- Pure: takes an explicit 'Int' seed and returns the chosen candidate plus the next +-- seed, so the search threads RNG state with no IO. A simple LCG drives the choice. +sampleParent :: Int -> [Candidate] -> Maybe (Candidate, Int) +sampleParent _ [] = Nothing +sampleParent seed cands@(c0 : _) = Just (go r (zip cands weights), seed') + where + cols = length (candPerExample c0) + colMax j = maximum [candPerExample c !! j | c <- cands] + weightOf c = max 1 (length [() | j <- [0 .. cols - 1], (candPerExample c !! j) >= colMax j]) + weights = map weightOf cands + total = max 1 (sum weights) + seed' = lcg seed + r = seed' `mod` total + go _ [] = c0 + go k ((c, w) : rest) = if k < w then c else go (k - w) rest + +-- | A small linear congruential generator (glibc constants), kept in-house so runs +-- are reproducible with no IO entropy. Returns a non-negative pseudo-random 'Int'. +lcg :: Int -> Int +lcg x = (1103515245 * abs x + 12345) `mod` 2147483648 diff --git a/src/Shikumi/Optimize/Propose.hs b/src/Shikumi/Optimize/Propose.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Propose.hs @@ -0,0 +1,42 @@ +-- | The grounded instruction proposer (EP-19), a single import for consumers +-- (instruction search today; MIPROv2 and COPRO next). Re-exports the public surface: +-- the 'proposeInstructions' driver, its request/result records, the per-node +-- field-metadata accessor (integration point #3), the instruction-history vocabulary, +-- and the tip bank. +module Shikumi.Optimize.Propose + ( -- * The proposer + proposeInstructions, + ProposeRequest (..), + ProposeResult (..), + + -- * Per-node field metadata (integration point #3) + NodeFieldNames (..), + programFieldNames, + moduleSignatureAt, + + -- * Instruction history + PastInstruction (..), + renderHistory, + + -- * Tips + tipBank, + tipAt, + + -- * Pseudo-code and dataset summary (reusable by GEPA, EP-22) + renderProgramPseudo, + datasetSummary, + ) +where + +import Shikumi.Optimize.Propose.Grounded (proposeInstructions) +import Shikumi.Optimize.Propose.Summarize (datasetSummary, renderProgramPseudo) +import Shikumi.Optimize.Propose.Tips (tipAt, tipBank) +import Shikumi.Optimize.Propose.Types + ( NodeFieldNames (..), + PastInstruction (..), + ProposeRequest (..), + ProposeResult (..), + moduleSignatureAt, + programFieldNames, + renderHistory, + ) diff --git a/src/Shikumi/Optimize/Propose/Grounded.hs b/src/Shikumi/Optimize/Propose/Grounded.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Propose/Grounded.hs @@ -0,0 +1,141 @@ +{-# LANGUAGE FlexibleContexts #-} + +-- | The assembled grounded proposer (EP-19): the final instruction-generating +-- predictor (DSPy's @GenerateSingleModuleInstruction@) and the 'proposeInstructions' +-- driver that gathers every signal and returns ranked candidate instructions for one +-- node — with the current instruction always retained, so a node can never degrade. +module Shikumi.Optimize.Propose.Grounded + ( GenerateInstructionIn (..), + GenerateInstructionOut (..), + instructionGenerator, + proposeInstructions, + ) +where + +import Control.Monad (forM) +import Data.Aeson (ToJSON) +import Data.Text (Text) +import Data.Text qualified as T +import Effectful (Eff, (:>)) +import Effectful.Error.Static (Error) +import GHC.Generics (Generic) +import Shikumi.Adapter (ToPrompt) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval (Dataset, datasetExamples) +import Shikumi.LLM (LLM) +import Shikumi.Module (predict) +import Shikumi.Optimize.Propose.Summarize + ( ModuleDescribeIn (..), + ModuleDescribeOut (..), + ProgramDescribeIn (..), + ProgramDescribeOut (..), + datasetSummary, + moduleDescriber, + programDescriber, + renderExampleRow, + renderProgramPseudo, + ) +import Shikumi.Optimize.Propose.Tips (tipAt) +import Shikumi.Optimize.Propose.Types + ( ProposeRequest (..), + ProposeResult (..), + moduleSignatureAt, + renderHistory, + ) +import Shikumi.Program (Program, runProgram) +import Shikumi.Schema (FromModel, ToSchema, Validatable) +import Shikumi.Signature (mkSignature) + +-- | The grounded proposer's input: every signal about the optimization target, +-- rendered into the prompt under its field name by the generic 'ToPrompt'. +data GenerateInstructionIn = GenerateInstructionIn + { datasetDescription :: Text, + programCode :: Text, + programDescription :: Text, + moduleSignature :: Text, + moduleDescription :: Text, + taskDemos :: Text, + previousInstructions :: Text, + basicInstruction :: Text, + tip :: Text + } + deriving stock (Generic, Show) + +instance FromModel GenerateInstructionIn + +instance ToPrompt GenerateInstructionIn + +newtype GenerateInstructionOut = GenerateInstructionOut {proposedInstruction :: Text} + deriving stock (Generic, Show) + +instance ToSchema GenerateInstructionOut + +instance FromModel GenerateInstructionOut + +instance ToPrompt GenerateInstructionOut + +instance Validatable GenerateInstructionOut + +instructionGenerator :: Program GenerateInstructionIn GenerateInstructionOut +instructionGenerator = + predict + ( mkSignature + "Use the information below to understand a task that is solved with calls to a language \ + \model, then write a NEW instruction that will prompt the model to solve the task better. \ + \Consider the dataset notes, the program outline and what it does, the target module's \ + \role and its fields, any demonstrations, the previously-tried instructions with their \ + \scores, the current instruction, and the tip. Write only the new instruction text in the \ + \`proposedInstruction` field." + ) + +-- | Gather every grounding signal and return ranked candidate instructions for the +-- request's target node. Runs, in order: dataset summary → program pseudo-code → +-- program description → module description → @numCandidates@ instruction generations +-- (each varying the tip so the candidates are distinct), then prepends the current +-- instruction and deduplicates preserving order. Makes exactly @4 + numCandidates@ LM +-- calls (2 for the summary, 1 program describe, 1 module describe, N generations). +proposeInstructions :: + (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es) => + Dataset i o -> + ProposeRequest i o -> + Eff es ProposeResult +proposeInstructions ds req = do + summary <- datasetSummary (viewBatch req) ds + let prog = program req + code = renderProgramPseudo prog + sigStr = moduleSignatureAt (targetNode req) prog + sampleRow = case take 1 (datasetExamples ds) of + (ex : _) -> renderExampleRow ex + [] -> "No example provided." + ProgramDescribeOut desc <- runProgram programDescriber (ProgramDescribeIn code sampleRow) + ModuleDescribeOut modDesc <- runProgram moduleDescriber (ModuleDescribeIn code desc sigStr) + proposed <- forM [0 .. max 0 (numCandidates req) - 1] $ \j -> do + let genIn = + GenerateInstructionIn + { datasetDescription = summary, + programCode = code, + programDescription = desc, + moduleSignature = sigStr, + moduleDescription = modDesc, + taskDemos = renderDemos (bootstrappedDemos req), + previousInstructions = renderHistory 5 (history req), + basicInstruction = currentInstruction req, + tip = tipAt (tipIndex req + j) + } + GenerateInstructionOut p <- runProgram instructionGenerator genIn + pure p + pure (ProposeResult (dedup (currentInstruction req : proposed))) + +-- | Render bootstrapped demos, or a placeholder when there are none. +renderDemos :: [Text] -> Text +renderDemos [] = "No task demos provided." +renderDemos xs = T.intercalate "\n" xs + +-- | Order-preserving de-duplication. +dedup :: (Eq a) => [a] -> [a] +dedup = go [] + where + go _ [] = [] + go seen (x : xs) + | x `elem` seen = go seen xs + | otherwise = x : go (x : seen) xs diff --git a/src/Shikumi/Optimize/Propose/Summarize.hs b/src/Shikumi/Optimize/Propose/Summarize.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Propose/Summarize.hs @@ -0,0 +1,246 @@ +{-# LANGUAGE FlexibleContexts #-} + +-- | The signal-gatherers of the grounded proposer (EP-19) that are themselves typed +-- Shikumi 'Program's — preserving V1's "the optimizer is written in the framework it +-- optimizes" pattern. Each mirrors a DSPy proposer sub-module: +-- +-- * 'renderProgramPseudo' / 'programDescriber' — DSPy's @DescribeProgram@: render +-- the whole program as deterministic pseudo-code and describe what it does. +-- * 'datasetDescriber' / 'observationSummarizer' / 'datasetSummary' — DSPy's +-- @DatasetDescriptor@ + @ObservationSummarizer@: observe patterns across sampled +-- rows, then condense them into a 2-3 sentence summary. +-- * 'moduleDescriber' — DSPy's @DescribeModule@: describe one node's role within +-- the broader program. +module Shikumi.Optimize.Propose.Summarize + ( -- * Program pseudo-code + renderProgramPseudo, + + -- * Row rendering (shared with the grounded driver) + renderExampleRow, + encodeJsonText, + + -- * The program describer + ProgramDescribeIn (..), + ProgramDescribeOut (..), + programDescriber, + + -- * The dataset summarizer + DatasetDescribeIn (..), + DatasetDescribeOut (..), + datasetDescriber, + SummarizeIn (..), + SummarizeOut (..), + observationSummarizer, + datasetSummary, + + -- * The module describer + ModuleDescribeIn (..), + ModuleDescribeOut (..), + moduleDescriber, + ) +where + +import Data.Aeson (ToJSON) +import Data.Aeson.Text (encodeToLazyText) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Lazy qualified as TL +import Effectful (Eff, (:>)) +import Effectful.Error.Static (Error) +import GHC.Generics (Generic) +import Shikumi.Adapter (ToPrompt) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval (Dataset, Example (..), datasetExamples) +import Shikumi.LLM (LLM) +import Shikumi.Module (predict) +import Shikumi.Optimize.Propose.Types (NodeFieldNames (..), programFieldNames) +import Shikumi.Program + ( Program, + ProgramShape (..), + programShape, + runProgram, + ) +import Shikumi.Schema (FromModel, ToSchema, Validatable) +import Shikumi.Signature (mkSignature) + +-- --------------------------------------------------------------------------- +-- Program pseudo-code +-- --------------------------------------------------------------------------- + +-- | Render a program as a short, deterministic, human-readable outline: one line +-- per node and combinator, with 'Predict' nodes shown as @predict(inputs) -> +-- outputs@ (using 'programFieldNames') and combinators shown by name. Shikumi's +-- analogue of DSPy's @get_dspy_source_code@. Deterministic, so tests can assert on it. +renderProgramPseudo :: Program i o -> Text +renderProgramPseudo prog = T.intercalate "\n" (fst (go (programShape prog) (programFieldNames prog))) + where + go :: ProgramShape -> [NodeFieldNames] -> ([Text], [NodeFieldNames]) + go (ShapePredict lbl) nfs = case nfs of + (nf : rest) -> + (["predict(" <> commas (inputFieldNames nf) <> ") -> " <> commas (outputFieldNames nf)], rest) + [] -> (["predict(?) -> " <> lbl], []) + go (ShapeCompose a b) nfs = combine "compose" [a, b] nfs + go (ShapeFMap a) nfs = combine "fmap" [a] nfs + go (ShapeMap _ a) nfs = combine "map" [a] nfs + go (ShapeParallel a b) nfs = combine "parallel" [a, b] nfs + go (ShapeRetry _ a) nfs = combine "retry" [a] nfs + go (ShapeRetryWhen _ a) nfs = combine "retryWhen" [a] nfs + go (ShapeValidate a) nfs = combine "validate" [a] nfs + go (ShapeMajorityVote _ _ a) nfs = combine "majorityVote" [a] nfs + go (ShapeEnsemble as) nfs = combine "ensemble" as nfs + go ShapeEmbed nfs = (["embed"], nfs) + combine name children nfs0 = + let (lns, rest) = foldl step ([], nfs0) children + step (acc, nf) c = let (l, nf') = go c nf in (acc ++ l, nf') + in ((name <> ":") : map (" " <>) lns, rest) + commas xs = if null xs then "?" else T.intercalate ", " xs + +-- | Render a single dataset example as @ => @. +renderExampleRow :: (ToJSON i, ToJSON o) => Example i o -> Text +renderExampleRow (Example i o) = encodeJsonText i <> " => " <> encodeJsonText o + +-- | Encode any 'ToJSON' value to compact 'Text'. +encodeJsonText :: (ToJSON a) => a -> Text +encodeJsonText = TL.toStrict . encodeToLazyText + +-- --------------------------------------------------------------------------- +-- The program describer +-- --------------------------------------------------------------------------- + +data ProgramDescribeIn = ProgramDescribeIn + { programCode :: Text, + programExample :: Text + } + deriving stock (Generic, Show) + +instance FromModel ProgramDescribeIn + +instance ToPrompt ProgramDescribeIn + +newtype ProgramDescribeOut = ProgramDescribeOut {programDescription :: Text} + deriving stock (Generic, Show) + +instance ToSchema ProgramDescribeOut + +instance FromModel ProgramDescribeOut + +instance ToPrompt ProgramDescribeOut + +instance Validatable ProgramDescribeOut + +programDescriber :: Program ProgramDescribeIn ProgramDescribeOut +programDescriber = + predict + ( mkSignature + "Below is pseudo-code for a pipeline that solves a task with calls to a language \ + \model. Describe, in two or three sentences, what task this program is designed to \ + \solve and how it goes about solving it. Write your answer in the `programDescription` field." + ) + +-- --------------------------------------------------------------------------- +-- The dataset summarizer +-- --------------------------------------------------------------------------- + +newtype DatasetDescribeIn = DatasetDescribeIn {examples :: Text} + deriving stock (Generic, Show) + +instance FromModel DatasetDescribeIn + +instance ToPrompt DatasetDescribeIn + +newtype DatasetDescribeOut = DatasetDescribeOut {observations :: Text} + deriving stock (Generic, Show) + +instance ToSchema DatasetDescribeOut + +instance FromModel DatasetDescribeOut + +instance ToPrompt DatasetDescribeOut + +instance Validatable DatasetDescribeOut + +datasetDescriber :: Program DatasetDescribeIn DatasetDescribeOut +datasetDescriber = + predict + ( mkSignature + "Given several examples from a dataset, write observations about trends that hold for \ + \most or all of the samples — topics, content, syntax, brevity — and guess at the task \ + \the dataset enables. Write them in the `observations` field." + ) + +newtype SummarizeIn = SummarizeIn {observations :: Text} + deriving stock (Generic, Show) + +instance FromModel SummarizeIn + +instance ToPrompt SummarizeIn + +newtype SummarizeOut = SummarizeOut {summary :: Text} + deriving stock (Generic, Show) + +instance ToSchema SummarizeOut + +instance FromModel SummarizeOut + +instance ToPrompt SummarizeOut + +instance Validatable SummarizeOut + +observationSummarizer :: Program SummarizeIn SummarizeOut +observationSummarizer = + predict + ( mkSignature + "Condense the dataset notes below into a brief two-to-three sentence summary highlighting \ + \only the most important details. Write it in the `summary` field." + ) + +-- | Sample up to @viewBatch@ rows, render them, observe trends, then condense into a +-- short summary. A single observation batch (DSPy's batched refinement loop is noted +-- as a future enhancement); the goal here is that the summary signal /reaches/ the +-- proposer. +datasetSummary :: + (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es) => + -- | viewBatch: how many rows to sample + Int -> + Dataset i o -> + Eff es Text +datasetSummary viewBatch ds = do + let rendered = T.intercalate "\n" (map renderExampleRow (take (max 1 viewBatch) (datasetExamples ds))) + DatasetDescribeOut obs <- runProgram datasetDescriber (DatasetDescribeIn rendered) + SummarizeOut s <- runProgram observationSummarizer (SummarizeIn obs) + pure s + +-- --------------------------------------------------------------------------- +-- The module describer +-- --------------------------------------------------------------------------- + +data ModuleDescribeIn = ModuleDescribeIn + { programCode :: Text, + programDescription :: Text, + moduleSignature :: Text + } + deriving stock (Generic, Show) + +instance FromModel ModuleDescribeIn + +instance ToPrompt ModuleDescribeIn + +newtype ModuleDescribeOut = ModuleDescribeOut {moduleDescription :: Text} + deriving stock (Generic, Show) + +instance ToSchema ModuleDescribeOut + +instance FromModel ModuleDescribeOut + +instance ToPrompt ModuleDescribeOut + +instance Validatable ModuleDescribeOut + +moduleDescriber :: Program ModuleDescribeIn ModuleDescribeOut +moduleDescriber = + predict + ( mkSignature + "Below is pseudo-code for a language-model pipeline and a description of what it does. \ + \Describe the role of the one specified module within the broader program. Write it in \ + \the `moduleDescription` field." + ) diff --git a/src/Shikumi/Optimize/Propose/Tips.hs b/src/Shikumi/Optimize/Propose/Tips.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Propose/Tips.hs @@ -0,0 +1,27 @@ +-- | The stylistic tip bank (EP-19), porting DSPy's @GroundedProposer.TIPS@. A tip +-- is a one-sentence suggestion to the proposing model. Selection is deterministic +-- (an index modulo the bank size), not a hidden RNG, so proposals are reproducible +-- and a search can vary the tip across rounds by varying the index. +module Shikumi.Optimize.Propose.Tips + ( tipBank, + tipAt, + ) +where + +import Data.Text (Text) + +-- | The stylistic tips, in a fixed order: none, creative, simple, descriptive, +-- high-stakes, persona. +tipBank :: [Text] +tipBank = + [ "", + "Don't be afraid to be creative when creating the new instruction!", + "Keep the instruction clear and concise.", + "Make sure your instruction is very informative and descriptive.", + "The instruction should describe a high-stakes scenario in which the model must solve the task.", + "Include a persona relevant to the task in the instruction (e.g. \"You are a ...\")." + ] + +-- | The tip at @i@ modulo the bank size (so an out-of-range index wraps). +tipAt :: Int -> Text +tipAt i = tipBank !! (i `mod` length tipBank) diff --git a/src/Shikumi/Optimize/Propose/Types.hs b/src/Shikumi/Optimize/Propose/Types.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Propose/Types.hs @@ -0,0 +1,102 @@ +-- | The shared types of the grounded instruction proposer (EP-19): the per-node +-- field-metadata accessor (integration point #3), the instruction-history vocabulary, +-- and the proposer's request/result records. +-- +-- This is the contract MIPROv2 (@docs/plans/20-miprov2-optimizer.md@) and COPRO +-- (@docs/plans/21-copro-instruction-optimizer.md@) both consume, so it lives in its +-- own module — neither optimizer drags in @instructionSearch@'s loop to use it. +module Shikumi.Optimize.Propose.Types + ( -- * Per-node field metadata (integration point #3) + NodeFieldNames (..), + programFieldNames, + moduleSignatureAt, + + -- * Instruction history + PastInstruction (..), + renderHistory, + + -- * The proposer request/result + ProposeRequest (..), + ProposeResult (..), + ) +where + +import Data.Text (Text) +import Data.Text qualified as T +import GHC.Generics (Generic) +import Shikumi.Program (NodeFields (NodeFields), Program, nodeFieldsIndexed) + +-- | A node's input/output field names, recovered structurally. A 'Predict' node +-- hides its @i@/@o@ types existentially, so this carries the field /names/ (plain +-- 'Text'), never a typed @Signature@ — exactly what EP-16's @nodeFieldsIndexed@ +-- returns. +data NodeFieldNames = NodeFieldNames + { inputFieldNames :: ![Text], + outputFieldNames :: ![Text] + } + deriving stock (Eq, Show, Generic) + +-- | One 'NodeFieldNames' per 'Predict' node, in @foldParams@/@mapParamsAt@ order +-- (integration point #3). Delegates to EP-16's @nodeFieldsIndexed@; the count and +-- ordering align with @foldParams@ by construction, so +-- @programFieldNames prog !! k@ describes the node @mapParamsAt k@ edits. +programFieldNames :: Program i o -> [NodeFieldNames] +programFieldNames = map convert . nodeFieldsIndexed + where + convert (NodeFields i o) = NodeFieldNames i o + +-- | A human-readable signature for node @k@: @predict(in1, in2) -> out1@. Used as +-- the @moduleSignature@ signal in a proposal prompt. An out-of-range index yields +-- @predict(?) -> ?@. +moduleSignatureAt :: Int -> Program i o -> Text +moduleSignatureAt k prog = case drop k (programFieldNames prog) of + (nf : _) -> "predict(" <> commas (inputFieldNames nf) <> ") -> " <> commas (outputFieldNames nf) + [] -> "predict(?) -> ?" + where + commas xs = if null xs then "?" else T.intercalate ", " xs + +-- | One past attempt for a node: the instruction text and the score it earned. +data PastInstruction = PastInstruction + { pastInstruction :: !Text, + pastScore :: !Double + } + deriving stock (Eq, Show, Generic) + +-- | Render the instruction history as lines @"score 0.83 :: "@, capped +-- at @maxInHistory@ entries. Empty history renders as @"No previous instructions."@ +-- so the proposal prompt is always well-formed. +renderHistory :: Int -> [PastInstruction] -> Text +renderHistory maxInHistory hist = case take (max 0 maxInHistory) hist of + [] -> "No previous instructions." + xs -> T.intercalate "\n" (map line xs) + where + line (PastInstruction ins sc) = "score " <> tshow sc <> " :: " <> ins + +-- | Everything the grounded proposer needs to propose instructions for one node. +data ProposeRequest i o = ProposeRequest + { -- | the program being optimized + program :: !(Program i o), + -- | @foldParams@ index of the node to propose for + targetNode :: !Int, + -- | the node's instruction now (always retained as a candidate) + currentInstruction :: !Text, + -- | prior attempts + scores for this node + history :: ![PastInstruction], + -- | rendered demos for this node (may be empty) + bootstrappedDemos :: ![Text], + -- | how many fresh candidates to ask for (N) + numCandidates :: !Int, + -- | selects a tip deterministically (added to the per-candidate offset) + tipIndex :: !Int, + -- | dataset rows to sample for the summary + viewBatch :: !Int + } + +-- | A ranked list of candidate instruction strings for one node, distinct and with +-- the current instruction always present (and first). The proposer does not score — +-- scoring belongs to the calling optimizer. +newtype ProposeResult = ProposeResult {rankedCandidates :: [Text]} + deriving stock (Eq, Show, Generic) + +tshow :: (Show a) => a -> Text +tshow = T.pack . show diff --git a/src/Shikumi/Optimize/RandomSearch.hs b/src/Shikumi/Optimize/RandomSearch.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/RandomSearch.hs @@ -0,0 +1,98 @@ +-- | Bootstrap few-shot with random search (EP-23, DSPy's +-- @BootstrapFewShotWithRandomSearch@): run V1's 'bootstrapFewShot' several times with +-- different deterministic seeds — each shuffling the trainset and picking a random +-- demo count — score each resulting program, and keep the best. A zero-shot baseline +-- candidate is always included, so the search can never do worse than zero-shot. +-- +-- This is a /wrapper/ over V1's bootstrap, not a re-implementation (MasterPlan +-- integration point #4). Randomness is a deterministic LCG (glibc constants), +-- mirroring "Shikumi.Optimize.Ensemble", so runs are reproducible with no IO entropy. +module Shikumi.Optimize.RandomSearch + ( bootstrapRandomSearch, + bootstrapRandomSearchWith, + RandomSearchConfig (..), + defaultRandomSearchConfig, + ) +where + +import Data.Aeson (ToJSON) +import Data.List (sortBy) +import Data.Ord (comparing) +import Shikumi.Compile.Types (compiledProgram) +import Shikumi.Eval (Dataset, dataset, datasetExamples) +import Shikumi.Optimize.Bootstrap (BootstrapConfig (..), bootstrapFewShotWith, defaultBootstrapConfig) +import Shikumi.Optimize.Search (freezeProgram, scoreOn, selectBest) +import Shikumi.Optimize.Types (Budget (..), Optimizer (..), Scored (..)) +import Shikumi.Program (Program) + +-- | Random-search tunables. (The bootstrap pass threshold is left at +-- 'defaultBootstrapConfig'\'s @1.0@ — keep only exactly-correct teacher runs.) +data RandomSearchConfig = RandomSearchConfig + { -- | lower bound on the per-seed random demo count + minDemos :: !Int, + -- | upper bound on the per-seed random demo count + maxDemos :: !Int + } + deriving stock (Eq, Show) + +-- | Demo count between 1 and 4. +defaultRandomSearchConfig :: RandomSearchConfig +defaultRandomSearchConfig = RandomSearchConfig {minDemos = 1, maxDemos = 4} + +-- | A deterministic linear-congruential stream (glibc constants), duplicated by +-- intent from "Shikumi.Optimize.Ensemble" (which keeps it module-private). +lcg :: Int -> [Int] +lcg s0 = drop 1 (iterate step s0) + where + step s = (1103515245 * s + 12345) `mod` 2147483648 + +-- | A deterministic seed-dependent permutation of the trainset (each example seen at +-- most once), so different seeds give different bootstrap subsets. +shuffle :: Int -> Dataset i o -> Dataset i o +shuffle seed train = + let exs = datasetExamples train + keys = take (length exs) (lcg (seed + 1)) + in dataset (map snd (sortBy (comparing fst) (zip keys exs))) + +-- | The per-seed demo count, in @[minDemos .. maxDemos]@. +sizeFor :: RandomSearchConfig -> Int -> Int +sizeFor cfg seed = + let lo = minDemos cfg + hi = max lo (maxDemos cfg) + span' = hi - lo + 1 + in case lcg seed of + (x : _) -> lo + (x `mod` span') + [] -> lo + +-- | 'bootstrapRandomSearch' with explicit tunables. +bootstrapRandomSearchWith :: + (ToJSON i, ToJSON o) => + RandomSearchConfig -> + -- | teacher + Program i o -> + -- | numCandidates: random seeds to try + Int -> + Budget -> + Optimizer i o +bootstrapRandomSearchWith cfg teacher numCandidates budget = Optimizer $ \train metric student -> do + let seeds = [1 .. max 1 numCandidates] + candidateFor seed = do + let cfg' = defaultBootstrapConfig {maxBootstrappedDemos = sizeFor cfg seed} + opt = bootstrapFewShotWith cfg' teacher budget + compiledProgram <$> runOptimizer opt (shuffle seed train) metric student + seeded <- mapM candidateFor seeds + let cands = student : seeded -- zero-shot baseline first + best <- selectBest budget (scoreOn train metric) cands + pure $ case best of + Nothing -> freezeProgram student + Just sc -> freezeProgram (candidate sc) + +-- | Run V1 bootstrap over @numCandidates@ random seeds plus a zero-shot baseline, +-- score each on the dataset, and keep the best-scoring 'CompiledProgram'. +bootstrapRandomSearch :: + (ToJSON i, ToJSON o) => + Program i o -> + Int -> + Budget -> + Optimizer i o +bootstrapRandomSearch = bootstrapRandomSearchWith defaultRandomSearchConfig diff --git a/src/Shikumi/Optimize/Search.hs b/src/Shikumi/Optimize/Search.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Search.hs @@ -0,0 +1,82 @@ +-- | The shared search-state plumbing every concrete optimizer reuses: the pure +-- 'selectBest' fold, the @evaluate@-backed 'scoreOn' scorer, and 'freezeProgram'. +-- +-- This module sits /below/ both 'Shikumi.Optimize' (which re-exports it) and the +-- four optimizer modules (which import it), so there is no import cycle: the +-- optimizers depend on the plumbing, not on the public driver. +-- +-- __No global mutable state.__ Candidates are threaded as ordinary values; the +-- only effectful step is scoring a candidate (running it over the dataset). +module Shikumi.Optimize.Search + ( selectBest, + scoreOn, + freezeProgram, + setNodeInstr, + instructionAt, + ) +where + +import Data.Text (Text) +import Effectful (Eff, (:>)) +import Effectful.Concurrent (Concurrent) +import Effectful.Error.Static (Error) +import Effectful.Prim (Prim) +import Shikumi.Compile.Types (CompiledProgram (..)) +import Shikumi.Effect.Time (Time) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval (Dataset, Metric, Report (aggregateScore), evaluatePure) +import Shikumi.LLM (LLM) +import Shikumi.Optimize.Types (Budget (..), Scored (..)) +import Shikumi.Program (Params (..), Program, foldParams, mapParamsAt) + +-- | Score every candidate (left to right), stopping once the candidate budget is +-- hit, and return the best by score (ties: earliest wins). The scorer is the +-- /only/ effectful part; the selection is a pure fold over the results, so the +-- search is reproducible and trivially testable by swapping in a pure scorer. +selectBest :: + (Monad m) => + Budget -> + -- | scorer (typically 'scoreOn') + (cand -> m Double) -> + -- | candidates, threaded explicitly + [cand] -> + m (Maybe (Scored cand)) +selectBest budget scorer cands = do + scored <- mapM scoreOne (take (max 0 (maxCandidates budget)) cands) + pure (bestOf scored) + where + scoreOne c = do + s <- scorer c + pure (Scored c s) + -- Fold to the maximum by score; @>@ (strict) keeps the earliest on ties. + bestOf [] = Nothing + bestOf (x : xs) = Just (foldl' (\b c -> if score c > score b then c else b) x xs) + +-- | Score a candidate program against a dataset and pure metric: run EP-8's +-- @evaluate@ and take the aggregate. One call scores one program over the whole +-- dataset (one LM call per example). +scoreOn :: + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + Dataset i o -> + Metric o -> + Program i o -> + Eff es Double +scoreOn ds m p = aggregateScore <$> evaluatePure ds m p + +-- | Mark a program compiled without changing its parameters. EP-9 stores each +-- node's parameters /on/ the node, so an optimizer that has already rewritten a +-- program's parameters just wraps the finished program — there is nothing to +-- freeze on the side. +freezeProgram :: Program i o -> CompiledProgram i o +freezeProgram = CompiledProgram + +-- | Set node @idx@'s instruction override (in @foldParams@ order). Shared by +-- @instructionSearch@ and @copro@. +setNodeInstr :: Int -> Text -> Program i o -> Program i o +setNodeInstr idx instr = mapParamsAt idx (\ps -> ps {instructionOverride = Just instr}) + +-- | The instruction override stored at node @idx@ (in @foldParams@ order), if any. +instructionAt :: Int -> Program i o -> Maybe Text +instructionAt idx prog = case drop idx (foldParams prog) of + (ps : _) -> instructionOverride ps + [] -> Nothing diff --git a/src/Shikumi/Optimize/Types.hs b/src/Shikumi/Optimize/Types.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Optimize/Types.hs @@ -0,0 +1,92 @@ +{-# LANGUAGE RankNTypes #-} + +-- | The central abstractions of the optimizer framework (EP-10): the 'Optimizer' +-- strategy object, the search 'Budget', and the 'Scored' candidate wrapper. +-- +-- An 'Optimizer' is a /search procedure/: given a training 'Dataset', a 'Metric', +-- 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. +-- +-- __Why a record of functions, not a typeclass.__ The four optimizers differ in +-- what extra inputs they carry (a teacher program, a budget, an inner optimizer +-- for ensembling) but share one driver signature. A @newtype@ whose field is the +-- driver lets each smart constructor ('Shikumi.Optimize.LabeledFewShot.labeledFewShot', +-- 'Shikumi.Optimize.Bootstrap.bootstrapFewShot', …) close over its own +-- configuration while presenting a uniform interface to +-- 'Shikumi.Optimize.optimize'. A typeclass would force the extra configuration +-- into associated types and make it impossible to build an optimizer at runtime +-- from CLI flags (which EP-12 needs). +-- +-- __The effect row.__ The rank-2 'runOptimizer' field quantifies over the exact +-- effect row that EP-8's @evaluate@ requires — +-- @(LLM, Concurrent, Error ShikumiError, Time, Prim)@ — because scoring a candidate +-- /is/ a call to @evaluate@ (via 'Shikumi.Optimize.scoreOn'). This is wider than +-- the plan's original @(LLM :> es)@ sketch: the delivered EP-8 runner threads +-- 'Concurrent' (bounded parallelism), 'Time' (monotonic per-example latency via +-- the 'Shikumi.Effect.Time' effect), and 'Prim' (the usage-accounting 'IORef' in +-- 'Shikumi.Eval.Usage.withUsageTotals'). 'IOE' is no longer required — it is +-- supplied only at the discharge edge by 'runEff' under 'runTime'/'runPrim'. See +-- the plan's Decision Log. +module Shikumi.Optimize.Types + ( Optimizer (..), + Budget (..), + defaultBudget, + Scored (..), + ) +where + +import Effectful (Eff, (:>)) +import Effectful.Concurrent (Concurrent) +import Effectful.Error.Static (Error) +import Effectful.Prim (Prim) +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.Program (Program) + +-- | A search strategy that, given a training dataset, a metric, and a starting +-- (student) program, searches for better node parameters and returns a compiled +-- program. The rank-2 field shares one driver signature across all four +-- optimizers (see the module header). +newtype Optimizer i o = Optimizer + { runOptimizer :: + forall es. + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + -- \| training set the search may fit to + Dataset i o -> + -- \| how candidates are scored + Metric o -> + -- \| the starting (student) program + Program i o -> + Eff es (CompiledProgram i o) + } + +-- | A hard, explicit bound on search cost. The optimizers that issue LM calls +-- count the calls they make (proposer calls + per-candidate scoring evaluations, +-- each scoring evaluation costing one LM call per dataset example) and stop — +-- returning the best candidate found /so far/ — before any bound is exceeded, so +-- they never silently produce an unscored program and never blow a cost ceiling. +data Budget = Budget + { -- | ceiling on raw LM calls the optimizer may make (proposals + scoring) + maxLmCalls :: !Int, + -- | ceiling on the number of candidate programs scored + maxCandidates :: !Int + } + deriving stock (Eq, Show) + +-- | A sensible default budget: up to 200 LM calls across at most 32 candidates. +defaultBudget :: Budget +defaultBudget = Budget {maxLmCalls = 200, maxCandidates = 32} + +-- | A candidate paired with its score. Threaded as a plain value through the pure +-- selection fold; there is no mutable search state. +data Scored a = Scored + { candidate :: a, + score :: !Double + } + deriving stock (Eq, Show) diff --git a/test/AcceptanceSpec.hs b/test/AcceptanceSpec.hs new file mode 100644 --- /dev/null +++ b/test/AcceptanceSpec.hs @@ -0,0 +1,100 @@ +{-# LANGUAGE TypeApplications #-} + +-- | M5 — the end-to-end acceptance that proves the whole point. A deliberately +-- underspecified sentiment program (empty instruction, no demos) answers a +-- /held-out/ set it never sees during search poorly (every answer is the +-- never-correct @"neutral"@, score 0). After @optimize@ over a disjoint training +-- set, the returned 'CompiledProgram' answers the same held-out set strictly +-- better (a perfect 1.0). The improvement is shown for three different strategies +-- (bootstrap, labeled few-shot, instruction search), so it is not specific to one. +-- +-- Everything runs offline against the deterministic stub LM — no network, no key, +-- reproducible. +module AcceptanceSpec (tests) where + +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +import Shikumi.Compile.Types (Compiler (runCompiler), compiledProgram) +import Shikumi.Compile.ZeroShot (zeroShot) +import Shikumi.Effect.Time (Time, runTime) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval (Dataset, dataset, exactMatch, example) +import Shikumi.LLM (LLM) +import Shikumi.Optimize +import Shikumi.Program (Program) +import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, sentimentProg) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase) + +runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act + +-- | Six labelled training examples spanning both classes (the @good@/@bad@ +-- lexicon), across three topics. +trainset :: Dataset Sentence Label +trainset = + dataset + [ example (Sentence "good film") (Label "positive"), + example (Sentence "good book") (Label "positive"), + example (Sentence "good song") (Label "positive"), + example (Sentence "bad film") (Label "negative"), + example (Sentence "bad book") (Label "negative"), + example (Sentence "bad song") (Label "negative") + ] + +-- | A disjoint held-out set (different topics) the optimizer never sees. +heldout :: Dataset Sentence Label +heldout = + dataset + [ example (Sentence "good movie") (Label "positive"), + example (Sentence "good meal") (Label "positive"), + example (Sentence "bad movie") (Label "negative"), + example (Sentence "bad meal") (Label "negative") + ] + +-- | The teacher for bootstrap: the student with the ground-truth-unlocking +-- instruction baked in, so it classifies the training set correctly. +teacher :: Program Sentence Label +teacher = runCompiler (zeroShot ruleInstruction) sentimentProg + +-- | The deliberately-underspecified student. +underspecified :: Program Sentence Label +underspecified = sentimentProg + +tests :: TestTree +tests = + testGroup + "M5 acceptance" + [ testCase "optimization strictly improves held-out score (3 strategies)" $ do + res <- runStub $ do + before <- scoreOn heldout exactMatch underspecified + afterBoot <- do + cp <- optimize (bootstrapFewShot teacher defaultBudget) trainset exactMatch underspecified + scoreOn heldout exactMatch (compiledProgram cp) + afterLabeled <- do + cp <- optimize (labeledFewShot 2) trainset exactMatch underspecified + scoreOn heldout exactMatch (compiledProgram cp) + afterInstr <- do + cp <- optimize (instructionSearch 3 defaultBudget) trainset exactMatch underspecified + scoreOn heldout exactMatch (compiledProgram cp) + pure (before, afterBoot, afterLabeled, afterInstr) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (before, afterBoot, afterLabeled, afterInstr) -> do + assertBool ("before should be low, got " <> show before) (before < 0.5) + improved "bootstrapFewShot" before afterBoot + improved "labeledFewShot" before afterLabeled + improved "instructionSearch" before afterInstr + ] + where + -- Strictly higher than before, and at least a floor of 0.75, so the test fails + -- loudly if a future change makes an optimizer a near-no-op. + improved name before after = do + assertBool + (name <> ": expected " <> show after <> " > " <> show before) + (after > before) + assertBool + (name <> ": expected " <> show after <> " >= 0.75") + (after >= 0.75) diff --git a/test/BootstrapSpec.hs b/test/BootstrapSpec.hs new file mode 100644 --- /dev/null +++ b/test/BootstrapSpec.hs @@ -0,0 +1,72 @@ +{-# LANGUAGE TypeApplications #-} + +-- | M2: bootstrap few-shot. Two tests. (a) 'recoverDemo' turns a typed +-- input/output pair into a JSON 'Demo' that round-trips through @fromModel@ — i.e. +-- the recovered demo is decodable back into the node's typed channel, which is +-- what the run-time adapter does. (b) @bootstrapFewShot@ runs a (rule-bearing) +-- teacher over a training set in which one example carries a deliberately wrong +-- label; the search attaches demos only for the examples the metric judged the +-- teacher correct on, and the mislabelled one is excluded. +module BootstrapSpec (tests) where + +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +import Shikumi.Compile.Types (CompiledProgram, Compiler (runCompiler), compiledProgram) +import Shikumi.Compile.ZeroShot (zeroShot) +import Shikumi.Effect.Time (Time, runTime) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval (Dataset, dataset, exactMatch, example) +import Shikumi.LLM (LLM) +import Shikumi.Optimize +import Shikumi.Program (Demo (..), Params (..), Program, programParams) +import Shikumi.Schema (fromModel) +import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, sentimentProg) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertFailure, testCase, (@?=)) + +runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act + +-- | A teacher that classifies by the ground-truth rule: the student program with +-- a @RULE@-bearing instruction baked in (and no demos). +teacher :: Program Sentence Label +teacher = runCompiler (zeroShot ruleInstruction) sentimentProg + +-- | Three examples; the third is mislabelled (the teacher will say "positive" for +-- "good soup", disagreeing with the "negative" label), so it must be dropped. +trainset :: Dataset Sentence Label +trainset = + dataset + [ example (Sentence "good film") (Label "positive"), + example (Sentence "bad film") (Label "negative"), + example (Sentence "good soup") (Label "negative") + ] + +-- | The demos baked into a compiled single-node program. +compiledDemos :: CompiledProgram Sentence Label -> [Demo] +compiledDemos cp = case programParams (compiledProgram cp) of + (ps : _) -> demos ps + [] -> [] + +tests :: TestTree +tests = + testGroup + "M2 bootstrap" + [ 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"), + testCase "bootstrap attaches demos only for metric-passing teacher runs" $ do + res <- runStub (optimize (bootstrapFewShot teacher defaultBudget) trainset exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right cp -> do + let ds = compiledDemos cp + ins = traverse (\d -> fromModel (input d)) ds :: Either ShikumiError [Sentence] + outs = traverse (\d -> fromModel (output d)) ds :: Either ShikumiError [Label] + length ds @?= 2 + ins @?= Right [Sentence "good film", Sentence "bad film"] + outs @?= Right [Label "positive", Label "negative"] + ] diff --git a/test/CoproSpec.hs b/test/CoproSpec.hs new file mode 100644 --- /dev/null +++ b/test/CoproSpec.hs @@ -0,0 +1,134 @@ +{-# LANGUAGE TypeApplications #-} + +-- | EP-21 acceptance: COPRO coordinate-ascent instruction optimization. Hermetic, +-- offline. Proves the held-out lift, depth monotonicity, the LM-call budget bound, +-- and the serialization round-trip — the breadth/depth control the MasterPlan tracks. +module CoproSpec (tests) where + +import Data.IORef (newIORef, readIORef) +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +import 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.Optimize + ( Budget (..), + CoproConfig (..), + copro, + defaultCoproConfig, + optimize, + scoreOn, + ) +import Shikumi.Program (Params (..), programParams) +import StubLM + ( Label (..), + Sentence (..), + ruleInstruction, + runStubLM, + runStubLMCounting, + sentimentProg, + ) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +-- --------------------------------------------------------------------------- +-- Fixtures and run helpers +-- --------------------------------------------------------------------------- + +trainset :: Dataset Sentence Label +trainset = + dataset + [ example (Sentence "good film") (Label "positive"), + example (Sentence "bad film") (Label "negative") + ] + +heldout :: Dataset Sentence Label +heldout = + dataset + [ example (Sentence "good book") (Label "positive"), + example (Sentence "bad book") (Label "negative") + ] + +runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act + +tests :: TestTree +tests = testGroup "Copro" [breadthSurfacesMagic, depthMonotone, budgetRespected, heldoutLift, roundTrips] + +-- --------------------------------------------------------------------------- +-- M1/M2 — breadth, depth, budget +-- --------------------------------------------------------------------------- + +breadthSurfacesMagic :: TestTree +breadthSurfacesMagic = + testCase "breadth surfaces and selects the magic RULE instruction" $ do + res <- runStub (optimize (copro defaultCoproConfig {depth = 1}) trainset exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right cp -> case programParams (compiledProgram cp) of + (ps : _) -> instructionOverride ps @?= Just ruleInstruction + [] -> assertFailure "expected one node" + +depthMonotone :: TestTree +depthMonotone = + testCase "deeper depth never scores worse than shallower on held-out" $ do + res <- + runStub $ do + cp1 <- optimize (copro defaultCoproConfig {depth = 1}) trainset exactMatch sentimentProg + s1 <- scoreOn heldout exactMatch (compiledProgram cp1) + cp3 <- optimize (copro defaultCoproConfig {depth = 3}) trainset exactMatch sentimentProg + s3 <- scoreOn heldout exactMatch (compiledProgram cp3) + pure (s1, s3) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (s1, s3) -> assertBool ("depth3 " <> show s3 <> " >= depth1 " <> show s1) (s3 >= s1) + +budgetRespected :: TestTree +budgetRespected = + testCase "respects the LM-call budget" $ do + ref <- newIORef (0 :: Int) + let cfg = defaultCoproConfig {budget = Budget {maxLmCalls = 8, maxCandidates = 32}} + res <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runStubLMCounting ref (optimize (copro cfg) trainset exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right _ -> pure () + count <- readIORef ref + assertBool ("expected 0 < calls <= 8, got " <> show count) (count > 0 && count <= 8) + +-- --------------------------------------------------------------------------- +-- M3 — held-out lift and serialization round-trip +-- --------------------------------------------------------------------------- + +heldoutLift :: TestTree +heldoutLift = + testCase "held-out score strictly improves after copro" $ do + res <- + runStub $ do + before <- scoreOn heldout exactMatch sentimentProg + cp <- optimize (copro defaultCoproConfig) trainset exactMatch sentimentProg + after <- scoreOn heldout exactMatch (compiledProgram cp) + pure (before, after) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (before, after) -> do + assertBool ("before should be low, got " <> show before) (before < 0.5) + assertBool ("after " <> show after <> " should exceed before " <> show before) (after > before) + assertBool ("after should reach the floor, got " <> show after) (after >= 0.75) + +roundTrips :: TestTree +roundTrips = + testCase "copro output round-trips through encodeCompiled/decodeCompiledOnto" $ do + res <- runStub (optimize (copro defaultCoproConfig) trainset exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right cp -> case decodeCompiledOnto sentimentProg (encodeCompiled cp) of + Left err -> assertFailure ("decode failed: " <> err) + Right cp' -> + programParams (compiledProgram cp') @?= programParams (compiledProgram cp) diff --git a/test/EnsembleSpec.hs b/test/EnsembleSpec.hs new file mode 100644 --- /dev/null +++ b/test/EnsembleSpec.hs @@ -0,0 +1,92 @@ +{-# LANGUAGE TypeApplications #-} + +-- | M4: ensemble search. Two tests. (a) The value: three members that each answer +-- 3 of 4 held-out items correctly, but err on /different/ items, combine by +-- majority vote (the exact construction 'ensembleSearch' uses — +-- 'Shikumi.Combinator.ensemble' under 'majorityReducer') into a program that +-- answers all 4 correctly — so the ensemble strictly beats its best member. (b) +-- The wiring: @ensembleSearch n@ runs the inner optimizer @n@ times and produces +-- an @n@-member ensemble (checked structurally via 'programShape'). +module EnsembleSpec (tests) where + +import Data.Aeson (toJSON) +import Data.Text (Text) +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +import Shikumi.Combinator (ensemble) +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.Optimize +import Shikumi.Program (Demo (..), Program, ProgramShape (..), programShape) +import StubLM (Label (..), Sentence (..), runStubLM, sentimentProg) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act + +-- | The held-out set: two positives (@good@) and two negatives (@bad@), across two +-- topics, so a member can be fooled on exactly one item by a poisoned demo. +heldout :: Dataset Sentence Label +heldout = + dataset + [ example (Sentence "good apple") (Label "positive"), + example (Sentence "good melon") (Label "positive"), + example (Sentence "bad apple") (Label "negative"), + example (Sentence "bad melon") (Label "negative") + ] + +mkMember :: [(Text, Text)] -> Program Sentence Label +mkMember pairs = + withDemos [Demo (toJSON (Sentence s)) (toJSON (Label l)) | (s, l) <- pairs] sentimentProg + +-- Each member is correct on 3/4 held-out items, wrong on a different one (a +-- poisoned demo drags one item to the wrong label under nearest-demo voting). +memberA, memberB, memberC :: Program Sentence Label +memberA = mkMember [("good apple", "positive"), ("bad apple", "negative"), ("bad melon", "positive")] +memberB = mkMember [("good apple", "positive"), ("bad melon", "negative"), ("bad apple", "positive")] +memberC = + mkMember + [ ("good apple", "positive"), + ("good melon", "negative"), + ("bad apple", "negative"), + ("bad melon", "negative") + ] + +trainset :: Dataset Sentence Label +trainset = + dataset + [ example (Sentence "good film") (Label "positive"), + example (Sentence "bad film") (Label "negative") + ] + +tests :: TestTree +tests = + testGroup + "M4 ensembleSearch" + [ testCase "majority vote scores at least as high as the best member" $ do + res <- runStub $ do + ms <- mapM (scoreOn heldout exactMatch) [memberA, memberB, memberC] + es <- scoreOn heldout exactMatch (ensemble [memberA, memberB, memberC] majorityReducer) + pure (ms, es) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (memberScores, ensembleScore) -> do + memberScores @?= [0.75, 0.75, 0.75] + ensembleScore @?= 1.0 + assertBool + ("ensemble " <> show ensembleScore <> " should beat best member " <> show (maximum memberScores)) + (ensembleScore > maximum memberScores), + testCase "ensembleSearch n produces an n-member ensemble" $ do + res <- runStub (optimize (ensembleSearch 3 (labeledFewShot 1)) trainset exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right cp -> case programShape (compiledProgram cp) of + ShapeEnsemble subs -> length subs @?= 3 + other -> assertFailure ("expected a 3-member ShapeEnsemble, got " <> show other) + ] diff --git a/test/GepaSpec.hs b/test/GepaSpec.hs new file mode 100644 --- /dev/null +++ b/test/GepaSpec.hs @@ -0,0 +1,165 @@ +{-# LANGUAGE TypeApplications #-} + +-- | EP-22 acceptance: GEPA reflective evolution. Hermetic, offline. Pure unit tests +-- for the Pareto frontier; effectful tests for per-node feedback capture, reflective +-- mutation, and the end-to-end held-out lift + serialization round-trip. +module GepaSpec (tests) where + +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +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, boolScore, dataset, exactMatch, example, predictionPrimary) +import Shikumi.LLM (LLM) +import Shikumi.Optimize + ( Candidate (..), + FeedbackMetric, + captureFeedback, + defaultBudget, + dominates, + gepa, + mutateNode, + optimize, + paretoFrontier, + reflectiveProposer, + sampleParent, + scoreOn, + ) +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 Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +-- --------------------------------------------------------------------------- +-- Fixtures and run helpers +-- --------------------------------------------------------------------------- + +trainset :: Dataset Sentence Label +trainset = + dataset + [ example (Sentence "good film") (Label "positive"), + example (Sentence "bad film") (Label "negative") + ] + +heldout :: Dataset Sentence Label +heldout = + dataset + [ example (Sentence "good book") (Label "positive"), + example (Sentence "bad book") (Label "negative") + ] + +-- | The feedback metric: correct → score 1, no critique; wrong → score 0 and the +-- critique "be more specific" (which the stub's reflective proposer keys off). +fbMetric :: FeedbackMetric Label +fbMetric expd p = + let correct = expd == predictionPrimary p + in (boolScore correct, if correct then "" else "be more specific") + +runGepa :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runGepa act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runGepaStubLM act + +tests :: TestTree +tests = testGroup "Gepa" [paretoPure, feedbackCapture, reflectiveMutation, heldoutLift, roundTrips] + +-- --------------------------------------------------------------------------- +-- Pure Pareto-frontier tests +-- --------------------------------------------------------------------------- + +paretoPure :: TestTree +paretoPure = + testGroup + "gepa-pareto" + [ testCase "dominance and the non-dominated frontier" $ do + let a = Candidate [] [1.0, 0.0] 0.5 + b = Candidate [] [0.0, 0.0] 0.0 + c = Candidate [] [0.0, 1.0] 0.5 + dominates a b @?= True + dominates a c @?= False + paretoFrontier [a, b, c] @?= [a, c], + testCase "sampleParent is deterministic for a fixed seed" $ do + let a = Candidate [] [1.0, 0.0] 0.5 + c = Candidate [] [0.0, 1.0] 0.5 + (fst <$> sampleParent 7 [a, c]) @?= (fst <$> sampleParent 7 [a, c]) + assertBool "samples a frontier member" (maybe False ((`elem` [a, c]) . fst) (sampleParent 7 [a, c])) + ] + +-- --------------------------------------------------------------------------- +-- M1 — feedback capture +-- --------------------------------------------------------------------------- + +feedbackCapture :: TestTree +feedbackCapture = + testCase "captureFeedback accumulates a per-node critique keyed by NodePath" $ do + res <- runGepa (captureFeedback trainset fbMetric sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (fblog, scores) -> do + length scores @?= 2 + case programNodePaths sentimentProg of + (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) + [] -> assertFailure "expected at least one node" + +-- --------------------------------------------------------------------------- +-- M2 — reflective mutation +-- --------------------------------------------------------------------------- + +reflectiveMutation :: TestTree +reflectiveMutation = + testCase "a node whose feedback says be-more-specific receives a mutated instruction" $ do + res <- + runGepa $ do + (fblog, _) <- captureFeedback trainset fbMetric sentimentProg + mutateNode + reflectiveProposer + "a one-node program" + "a tiny dataset" + (nodeFieldsIndexed sentimentProg) + fblog + (programNodePaths sentimentProg) + 0 + sentimentProg + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right mutated -> case foldParams mutated of + -- the original node had no instruction override (Nothing); a Just RULE proves it mutated + (ps : _) -> instructionOverride ps @?= Just ruleInstruction + [] -> assertFailure "expected one node" + +-- --------------------------------------------------------------------------- +-- M3 — held-out lift and serialization round-trip +-- --------------------------------------------------------------------------- + +heldoutLift :: TestTree +heldoutLift = + testCase "reflective evolution lifts the held-out score" $ do + res <- + runGepa $ do + before <- scoreOn heldout exactMatch sentimentProg + cp <- optimize (gepa reflectiveProposer fbMetric defaultBudget) trainset exactMatch sentimentProg + after <- scoreOn heldout exactMatch (compiledProgram cp) + pure (before, after) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (before, after) -> do + assertBool ("before should be low, got " <> show before) (before < 0.5) + assertBool ("after " <> show after <> " should exceed before " <> show before) (after > before) + assertBool ("after should reach the floor, got " <> show after) (after >= 0.75) + +roundTrips :: TestTree +roundTrips = + testCase "gepa output round-trips through encodeCompiled/decodeCompiledOnto" $ do + res <- runGepa (optimize (gepa reflectiveProposer fbMetric defaultBudget) trainset exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right cp -> case decodeCompiledOnto sentimentProg (encodeCompiled cp) of + Left err -> assertFailure ("decode failed: " <> err) + Right cp' -> programParams (compiledProgram cp') @?= programParams (compiledProgram cp) diff --git a/test/InstructionSpec.hs b/test/InstructionSpec.hs new file mode 100644 --- /dev/null +++ b/test/InstructionSpec.hs @@ -0,0 +1,61 @@ +{-# LANGUAGE TypeApplications #-} + +-- | M3: instruction search. Two tests. (a) Among the proposer's offerings, exactly +-- one (variant 0) carries the "magic" @RULE@-bearing instruction that unlocks +-- correct answers under the stub; the search selects it for the node. (b) Under a +-- deliberately tight 'Budget', the recorded LM-call count (proposer calls plus +-- per-candidate scoring calls) never exceeds @maxLmCalls@ — the search stops early +-- and returns the best found so far. +module InstructionSpec (tests) where + +import Data.IORef (newIORef, readIORef) +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +import 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.Optimize +import Shikumi.Program (Params (..), programParams) +import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, runStubLMCounting, sentimentProg) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act + +trainset :: Dataset Sentence Label +trainset = + dataset + [ example (Sentence "good film") (Label "positive"), + example (Sentence "bad film") (Label "negative") + ] + +tests :: TestTree +tests = + testGroup + "M3 instructionSearch" + [ testCase "selects the best proposed instruction for the node" $ do + res <- runStub (optimize (instructionSearch 3 defaultBudget) trainset exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right cp -> case programParams (compiledProgram cp) of + (ps : _) -> instructionOverride ps @?= Just ruleInstruction + [] -> assertFailure "expected one node", + testCase "respects the LM-call budget" $ do + ref <- newIORef (0 :: Int) + let budget = Budget {maxLmCalls = 6, maxCandidates = 32} + res <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runStubLMCounting ref (optimize (instructionSearch 3 budget) trainset exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right _ -> pure () + count <- readIORef ref + assertBool + ("expected 0 < calls <= 6, got " <> show count) + (count > 0 && count <= 6) + ] diff --git a/test/KNNSpec.hs b/test/KNNSpec.hs new file mode 100644 --- /dev/null +++ b/test/KNNSpec.hs @@ -0,0 +1,129 @@ +{-# LANGUAGE TypeApplications #-} + +-- | EP-23 acceptance (KNN): under a known stub-embedder geometry, the demos +-- attached are the input-nearest training examples (run-time form) and the +-- centroid-nearest examples (compile-time fallback). Plus the serialization +-- round-trip: the run-time Embed form carries no Params (like @react@); the centroid +-- form's baked demos persist. +module KNNSpec (tests) where + +import Data.Aeson (toJSON) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Vector (Vector) +import Data.Vector qualified as V +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +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, datasetExamples, exactMatch, example) +import Shikumi.LLM (LLM) +import Shikumi.Optimize (knnDemos, knnFewShotCentroid, nearestDemos, optimize) +import Shikumi.Optimize.Search (freezeProgram) +import Shikumi.Program (Demo (..), Params (..), programParams, runProgram) +import StubLM (Label (..), Sentence (..), runStubLM, sentimentProg) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +-- --------------------------------------------------------------------------- +-- Stub embedder (a known 2-D geometry) and fixtures +-- --------------------------------------------------------------------------- + +-- | Geography text → near @[1, 0]@; arithmetic text → near @[0, 1]@; else the +-- diagonal. The two clusters are near-orthogonal, so "nearest" is unambiguous. +stubEmbed :: Text -> Vector Double +stubEmbed t + | any (`T.isInfixOf` lt) ["france", "paris", "capital", "geography", "city"] = V.fromList [1.0, 0.1] + | any (`T.isInfixOf` lt) ["plus", "sum", "arithmetic", "two", "three"] = V.fromList [0.1, 1.0] + | otherwise = V.fromList [0.5, 0.5] + where + lt = T.toLower t + +-- A balanced training set: two geography, two arithmetic (foldParams index order). +geoA, geoB, ariA, ariB :: Sentence +geoA = Sentence "capital of france" +geoB = Sentence "capital city geography" +ariA = Sentence "two plus two" +ariB = Sentence "three plus three sum" + +balanced :: Dataset Sentence Label +balanced = + dataset + [ example geoA (Label "paris"), + example ariA (Label "four"), + example geoB (Label "city"), + example ariB (Label "six") + ] + +-- A geography-skewed set, so the centroid lands in the geography cluster. +skewed :: Dataset Sentence Label +skewed = + dataset + [ example geoA (Label "paris"), + example geoB (Label "city"), + example (Sentence "capital paris france") (Label "fr"), + example ariA (Label "four") + ] + +runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act + +tests :: TestTree +tests = testGroup "KNN" [inputNearest, centroidNearest, runtimeRuns, roundTrips] + +-- --------------------------------------------------------------------------- +-- M1 — selection +-- --------------------------------------------------------------------------- + +inputNearest :: TestTree +inputNearest = + testCase "nearestDemos attaches the input-nearest training examples" $ do + let geoQuery = nearestDemos stubEmbed 2 (datasetExamples balanced) (toPrompt (Sentence "france capital quiz")) + ariQuery = nearestDemos stubEmbed 2 (datasetExamples balanced) (toPrompt (Sentence "sum two plus three")) + [i | Demo i _ <- geoQuery] @?= [toJSON geoA, toJSON geoB] + [i | Demo i _ <- ariQuery] @?= [toJSON ariA, toJSON ariB] + +centroidNearest :: TestTree +centroidNearest = + testCase "knnFewShotCentroid bakes the centroid-nearest examples as demos" $ do + res <- runStub (optimize (knnFewShotCentroid stubEmbed 2) skewed exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right cp -> case programParams (compiledProgram cp) of + (ps : _) -> [i | Demo i _ <- demos ps] @?= [toJSON geoA, toJSON geoB] + [] -> assertFailure "expected one node" + +runtimeRuns :: TestTree +runtimeRuns = + testCase "the run-time KNN program runs end-to-end under the stub LM" $ do + let prog = knnDemos stubEmbed 2 balanced sentimentProg + res <- runStub (runProgram prog (Sentence "france capital quiz")) + assertBool "the Embed-wrapped KNN program runs" (either (const False) (const True) res) + +-- --------------------------------------------------------------------------- +-- M3 — serialization round-trip +-- --------------------------------------------------------------------------- + +roundTrips :: TestTree +roundTrips = + testGroup + "round-trips" + [ testCase "run-time KNN serializes as the empty Params vector (like react)" $ do + let compiled = freezeProgram (knnDemos stubEmbed 2 balanced sentimentProg) + programParams (compiledProgram compiled) @?= [] + case decodeCompiledOnto (knnDemos stubEmbed 2 balanced sentimentProg) (encodeCompiled compiled) of + Left err -> assertFailure ("decode failed: " <> err) + Right cp' -> programParams (compiledProgram cp') @?= [], + testCase "centroid KNN baked demos round-trip onto the bare student" $ do + res <- runStub (optimize (knnFewShotCentroid stubEmbed 2) skewed exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right cp -> case decodeCompiledOnto sentimentProg (encodeCompiled cp) of + Left err -> assertFailure ("decode failed: " <> err) + Right cp' -> programParams (compiledProgram cp') @?= programParams (compiledProgram cp) + ] diff --git a/test/LabeledFewShotSpec.hs b/test/LabeledFewShotSpec.hs new file mode 100644 --- /dev/null +++ b/test/LabeledFewShotSpec.hs @@ -0,0 +1,74 @@ +{-# LANGUAGE TypeApplications #-} + +-- | M1: labeled few-shot demo selection. The optimizer enumerates every size-@k@ +-- demo set drawn from the training examples, scores each, and keeps the best. The +-- test proves the /chosen/ set is the highest-scoring of those tried (not merely +-- "some demos were attached"): it recomputes the score of every candidate set the +-- optimizer considered and asserts the chosen set attains the maximum, that the +-- candidates are not all equal (so the choice is meaningful), and that the maximum +-- is a perfect 1.0. +module LabeledFewShotSpec (tests) where + +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +import Shikumi.Compile.Types (CompiledProgram, compiledProgram) +import Shikumi.Effect.Time (Time, runTime) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval (Dataset, dataset, exactMatch, example) +import Shikumi.LLM (LLM) +import Shikumi.Optimize +import Shikumi.Program (Demo, Params (..), programParams) +import StubLM (Label (..), Sentence (..), runStubLM, sentimentProg) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act + +-- | A balanced training set: two positives (containing @good@) and two negatives +-- (containing @bad@). A single demo cannot classify all four, so a winning set +-- must pair a @good@-demo with a @bad@-demo whose other words align with the +-- inputs — only some size-2 sets do. +trainset :: Dataset Sentence Label +trainset = + dataset + [ example (Sentence "good movie") (Label "positive"), + example (Sentence "good meal") (Label "positive"), + example (Sentence "bad movie") (Label "negative"), + example (Sentence "bad meal") (Label "negative") + ] + +-- | The demos baked into a compiled single-node program. +chosenDemos :: CompiledProgram Sentence Label -> [Demo] +chosenDemos cp = case programParams (compiledProgram cp) of + (ps : _) -> demos ps + [] -> [] + +tests :: TestTree +tests = + testGroup + "M1 labeledFewShot" + [ testCase "selects the highest-scoring demo set of those tried" $ do + res <- runStub $ do + cp <- optimize (labeledFewShot 2) trainset exactMatch sentimentProg + let chosen = chosenDemos cp + chosenScore <- scoreOn trainset exactMatch (withDemos chosen sentimentProg) + candScores <- + mapM + (\ds -> scoreOn trainset exactMatch (withDemos ds sentimentProg)) + (labeledCandidateSets 2 trainset) + pure (chosen, chosenScore, candScores) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (chosen, chosenScore, candScores) -> do + length chosen @?= 2 + chosenScore @?= maximum candScores + assertBool + ("candidate scores should differ, got " <> show candScores) + (maximum candScores > minimum candScores) + assertBool + ("best candidate should be perfect, got " <> show chosenScore) + (chosenScore == 1.0) + ] diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,39 @@ +-- | EP-10 acceptance suite for @shikumi-optimize@: the M0 scaffold, the four +-- optimizers (M1 labeled few-shot, M2 bootstrap, M3 instruction search, M4 +-- ensemble), and the M5 end-to-end held-out-improvement acceptance. Every group +-- is hermetic — it runs against the deterministic stub LM with no network and no +-- API key. +module Main (main) where + +import AcceptanceSpec qualified +import BootstrapSpec qualified +import CoproSpec qualified +import EnsembleSpec qualified +import GepaSpec qualified +import InstructionSpec qualified +import KNNSpec qualified +import LabeledFewShotSpec qualified +import Miprov2Spec qualified +import OptimizeSpec qualified +import ProposeSpec qualified +import RandomSearchSpec qualified +import Test.Tasty (defaultMain, testGroup) + +main :: IO () +main = + defaultMain $ + testGroup + "shikumi-optimize" + [ OptimizeSpec.tests, + LabeledFewShotSpec.tests, + BootstrapSpec.tests, + InstructionSpec.tests, + EnsembleSpec.tests, + ProposeSpec.tests, + Miprov2Spec.tests, + CoproSpec.tests, + GepaSpec.tests, + KNNSpec.tests, + RandomSearchSpec.tests, + AcceptanceSpec.tests + ] diff --git a/test/Miprov2Spec.hs b/test/Miprov2Spec.hs new file mode 100644 --- /dev/null +++ b/test/Miprov2Spec.hs @@ -0,0 +1,202 @@ +{-# LANGUAGE TypeApplications #-} + +-- | EP-20 acceptance: MIPROv2. Hermetic, offline. The headline (M3) is that on a +-- joint-win task — where a sentence is answered correctly only if /both/ a @RULE@ +-- instruction and a covering demo are present — MIPROv2's joint search lifts the +-- held-out score to 1.0, strictly above the 0.5 that single-axis @instructionSearch@ +-- can reach. M1 checks candidate generation, M2 the search surrogate (budget + +-- joint selection), and M3 the lift plus the serialization round-trip. +module Miprov2Spec (tests) where + +import Data.IORef (newIORef, readIORef) +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +import 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.Optimize + ( Budget (..), + Miprov2Auto (..), + Miprov2Config (..), + bootstrapDemoCandidates, + defaultBudget, + instructionSearch, + miprov2Auto, + miprov2With, + optimize, + proposeInstructionCandidates, + scoreOn, + searchJoint, + ) +import Shikumi.Program (Params (..), mapParamsAt) +import StubLM + ( Label (..), + Sentence (..), + ruleInstruction, + runJointStubLM, + runJointStubLMCounting, + sentimentProg, + ) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +-- --------------------------------------------------------------------------- +-- Fixtures and run helpers +-- --------------------------------------------------------------------------- + +-- | Region A (good/bad) needs a RULE instruction; region B (great/terrible) needs a +-- covering demo. See "StubLM".'StubLM.runJointStubLM'. +jointTrain :: Dataset Sentence Label +jointTrain = + dataset + [ example (Sentence "good movie") (Label "positive"), + example (Sentence "bad movie") (Label "negative"), + example (Sentence "great film") (Label "positive"), + example (Sentence "terrible film") (Label "negative") + ] + +jointHeldout :: Dataset Sentence Label +jointHeldout = + dataset + [ example (Sentence "good show") (Label "positive"), + example (Sentence "bad show") (Label "negative"), + example (Sentence "great show") (Label "positive"), + example (Sentence "terrible show") (Label "negative") + ] + +-- | Run an action under the joint-task stub with the five optimizer effects. +runJoint :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runJoint act = + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runJointStubLM act + +cfgLight :: Miprov2Config +cfgLight = miprov2Auto Miprov2Light + +tests :: TestTree +tests = testGroup "Miprov2" [candidates, search, acceptance, serialize] + +-- --------------------------------------------------------------------------- +-- M1 — candidate generation +-- --------------------------------------------------------------------------- + +candidates :: TestTree +candidates = + testGroup + "candidates" + [ testCase "proposeInstructionCandidates returns the current instruction first plus a RULE candidate" $ do + res <- + runJoint $ do + demoCands <- bootstrapDemoCandidates cfgLight sentimentProg jointTrain exactMatch sentimentProg + proposeInstructionCandidates cfgLight sentimentProg jointTrain demoCands + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right [node0] -> do + length node0 @?= numInstructCandidates cfgLight + take 1 node0 @?= [""] -- the node's current (empty) instruction, first + assertBool "contains the RULE instruction" (ruleInstruction `elem` node0) + Right other -> assertFailure ("expected one node, got " <> show (length other)), + testCase "bootstrapDemoCandidates yields an empty set then a non-empty labelled set" $ do + res <- runJoint (bootstrapDemoCandidates cfgLight sentimentProg jointTrain exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right [sets] -> do + assertBool "candidate 0 is the empty set" (take 1 sets == [[]]) + assertBool "at least one non-empty candidate set" (any (not . null) sets) + Right other -> assertFailure ("expected one node, got " <> show (length other)) + ] + +-- --------------------------------------------------------------------------- +-- M2 — the search surrogate +-- --------------------------------------------------------------------------- + +search :: TestTree +search = + testGroup + "search" + [ testCase "searchJoint selects both axes, beating any single-axis choice" $ do + res <- + runJoint $ do + demoCands <- bootstrapDemoCandidates cfgLight sentimentProg jointTrain exactMatch sentimentProg + instrCands <- proposeInstructionCandidates cfgLight sentimentProg jointTrain demoCands + best <- searchJoint cfgLight jointTrain exactMatch sentimentProg instrCands demoCands + joint <- scoreOn jointHeldout exactMatch best + ruleOnly <- + scoreOn + jointHeldout + exactMatch + (mapParamsAt 0 (\ps -> ps {instructionOverride = Just ruleInstruction}) sentimentProg) + pure (joint, ruleOnly) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (joint, ruleOnly) -> do + ruleOnly @?= 0.5 + assertBool ("joint " <> show joint <> " should beat single-axis " <> show ruleOnly) (joint > ruleOnly) + joint @?= 1.0, + testCase "searchJoint honors the LM-call budget" $ do + ref <- newIORef (0 :: Int) + let cfg = cfgLight {numTrials = 5, minibatchSize = 2, budget = Budget {maxLmCalls = 8, maxCandidates = 32}} + instrCands = [["", ruleInstruction]] + demoCands = [[[], []]] + res <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runJointStubLMCounting ref (searchJoint cfg jointTrain exactMatch sentimentProg instrCands demoCands) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right _ -> pure () + count <- readIORef ref + assertBool ("expected 0 < calls <= 8, got " <> show count) (count > 0 && count <= 8) + ] + +-- --------------------------------------------------------------------------- +-- M3 — the held-out lift and serialization round-trip +-- --------------------------------------------------------------------------- + +acceptance :: TestTree +acceptance = + testGroup + "acceptance" + [ testCase "miprov2 lifts held-out score above before and above instructionSearch" $ do + res <- + runJoint $ do + before <- scoreOn jointHeldout exactMatch sentimentProg + cpInstr <- optimize (instructionSearch 3 defaultBudget) jointTrain exactMatch sentimentProg + afterInstr <- scoreOn jointHeldout exactMatch (compiledProgram cpInstr) + cpMipro <- optimize (miprov2With cfgLight sentimentProg) jointTrain exactMatch sentimentProg + afterMipro <- scoreOn jointHeldout exactMatch (compiledProgram cpMipro) + pure (before, afterInstr, afterMipro) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (before, afterInstr, afterMipro) -> do + assertBool ("before should be low, got " <> show before) (before < 0.5) + assertBool ("miprov2 should lift above before: " <> show (before, afterMipro)) (afterMipro > before) + assertBool + ("miprov2 " <> show afterMipro <> " should beat instructionSearch " <> show afterInstr) + (afterMipro > afterInstr) + assertBool ("miprov2 should reach the floor, got " <> show afterMipro) (afterMipro >= 0.75) + ] + +serialize :: TestTree +serialize = + testGroup + "serialize" + [ testCase "miprov2 output round-trips through encodeCompiled/decodeCompiledOnto" $ do + res <- + runJoint $ do + cp <- optimize (miprov2With cfgLight sentimentProg) jointTrain exactMatch sentimentProg + let bytes = encodeCompiled cp + case decodeCompiledOnto sentimentProg bytes of + Left err -> pure (Left err) + Right decoded -> do + a <- scoreOn jointHeldout exactMatch (compiledProgram cp) + b <- scoreOn jointHeldout exactMatch (compiledProgram decoded) + pure (Right (a, b)) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (Left err) -> assertFailure ("decode failed: " <> err) + Right (Right (a, b)) -> a @?= b + ] diff --git a/test/OptimizeSpec.hs b/test/OptimizeSpec.hs new file mode 100644 --- /dev/null +++ b/test/OptimizeSpec.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE TypeApplications #-} + +-- | M0: the package scaffold and the central abstraction. A /no-op optimizer/ +-- (one that returns its input program compiled unchanged) driven through +-- 'optimize' yields a 'CompiledProgram' whose extracted program is structurally +-- identical to the input — proving the integration surface (the 'Optimizer' +-- record, the @evaluate@-shaped effect row, EP-9's 'CompiledProgram') compiles +-- and round-trips before any real search is written. +module OptimizeSpec (tests) where + +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +import Shikumi.Compile.Types (compiledProgram) +import Shikumi.Effect.Time (Time, runTime) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval (dataset, exactMatch, example) +import Shikumi.LLM (LLM) +import Shikumi.Optimize +import Shikumi.Program (programParams, programShape) +import StubLM (Label (..), Sentence (..), runStubLM, sentimentProg) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertFailure, testCase, (@?=)) + +-- | Run an optimizer/scoring action under the deterministic stub and the +-- @evaluate@ effect row, returning either a 'ShikumiError' or the result. +runStub :: + Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> + IO (Either ShikumiError a) +runStub act = + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act + +-- | The no-op optimizer: ignore the dataset and metric, return the program +-- compiled unchanged. +noopOptimizer :: Optimizer Sentence Label +noopOptimizer = Optimizer (\_ _ p -> pure (freezeProgram p)) + +tests :: TestTree +tests = + testGroup + "M0 scaffold" + [ testCase "no-op optimizer round-trips a program through optimize" $ do + let ds = dataset [example (Sentence "good film") (Label "positive")] + res <- runStub (optimize noopOptimizer ds exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right cp -> do + -- The compiled program is structurally identical to the input. + programShape (compiledProgram cp) @?= programShape sentimentProg + programParams (compiledProgram cp) @?= programParams sentimentProg + ] diff --git a/test/ProposeSpec.hs b/test/ProposeSpec.hs new file mode 100644 --- /dev/null +++ b/test/ProposeSpec.hs @@ -0,0 +1,178 @@ +{-# LANGUAGE TypeApplications #-} + +-- | EP-19 acceptance: the grounded instruction proposer. Three hermetic groups +-- (M1 field metadata + program describer, M2 dataset summary + tips + history, M3 +-- the assembled @proposeInstructions@), all offline against the deterministic stub +-- LM. The headline check (M3) reads back the rendered prompt the final instruction +-- generator received and asserts it carries the dataset summary, the node's real +-- field names, and the selected tip — proving the proposer is genuinely grounded. +module ProposeSpec (tests) where + +import Control.Monad (forM) +import Data.IORef (newIORef, readIORef) +import Data.List (nub) +import Data.Text (Text) +import Data.Text qualified as T +import Effectful (Eff, IOE, runEff, (:>)) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval (Dataset, dataset, example) +import Shikumi.LLM (LLM) +import Shikumi.Optimize.Propose + ( NodeFieldNames (..), + PastInstruction (..), + ProposeRequest (..), + ProposeResult (..), + datasetSummary, + programFieldNames, + proposeInstructions, + renderHistory, + renderProgramPseudo, + tipAt, + tipBank, + ) +import Shikumi.Optimize.Propose.Summarize + ( ProgramDescribeIn (..), + ProgramDescribeOut (..), + programDescriber, + ) +import Shikumi.Program (Program, foldParams, runProgram) +import StubLM + ( Label (..), + Sentence (..), + runStubLM, + runStubLMCapturing, + sentimentProg, + ) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +-- --------------------------------------------------------------------------- +-- Fixtures and run helpers +-- --------------------------------------------------------------------------- + +trainset :: Dataset Sentence Label +trainset = + dataset + [ example (Sentence "good film") (Label "positive"), + example (Sentence "bad film") (Label "negative") + ] + +runStub :: Eff '[LLM, Error ShikumiError, IOE] a -> IO (Either ShikumiError a) +runStub act = runEff . runErrorNoCallStack @ShikumiError $ runStubLM act + +runCapturing :: Eff '[LLM, Error ShikumiError, IOE] a -> IO (Either ShikumiError a, [Text]) +runCapturing act = do + ref <- newIORef [] + res <- runEff . runErrorNoCallStack @ShikumiError $ runStubLMCapturing ref act + texts <- readIORef ref + pure (res, texts) + +tests :: TestTree +tests = testGroup "Propose" [m1, m2, m3] + +-- --------------------------------------------------------------------------- +-- M1 — field metadata + program describer +-- --------------------------------------------------------------------------- + +m1 :: TestTree +m1 = + testGroup + "Propose.M1" + [ testCase "programFieldNames recovers each node's input/output field names" $ + programFieldNames sentimentProg + @?= [NodeFieldNames {inputFieldNames = ["text"], outputFieldNames = ["sentiment"]}], + testCase "renderProgramPseudo renders the node deterministically" $ + renderProgramPseudo sentimentProg @?= "predict(text) -> sentiment", + testCase "programDescriber returns a description and the prompt carried the pseudo-code" $ do + let code = renderProgramPseudo sentimentProg + (res, texts) <- + runCapturing (runProgram programDescriber (ProgramDescribeIn code "good film => positive")) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (ProgramDescribeOut d) -> assertBool "non-empty description" (not (T.null d)) + assertBool + "the describer prompt carried the program pseudo-code" + (any (T.isInfixOf "predict(text)") texts) + ] + +-- --------------------------------------------------------------------------- +-- M2 — dataset summary, tips, instruction history +-- --------------------------------------------------------------------------- + +m2 :: TestTree +m2 = + testGroup + "Propose.M2" + [ testCase "datasetSummary returns a summary and the describer prompt saw a training row" $ do + (res, texts) <- runCapturing (datasetSummary 2 trainset) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right s -> assertBool "non-empty summary" (not (T.null s)) + assertBool + "a sampled training row reached the describer" + (any (T.isInfixOf "good film") texts), + testCase "tipAt selects deterministically and wraps" $ do + tipAt 1 @?= (tipBank !! 1) + tipAt (length tipBank) @?= tipAt 0, + testCase "renderHistory formats entries and empties" $ do + renderHistory 5 [PastInstruction "a" 0.5, PastInstruction "b" 0.9] + @?= "score 0.5 :: a\nscore 0.9 :: b" + renderHistory 5 [] @?= "No previous instructions." + ] + +-- --------------------------------------------------------------------------- +-- M3 — the assembled proposeInstructions surface +-- --------------------------------------------------------------------------- + +m3 :: TestTree +m3 = + testGroup + "Propose.M3" + [ testCase "proposeInstructions returns distinct candidates retaining the current one" $ do + let req = ProposeRequest sentimentProg 0 "Classify it." [] [] 3 0 2 + res <- runStub (proposeInstructions trainset req) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (ProposeResult cs) -> do + assertBool "current instruction retained" ("Classify it." `elem` cs) + assertBool "candidates are distinct" (cs == nub cs) + assertBool "more than one candidate" (length cs >= 2), + testCase "the grounded signals reach the instruction generator" $ do + let req = ProposeRequest sentimentProg 0 "Classify it." [] [] 3 0 2 + (res, texts) <- runCapturing (proposeInstructions trainset req) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right _ -> pure () + let genReqs = filter (T.isInfixOf "## proposedInstruction ##") texts + assertBool "at least one generator request" (not (null genReqs)) + assertBool + "dataset summary reached the generator" + (any (T.isInfixOf "A sentiment dataset") genReqs) + assertBool + "the node's field names reached the generator (via the module signature)" + (any (T.isInfixOf "predict(text) -> sentiment") genReqs) + assertBool + "the creative tip reached the generator" + (any (T.isInfixOf (tipAt 1)) genReqs), + testCase "the surface drives a MIPROv2-shaped caller" $ do + res <- runStub (miproShapedPropose trainset sentimentProg 3) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right perNode -> case perNode of + [node0] -> assertBool "the node got candidates" (not (null node0)) + _ -> assertFailure ("expected one node, got " <> show (length perNode)) + ] + +-- | A stand-in for MIPROv2's proposal stage: for each node, ask the grounded +-- proposer for @n@ candidates. Proves the shared contract is usable before MIPROv2 +-- exists — it compiles against 'proposeInstructions' exactly as EP-20 will. +miproShapedPropose :: + (LLM :> es, Error ShikumiError :> es) => + Dataset Sentence Label -> + Program Sentence Label -> + Int -> + Eff es [[Text]] +miproShapedPropose ds prog n = + forM [0 .. length (foldParams prog) - 1] $ \k -> + rankedCandidates <$> proposeInstructions ds (ProposeRequest prog k "" [] [] n 0 2) diff --git a/test/RandomSearchSpec.hs b/test/RandomSearchSpec.hs new file mode 100644 --- /dev/null +++ b/test/RandomSearchSpec.hs @@ -0,0 +1,108 @@ +{-# LANGUAGE TypeApplications #-} + +-- | EP-23 acceptance (random search): the best-of-N bootstrap held-out score is +-- @>=@ a single bootstrap run on the same fixture; the search is reproducible; and +-- the best compiled output round-trips through serialization. +module RandomSearchSpec (tests) where + +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +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.Optimize + ( bootstrapFewShot, + bootstrapRandomSearch, + defaultBudget, + optimize, + scoreOn, + setNodeInstr, + ) +import Shikumi.Program (Program) +import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, sentimentProg) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +-- --------------------------------------------------------------------------- +-- Fixtures and run helpers +-- --------------------------------------------------------------------------- + +trainset :: Dataset Sentence Label +trainset = + dataset + [ example (Sentence "good film") (Label "positive"), + example (Sentence "bad film") (Label "negative"), + example (Sentence "good show") (Label "positive"), + example (Sentence "bad play") (Label "negative") + ] + +heldout :: Dataset Sentence Label +heldout = + dataset + [ example (Sentence "good book") (Label "positive"), + example (Sentence "bad book") (Label "negative") + ] + +-- A teacher that answers correctly under the stub (a RULE-bearing instruction), so +-- bootstrap can harvest good demos from its passing runs. +teacher :: Program Sentence Label +teacher = setNodeInstr 0 ruleInstruction sentimentProg + +runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act + +tests :: TestTree +tests = testGroup "RandomSearch" [bestOfNAtLeastSingle, reproducible, roundTrips] + +bestOfNAtLeastSingle :: TestTree +bestOfNAtLeastSingle = + testCase "best-of-N held-out score >= single bootstrap" $ do + res <- + runStub $ do + single <- optimize (bootstrapFewShot teacher defaultBudget) trainset exactMatch sentimentProg + singleScore <- scoreOn heldout exactMatch (compiledProgram single) + rs <- optimize (bootstrapRandomSearch teacher 8 defaultBudget) trainset exactMatch sentimentProg + rsScore <- scoreOn heldout exactMatch (compiledProgram rs) + pure (singleScore, rsScore) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (singleScore, rsScore) -> + assertBool + ("random search " <> show rsScore <> " should be >= single bootstrap " <> show singleScore) + (rsScore >= singleScore) + +reproducible :: TestTree +reproducible = + testCase "bootstrapRandomSearch is reproducible" $ do + res <- + runStub $ do + a <- optimize (bootstrapRandomSearch teacher 8 defaultBudget) trainset exactMatch sentimentProg + sa <- scoreOn heldout exactMatch (compiledProgram a) + b <- optimize (bootstrapRandomSearch teacher 8 defaultBudget) trainset exactMatch sentimentProg + sb <- scoreOn heldout exactMatch (compiledProgram b) + pure (sa, sb) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (sa, sb) -> sa @?= sb + +roundTrips :: TestTree +roundTrips = + testCase "best compiled output round-trips, scoring identically" $ do + res <- + runStub $ do + cp <- optimize (bootstrapRandomSearch teacher 8 defaultBudget) trainset exactMatch sentimentProg + before <- scoreOn heldout exactMatch (compiledProgram cp) + case decodeCompiledOnto sentimentProg (encodeCompiled cp) of + Left err -> pure (Left err) + Right cp' -> do + after <- scoreOn heldout exactMatch (compiledProgram cp') + pure (Right (before, after)) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (Left err) -> assertFailure ("decode failed: " <> err) + Right (Right (before, after)) -> after @?= before diff --git a/test/StubLM.hs b/test/StubLM.hs new file mode 100644 --- /dev/null +++ b/test/StubLM.hs @@ -0,0 +1,394 @@ +{-# LANGUAGE LambdaCase #-} + +-- | Shared, network-free fixtures and a deterministic stub @LLM@ interpreter for +-- the EP-10 optimizer suite. +-- +-- The whole point of an optimizer test is that __changing a program's parameters +-- changes its score__ — otherwise an optimizer cannot demonstrably improve +-- anything. So the stub is not a constant: it inspects the rendered request +-- (system prompt + messages) and answers accordingly, by a rule that is monotone +-- in parameter quality. The task is binary sentiment classification of a +-- 'Sentence' into a 'Label' (@"positive"@ / @"negative"@). The ground-truth rule +-- ('goldLabel') is lexical: a sentence containing the word @good@ is positive, +-- one containing @bad@ is negative. +-- +-- The stub answers a sentiment request as follows: +-- +-- * __with demonstrations present__ — nearest-demo classification: echo the +-- label of the demo sentence sharing the most words with the input (ties and +-- zero-overlap fall back to the demos' majority label). So better demo sets +-- (covering both classes' keywords) yield higher scores — which is exactly +-- what labeled-few-shot and bootstrap exploit. +-- * __no demos but an instruction mentioning @RULE@__ — apply 'goldLabel' +-- directly. So an instruction that describes the task unlocks correct answers +-- — which is what instruction search exploits. +-- * __no demos and no @RULE@ instruction__ — answer the constant @"neutral"@, +-- which never matches a real label, so a deliberately-underspecified program +-- scores zero. This is the "before" state the acceptance test improves on. +-- +-- A second request shape, the instruction /proposer/ (recognised by its +-- @proposedInstruction@ output field), returns a candidate instruction chosen by +-- the @variant:N@ marker the optimizer embeds in the request: variant @0@ is the +-- "magic" @RULE@-bearing instruction, others are bland. This lets instruction +-- search demonstrably select the best proposal deterministically. +module StubLM + ( -- * Task records + Sentence (..), + Label (..), + + -- * Signature and program + sentimentSig, + sentimentProg, + + -- * Ground truth and helpers + goldLabel, + ruleInstruction, + blandInstruction, + + -- * The stub interpreters + runStubLM, + runStubLMCounting, + runStubLMCapturing, + + -- * The joint instruction×demo task (EP-20) + runJointStubLM, + runJointStubLMCounting, + + -- * The reflective task (EP-22, GEPA) + runGepaStubLM, + ) +where + +import Baikai + ( AssistantContent (..), + Context, + Response, + TextContent (..), + _Response, + _TextContent, + ) +import Baikai.Content (UserContent (..)) +import Baikai.Message (AssistantPayload (..), Message (..), UserPayload (..)) +import Control.Lens ((&), (.~), (^.)) +import Data.Aeson (ToJSON) +import Data.Generics.Labels () +import Data.IORef (IORef, modifyIORef') +import Data.Maybe (fromMaybe) +import Data.Set qualified as Set +import Data.Text (Text) +import Data.Text qualified as T +import Data.Vector qualified as V +import Effectful (Eff, IOE, liftIO, (:>)) +import Effectful.Dispatch.Dynamic (interpret) +import GHC.Generics (Generic) +import Shikumi.Adapter (ToPrompt) +import Shikumi.LLM (LLM (..)) +import Shikumi.Module (predict) +import Shikumi.Program (Program) +import Shikumi.Schema (FromModel, ToSchema, Validatable) +import Shikumi.Signature (Signature, mkSignature) + +-- --------------------------------------------------------------------------- +-- Task records +-- --------------------------------------------------------------------------- + +-- | A sentence to classify. +newtype Sentence = Sentence {text :: Text} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToJSON) + +instance ToSchema Sentence + +instance FromModel Sentence + +instance ToPrompt Sentence + +instance Validatable Sentence + +-- | A sentiment label (@"positive"@ / @"negative"@). +newtype Label = Label {sentiment :: Text} + deriving stock (Generic, Show, Eq, Ord) + deriving anyclass (ToJSON) + +instance ToSchema Label + +instance FromModel Label + +instance ToPrompt Label + +instance Validatable Label + +-- --------------------------------------------------------------------------- +-- Signature and program +-- --------------------------------------------------------------------------- + +-- | The deliberately-underspecified signature: an /empty/ instruction and no +-- demos. A program built from it scores zero under the stub until an optimizer +-- gives it useful parameters. +sentimentSig :: Signature Sentence Label +sentimentSig = mkSignature "" + +-- | The single-node sentiment program. +sentimentProg :: Program Sentence Label +sentimentProg = predict sentimentSig + +-- --------------------------------------------------------------------------- +-- Ground truth and helpers +-- --------------------------------------------------------------------------- + +-- | The ground-truth labelling rule: @good@ → positive, @bad@ → negative, +-- otherwise negative (the conservative default). +goldLabel :: Text -> Text +goldLabel s + | "good" `elem` ws = "positive" + | "bad" `elem` ws = "negative" + | otherwise = "negative" + where + ws = T.words (T.toLower s) + +-- | An instruction that unlocks correct answers under the stub (it mentions +-- @RULE@). The "magic" instruction the proposer offers as variant 0. +ruleInstruction :: Text +ruleInstruction = "RULE: classify the sentiment as positive or negative." + +-- | An instruction that does /not/ unlock correct answers (no @RULE@). +blandInstruction :: Text +blandInstruction = "Please answer." + +-- --------------------------------------------------------------------------- +-- The stub interpreters +-- --------------------------------------------------------------------------- + +-- | Interpret the @LLM@ effect deterministically and offline (see module header). +runStubLM :: Eff (LLM : es) a -> Eff es a +runStubLM = interpret $ \_ -> \case + Complete _ ctx _ -> pure (mkResponse (respondTo ctx)) + Stream {} -> pure [] + +-- | Like 'runStubLM' but increments a counter on every completion, for the +-- budget-enforcement test. +runStubLMCounting :: (IOE :> es) => IORef Int -> Eff (LLM : es) a -> Eff es a +runStubLMCounting ref = interpret $ \_ -> \case + Complete _ ctx _ -> do + liftIO (modifyIORef' ref (+ 1)) + pure (mkResponse (respondTo ctx)) + Stream {} -> pure [] + +-- | Like 'runStubLM' but records each request's full rendered text (system prompt + +-- user messages) into a test-readable 'IORef', so a test can assert which signals +-- (dataset summary, field names, tip) reached the model. +runStubLMCapturing :: (IOE :> es) => IORef [Text] -> Eff (LLM : es) a -> Eff es a +runStubLMCapturing ref = interpret $ \_ -> \case + Complete _ ctx _ -> do + liftIO (modifyIORef' ref (++ [fullRequestText ctx])) + pure (mkResponse (respondTo ctx)) + Stream {} -> pure [] + +-- | The joint instruction×demo task (EP-20). The held-out score is a strictly +-- increasing function of getting /both/ axes right: "region A" sentences (those with +-- @good@/@bad@) are answered correctly only when a @RULE@ instruction is present, and +-- "region B" sentences (e.g. @great@/@terrible@) only when a covering demo is present. +-- A @RULE@-only program (what instruction search reaches) scores 0.5; a joint program +-- (instruction + demos, what MIPROv2 reaches) scores 1.0. +runJointStubLM :: Eff (LLM : es) a -> Eff es a +runJointStubLM = interpret $ \_ -> \case + Complete _ ctx _ -> pure (mkResponse (respondWith answerJoint ctx)) + Stream {} -> pure [] + +-- | Like 'runJointStubLM' but counts completions, for the budget test. +runJointStubLMCounting :: (IOE :> es) => IORef Int -> Eff (LLM : es) a -> Eff es a +runJointStubLMCounting ref = interpret $ \_ -> \case + Complete _ ctx _ -> do + liftIO (modifyIORef' ref (+ 1)) + pure (mkResponse (respondWith answerJoint ctx)) + Stream {} -> pure [] + +-- | The joint task's classification rule (see 'runJointStubLM'). +answerJoint :: Context -> Text +answerJoint ctx + | "good" `elem` ws || "bad" `elem` ws = + -- region A: correct only with a RULE instruction + if instructionHasRule ctx then goldLabel s else "neutral" + | otherwise = + -- region B: correct only with a covering demo + let demos = demoPairs (V.toList (ctx ^. #messages)) + in if any (\(ds, _) -> overlap s ds > 0) demos then nnLabel s demos else "neutral" + where + s = parseSentence (lastUserText ctx) + ws = T.words (T.toLower s) + +-- | The full rendered request: system prompt followed by every user message's text. +fullRequestText :: Context -> Text +fullRequestText ctx = + T.intercalate "\n" (maybe [] pure (ctx ^. #systemPrompt) ++ allUserText ctx) + +-- | All user message texts in a request. +allUserText :: Context -> [Text] +allUserText ctx = [userPayloadText u | UserMessage u <- V.toList (ctx ^. #messages)] + +-- | Decide the response body for a request, with a pluggable sentiment-answer rule. +-- The grounded proposer (EP-19) issues several distinct sub-program requests; each is +-- recognised by the output-field marker its guide carries in the system prompt +-- (@[[ ## name ## ]]@). The @answerFn@ supplies the task-specific classification. +respondWith :: (Context -> Text) -> Context -> Text +respondWith answerFn ctx + | sysHasMarker "programDescription" ctx = + markerBody [("programDescription", "This program classifies the sentiment of a sentence.")] + | sysHasMarker "observations" ctx = + markerBody [("observations", "The rows pair short film reviews with positive or negative labels.")] + | sysHasMarker "summary" ctx = + markerBody [("summary", "A sentiment dataset of short film reviews labelled positive or negative.")] + | sysHasMarker "moduleDescription" ctx = + markerBody [("moduleDescription", "This module assigns a sentiment label to the input sentence.")] + | sysHasMarker "proposedInstruction" ctx = + markerBody [("proposedInstruction", groundedInstruction ctx)] + | otherwise = markerBody [("sentiment", answerFn ctx)] + +-- | Whether the request's system prompt carries a @[[ ## name ## ]]@ output marker. +sysHasMarker :: Text -> Context -> Bool +sysHasMarker name ctx = maybe False (T.isInfixOf ("## " <> name <> " ##")) (ctx ^. #systemPrompt) + +-- | The default sentiment task's response decision. +respondTo :: Context -> Text +respondTo = respondWith answerSentiment + +-- | GEPA's reflective task. Like the sentiment task, but the reflective proposer +-- (recognised by its @proposedInstruction@ output) returns the magic @RULE@ +-- instruction whenever the request's feedback asks the node to be more /specific/, +-- and a bland instruction otherwise. So a node whose captured critique is "be more +-- specific" gets mutated to the instruction that unlocks correct answers. +runGepaStubLM :: Eff (LLM : es) a -> Eff es a +runGepaStubLM = interpret $ \_ -> \case + Complete _ ctx _ -> pure (mkResponse (respondGepa ctx)) + Stream {} -> pure [] + +respondGepa :: Context -> Text +respondGepa ctx + | sysHasMarker "proposedInstruction" ctx = + markerBody + [ ( "proposedInstruction", + if T.isInfixOf "specific" (lastUserText ctx) then ruleInstruction else blandInstruction + ) + ] + | otherwise = markerBody [("sentiment", answerSentiment ctx)] + +-- | The grounded proposer's instruction choice: the "magic" @RULE@-bearing +-- instruction when the request carries the /creative/ tip (so exactly one candidate +-- per node proposes it), otherwise a bland one. +groundedInstruction :: Context -> Text +groundedInstruction ctx + | T.isInfixOf "creative" (lastUserText ctx) = ruleInstruction + | otherwise = blandInstruction + +-- | Classify the actual input given the demos and instruction in the context. +answerSentiment :: Context -> Text +answerSentiment ctx = + case demos of + [] -> if instructionHasRule ctx then goldLabel s else "neutral" + _ -> nnLabel s demos + where + msgs = V.toList (ctx ^. #messages) + demos = demoPairs msgs + s = parseSentence (lastUserText ctx) + +-- | Whether the system prompt carries a @RULE@-bearing instruction. +instructionHasRule :: Context -> Bool +instructionHasRule ctx = maybe False (T.isInfixOf "RULE") (ctx ^. #systemPrompt) + +-- | Nearest-demo classification: the label of the demo sentence with the largest +-- word overlap with @s@; ties and zero-overlap fall back to the majority label. +nnLabel :: Text -> [(Text, Text)] -> Text +nnLabel s demos + | bestOverlap == 0 = majorityLabel (map snd demos) + | otherwise = bestLabel + where + scored = [(overlap s ds, lbl) | (ds, lbl) <- demos] + (bestOverlap, bestLabel) = foldl1Keep scored + -- keep the earliest maximum (strictly-greater replaces) + foldl1Keep (x : xs) = foldl' (\b c -> if fst c > fst b then c else b) x xs + foldl1Keep [] = (0, "negative") + +-- | Word-overlap (intersection size) of two strings, case-insensitive. +overlap :: Text -> Text -> Int +overlap a b = Set.size (Set.intersection (wordSet a) (wordSet b)) + where + wordSet = Set.fromList . T.words . T.toLower + +-- | The most frequent label, ties broken by first appearance. +majorityLabel :: [Text] -> Text +majorityLabel ls = case [(l, length (filter (== l) ls)) | l <- nubOrd ls] of + [] -> "negative" + (t : ts) -> fst (foldl' pick t ts) + where + pick best cur = if snd cur > snd best then cur else best + +-- | Order-preserving de-duplication. +nubOrd :: [Text] -> [Text] +nubOrd = go Set.empty + where + go _ [] = [] + go seen (x : xs) + | x `Set.member` seen = go seen xs + | otherwise = x : go (Set.insert x seen) xs + +-- --------------------------------------------------------------------------- +-- Request parsing (reading the rendered context) +-- --------------------------------------------------------------------------- + +-- | The demonstration @(sentence, label)@ pairs in a request: the fallback +-- adapter renders each demo as a @user@ (input) then @assistant@ (output) +-- message, before the final @user@ message carrying the actual input. +demoPairs :: [Message] -> [(Text, Text)] +demoPairs (UserMessage u : AssistantMessage a : rest) = + (parseSentence (userPayloadText u), parseLabel (assistantPayloadText a)) + : demoPairs rest +demoPairs _ = [] + +-- | The text of the last user message — the actual input being classified. +lastUserText :: Context -> Text +lastUserText ctx = + case [userPayloadText u | UserMessage u <- V.toList (ctx ^. #messages)] of + [] -> "" + xs -> last xs + +-- | Concatenate the text blocks of a user message. +userPayloadText :: UserPayload -> Text +userPayloadText u = T.concat [t | UserText (TextContent t) <- V.toList (u ^. #content)] + +-- | Concatenate the text blocks of an assistant message. +assistantPayloadText :: AssistantPayload -> Text +assistantPayloadText a = T.concat [t | AssistantText (TextContent t) <- V.toList (a ^. #content)] + +-- | Strip the @"text:"@ field label that @toPrompt@ renders for a 'Sentence'. +parseSentence :: Text -> Text +parseSentence t = T.strip (fromMaybe stripped (T.stripPrefix "text:" stripped)) + where + stripped = T.strip t + +-- | Read the @sentiment@ marker section out of a rendered demo output. +parseLabel :: Text -> Text +parseLabel = markerValue "sentiment" + +-- | The value line following a @[[ ## name ## ]]@ marker, trimmed. +markerValue :: Text -> Text -> Text +markerValue name body = + case dropWhile (not . isMarker) (T.lines body) of + (_ : v : _) -> T.strip v + _ -> "" + where + isMarker l = T.strip l == "[[ ## " <> name <> " ## ]]" + +-- --------------------------------------------------------------------------- +-- Response construction +-- --------------------------------------------------------------------------- + +-- | Build a fallback-style @[[ ## field ## ]]@ response body. +markerBody :: [(Text, Text)] -> Text +markerBody fields = T.unlines (concatMap sect fields ++ ["[[ ## completed ## ]]"]) + where + sect (k, v) = ["[[ ## " <> k <> " ## ]]", v] + +-- | An assistant 'Response' carrying @body@ as its single text block. +mkResponse :: Text -> Response +mkResponse body = + _Response & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ body))