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
-- normally changes parameters while preserving boundary types.
--
-- The additive 'Shikumi.Optimize.Structure.structureSearchWith' API selects among
-- caller-registered typed implementations using shared validation and admission;
-- its experimental structure artifacts restore through that same registry.
-- Two legacy optimizers are also explicit structure-changing exceptions. 'Shikumi.Optimize.KNN.knnFewShot'
-- returns an @Embed@ wrapper that selects demos at run time from an opaque closure;
-- persist the underlying student or use 'Shikumi.Optimize.KNN.knnFewShotCentroid'
-- when a plain parameter artifact is required. 'Shikumi.Optimize.Ensemble.ensembleSearch'
-- 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 (..),
    ConfiguredOptimizer (..),
    fromLegacyOptimizer,
    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.Optimize.Execution (SearchSession, markLegacy, sessionStopped)
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)

-- | A strategy receives an explicit, shared execution session.
newtype ConfiguredOptimizer i o = ConfiguredOptimizer
  { runConfiguredOptimizer ::
      forall es.
      (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) =>
      SearchSession es -> Dataset i o -> Metric o -> Program i o -> Eff es (CompiledProgram i o)
  }

fromLegacyOptimizer :: Optimizer i o -> ConfiguredOptimizer i o
fromLegacyOptimizer opt = ConfiguredOptimizer $ \session ds metric prog -> do
  halted <- sessionStopped session
  if halted
    then pure (CompiledProgram prog)
    else do
      markLegacy session
      runOptimizer opt ds metric prog