Types.hs

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