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 the shikumi core runtime.
+- Typed program and signature APIs, structured schema decoding, model routing, retries, budget tracking, multimodal helpers, streaming, refinement, rewards, and program combinators.
diff --git a/shikumi.cabal b/shikumi.cabal
new file mode 100644
--- /dev/null
+++ b/shikumi.cabal
@@ -0,0 +1,131 @@
+cabal-version: 3.4
+name: shikumi
+version: 0.1.0.0
+synopsis: Typed, structured, evaluable LM programs over baikai
+category: AI
+description:
+ Shikumi is a Haskell-native framework for building typed language-model programs.
+ This package provides the runtime substrate: the LLM effect over baikai, the
+ shikumi error type, and resilience (retries, rate limiting, budget control).
+
+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.Adapter
+ Shikumi.Combinator
+ Shikumi.Effect.Time
+ Shikumi.Error
+ Shikumi.LLM
+ Shikumi.LLM.Budget
+ Shikumi.Module
+ Shikumi.Multimodal
+ Shikumi.Prelude
+ Shikumi.Program
+ Shikumi.Refine
+ Shikumi.Reward
+ Shikumi.Routing
+ Shikumi.Schema
+ Shikumi.Schema.Types
+ Shikumi.Signature
+ Shikumi.Stream
+
+ build-depends:
+ , aeson
+ , baikai
+ , baikai-claude
+ , baikai-effectful
+ , baikai-openai
+ , base >=4.20 && <5
+ , base64-bytestring
+ , bytestring
+ , containers
+ , effectful
+ , filepath
+ , generic-lens
+ , lens ^>=5.3
+ , scientific
+ , stm
+ , text ^>=2.1
+ , time
+ , vector
+
+test-suite shikumi-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:
+ AdapterSpec
+ CombinatorSpec
+ ConstraintSpec
+ EndToEndSpec
+ ErrorSpec
+ Fixtures
+ LiveSpec
+ LLMSpec
+ ModuleSpec
+ MultimodalAdapterSpec
+ MultimodalEndToEndSpec
+ MultimodalSpec
+ ProgramAcceptanceSpec
+ ProgramFixtures
+ ProgramSpec
+ RefineSpec
+ RefineStub
+ ResilienceSpec
+ RoutingSpec
+ SchemaSpec
+ SerializeSpec
+ Shikumi.Effect.TimeSpec
+ Shikumi.LLM.Mock
+ SignatureSpec
+ StreamSpec
+ StubProvider
+ TwoStepSpec
+ XmlAdapterSpec
+
+ build-depends:
+ , aeson
+ , baikai
+ , baikai-claude
+ , baikai-effectful
+ , baikai-openai
+ , base
+ , base64-bytestring
+ , bytestring
+ , containers
+ , directory
+ , effectful
+ , generic-lens
+ , lens
+ , QuickCheck
+ , shikumi ^>=0.1.0.0
+ , stm
+ , streamly-core
+ , tasty
+ , tasty-hunit
+ , tasty-quickcheck
+ , text
+ , vector
diff --git a/src/Shikumi/Adapter.hs b/src/Shikumi/Adapter.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Adapter.hs
@@ -0,0 +1,471 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE DefaultSignatures #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE UndecidableInstances #-}
+
+-- | The 'Adapter' seam between a typed 'Signature' + input and the wire.
+-- @render@ builds a baikai @Context@+@Options@; @parse@ decodes a baikai
+-- @Response@ into the typed output via "Shikumi.Schema".
+--
+-- Two adapters ship. The native-schema adapter is the reliable path (the provider
+-- enforces the JSON schema); the prompt-based fallback renders @[[ ## field ## ]]@
+-- sections and re-parses them, for models without native structured output.
+-- 'capabilityFor' selects per model.
+--
+-- Native structured output is wired through a private /metadata channel/ (EP-14).
+-- Because a 'Program' renders before the ambient model is known (the model is
+-- supplied by an interpreter below 'Shikumi.Program.runProgram' — see
+-- "Shikumi.Routing"), 'render' cannot itself decide native-vs-fallback or build the
+-- final @responseFormat@. Instead, 'attachSchema' stamps the derived JSON schema
+-- under the reserved 'metaResponseSchemaKey' on @Options.metadata@, and
+-- 'stampTemperature' stamps a per-sample temperature under 'metaTemperatureKey';
+-- the router ("Shikumi.Routing".@routeLLM@) translates those into
+-- @responseFormat@\/@temperature@ against the real model and strips the keys before
+-- transport.
+module Shikumi.Adapter
+ ( -- * Input rendering
+ ToPrompt (..),
+
+ -- * The seam
+ Adapter (..),
+ ModelCapability (..),
+ capabilityFor,
+ nativeAdapter,
+ fallbackAdapter,
+ xmlAdapter,
+ adapterFor,
+ attachSchema,
+ responseText,
+
+ -- * The private request-metadata channel (consumed by "Shikumi.Routing")
+ stampTemperature,
+ metaResponseSchemaKey,
+ metaTemperatureKey,
+ )
+where
+
+import Baikai
+ ( Api (..),
+ AssistantContent (..),
+ Context,
+ Message,
+ Model,
+ Options,
+ Response,
+ TextContent (..),
+ assistant,
+ flattenAssistantBlocks,
+ user,
+ userImage,
+ _Context,
+ _Options,
+ )
+import Control.Lens ((&), (.~), (^.))
+import Data.Aeson (Object, Value (..), eitherDecodeStrict, toJSON)
+import Data.Aeson.Key qualified as Key
+import Data.Aeson.KeyMap qualified as KM
+import Data.Generics.Labels ()
+import Data.Kind (Type)
+import Data.Map.Strict (Map)
+import Data.Map.Strict qualified as Map
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Text.Encoding (encodeUtf8)
+import Data.Vector qualified as V
+import GHC.Generics
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.Multimodal (GImageFieldNames (..), GImageFields (..), Image, imageToContent)
+import Shikumi.Schema (FromModel, ToSchema, Validatable, deriveSchema, fromModelChecked)
+import Shikumi.Schema.Types (Constrained (..), FieldMeta (..))
+import Shikumi.Signature (Demo (..), Signature, getDemos, getInstruction, outputFields)
+
+-- ---------------------------------------------------------------------------
+-- ToPrompt: render an input record (and demo outputs) as labeled text
+-- ---------------------------------------------------------------------------
+
+-- | Render a record as labeled text for the prompt. 'toPromptFields' yields the
+-- @(fieldName, valueText)@ pairs; 'toPrompt' joins them as @"name: value"@ lines.
+class ToPrompt a where
+ toPromptFields :: a -> [(Text, Text)]
+ default toPromptFields :: (Generic a, GToPromptFields (Rep a)) => a -> [(Text, Text)]
+ toPromptFields = gToPromptFields . from
+
+ toPrompt :: a -> Text
+ toPrompt = T.intercalate "\n" . map (\(k, v) -> k <> ": " <> v) . toPromptFields
+
+ -- | The 'Image' values an /input/ record carries, in field order (EP-24). The
+ -- adapter lowers these to baikai @UserImage@ blocks in 'userTurn'. The generic
+ -- default walks the record, so any 'Generic' input gets it for free; a record
+ -- with no image field yields @[]@ and renders byte-for-byte as before. A
+ -- hand-written instance whose fields are polymorphic (e.g.
+ -- "Shikumi.Module".@WithReasoning@) cannot be walked generically and overrides
+ -- this to @[]@.
+ imageFields :: a -> [Image]
+ default imageFields :: (Generic a, GImageFields (Rep a)) => a -> [Image]
+ imageFields = gImageFields . from
+
+ -- | The names of an input record's image fields, in field order (EP-24). The
+ -- text renderer in 'userTurn' drops these so an image is not also rendered as
+ -- prose. Same defaulting story as 'imageFields'.
+ imageFieldNames :: a -> [Text]
+ default imageFieldNames :: (Generic a, GImageFieldNames (Rep a)) => a -> [Text]
+ imageFieldNames = gImageFieldNames . from
+
+class GToPromptFields (f :: Type -> Type) where
+ gToPromptFields :: f p -> [(Text, Text)]
+
+instance (GToPromptFields cs) => GToPromptFields (D1 d cs) where
+ gToPromptFields (M1 x) = gToPromptFields x
+
+instance (GToPromptFields cs) => GToPromptFields (C1 c cs) where
+ gToPromptFields (M1 x) = gToPromptFields x
+
+instance (GToPromptFields a, GToPromptFields b) => GToPromptFields (a :*: b) where
+ gToPromptFields (a :*: b) = gToPromptFields a ++ gToPromptFields b
+
+instance (Selector s, PromptValue t) => GToPromptFields (S1 s (K1 R t)) where
+ gToPromptFields m@(M1 (K1 v)) = [(T.pack (selName m), promptValue v)]
+
+-- | Render a leaf value to text. Text stays as-is; lists join with commas;
+-- 'Maybe' renders its contents or @""@; a 'Shikumi.Schema.Types.Field' unwraps;
+-- everything else falls back to 'Show'.
+class PromptValue a where
+ promptValue :: a -> Text
+
+instance {-# OVERLAPPABLE #-} (Show a) => PromptValue a where
+ promptValue = T.pack . show
+
+instance PromptValue Text where
+ promptValue = id
+
+instance (PromptValue a) => PromptValue [a] where
+ promptValue = T.intercalate ", " . map promptValue
+
+instance (PromptValue a) => PromptValue (Maybe a) where
+ promptValue = maybe "" promptValue
+
+-- | A 'Constrained' field renders as its inner value (the constraint is a
+-- compile-time/decode-time concern, not a prompt concern). EP-26.
+instance (PromptValue a) => PromptValue (Constrained cs a) where
+ promptValue = promptValue . unConstrained
+
+-- ---------------------------------------------------------------------------
+-- The seam
+-- ---------------------------------------------------------------------------
+
+-- | A record of two functions: format a request, parse a response.
+data Adapter i o = Adapter
+ { render :: Signature i o -> i -> (Context, Options),
+ parse :: Signature i o -> Response -> Either ShikumiError o
+ }
+
+-- | Whether a model supports provider-native structured output.
+data ModelCapability = NativeSchema | PromptFallback
+ deriving stock (Eq, Show)
+
+-- | A pure capability check over a baikai 'Model'. OpenAI/Anthropic on their
+-- non-CLI APIs are native-capable; CLI APIs and unknown @Custom@ hosts use the
+-- fallback. Refine as more models gain native support.
+capabilityFor :: Model -> ModelCapability
+capabilityFor m = case (m ^. #provider, m ^. #api) of
+ ("openai", OpenAIChatCompletions) -> NativeSchema
+ ("anthropic", AnthropicMessages) -> NativeSchema
+ _ -> PromptFallback
+
+-- | Select the adapter for a model from its capability.
+adapterFor ::
+ forall i o.
+ (ToSchema o, FromModel o, Validatable o, ToPrompt i, ToPrompt o) =>
+ Model ->
+ Adapter i o
+adapterFor m = case capabilityFor m of
+ NativeSchema -> nativeAdapter
+ PromptFallback -> fallbackAdapter
+
+-- | The reserved 'Options.metadata' key under which 'attachSchema' stamps the
+-- derived JSON schema 'Value'. "Shikumi.Routing".@routeLLM@ reads it, turns it into
+-- @Options.responseFormat@ for native-capable models, and strips it before
+-- transport. Defined here (the lowest module both the adapter and the router share)
+-- so it is written and read against a single source of truth.
+metaResponseSchemaKey :: Text
+metaResponseSchemaKey = "shikumi.responseSchema"
+
+-- | The reserved 'Options.metadata' key under which 'stampTemperature' stamps a
+-- per-sample temperature (a JSON number). The router reads it, sets
+-- @Options.temperature@, and strips it before transport.
+metaTemperatureKey :: Text
+metaTemperatureKey = "shikumi.temperature"
+
+-- | Stamp the derived JSON schema onto a request's private metadata channel. This
+-- is no longer a no-op: it records the schema under 'metaResponseSchemaKey' so the
+-- router can attach a real @responseFormat@ once the ambient model is known. (The
+-- router only honours it for native-capable models; for fallback models it is
+-- simply stripped.)
+attachSchema :: Value -> Options -> Options
+attachSchema schema opts =
+ opts & #metadata .~ Map.insert metaResponseSchemaKey schema (opts ^. #metadata)
+
+-- | Stamp a per-sample temperature onto a request's private metadata channel under
+-- 'metaTemperatureKey'. Used by 'Shikumi.Program.runProgram' to thread a
+-- 'Shikumi.Program.MajorityVote' sample's temperature down to its 'Predict' nodes;
+-- the router turns it into @Options.temperature@.
+stampTemperature :: Double -> Options -> Options
+stampTemperature t opts =
+ opts & #metadata .~ Map.insert metaTemperatureKey (toJSON t) (opts ^. #metadata)
+
+-- | The native-schema adapter. @render@ stamps the derived schema onto the request
+-- metadata channel (see 'attachSchema'); @parse@ reads the structured JSON from the
+-- assistant text.
+nativeAdapter ::
+ forall i o.
+ (ToSchema o, FromModel o, Validatable o, ToPrompt i, ToPrompt o) =>
+ Adapter i o
+nativeAdapter =
+ Adapter
+ { render = \sig i ->
+ let sys = systemHeader sig <> nativeOutputGuide sig
+ ctx = buildContext sys (demoMessages sig ++ [userTurn i])
+ opts = attachSchema (deriveSchema @o) _Options
+ in (ctx, opts),
+ parse = \_sig resp -> assistantJSON resp >>= fromModelChecked
+ }
+
+-- | The prompt-based fallback adapter. @render@ asks for @[[ ## field ## ]]@
+-- sections; @parse@ splits them, coerces each to its schema type, and decodes.
+fallbackAdapter ::
+ forall i o.
+ (ToSchema o, FromModel o, Validatable o, ToPrompt i, ToPrompt o) =>
+ Adapter i o
+fallbackAdapter =
+ Adapter
+ { render = \sig i ->
+ let sys = systemHeader sig <> fallbackOutputGuide sig
+ ctx = buildContext sys (demoMessages sig ++ [userTurn i])
+ in (ctx, _Options),
+ parse = \_sig resp ->
+ let sections = parseMarkers (responseText resp)
+ obj = sectionsToObject (deriveSchema @o) sections
+ in fromModelChecked obj
+ }
+
+-- | The XML adapter (EP-26). A third wire format on the same typed seam: @render@
+-- asks the model to wrap each output field in @\…\@ tags, and
+-- @parse@ reads those tags back. Some models follow an XML shape more reliably than
+-- JSON or the @[[ ## … ## ]]@ markers. Opt-in — a caller selects it explicitly;
+-- 'adapterFor' does not auto-select it (XML is a caller choice, not a detectable
+-- model capability). Reuses the same 'sectionsToObject' + 'fromModelChecked'
+-- decode path as 'fallbackAdapter', so nested records and lists in tags coerce the
+-- same way.
+xmlAdapter ::
+ forall i o.
+ (ToSchema o, FromModel o, Validatable o, ToPrompt i, ToPrompt o) =>
+ Adapter i o
+xmlAdapter =
+ Adapter
+ { render = \sig i ->
+ let sys = systemHeader sig <> xmlOutputGuide sig
+ ctx = buildContext sys (xmlDemoMessages sig ++ [userTurn i])
+ in (ctx, _Options),
+ parse = \sig resp ->
+ let names = map fieldName (outputFields sig)
+ sections = parseXmlTags names (responseText resp)
+ obj = sectionsToObject (deriveSchema @o) sections
+ in fromModelChecked obj
+ }
+
+-- ---------------------------------------------------------------------------
+-- Rendering helpers
+-- ---------------------------------------------------------------------------
+
+buildContext :: Text -> [Message] -> Context
+buildContext sys msgs =
+ _Context & #systemPrompt .~ Just sys & #messages .~ V.fromList msgs
+
+-- | Build the final user turn for an input. If the input has an image field, it is
+-- lowered to a baikai 'userImage' block, with the remaining (text) fields rendered
+-- as a leading text block and the image field name dropped from that text. With no
+-- image field this is exactly @user (toPrompt i)@, so the all-text path is
+-- byte-for-byte unchanged (the EP-26 coexistence + regression invariant).
+--
+-- Only the /first/ image is attached: baikai's 'userImage' carries one image block,
+-- and the headline use case is one image per input. A multi-image input is future
+-- work (it would assemble the @UserPayload@ content vector by hand).
+userTurn :: forall i. (ToPrompt i) => i -> Message
+userTurn i = case imageFields i of
+ [] -> user (toPrompt i)
+ (img : _) ->
+ let dropped = imageFieldNames i
+ textBody =
+ T.intercalate "\n" [k <> ": " <> v | (k, v) <- toPromptFields i, k `notElem` dropped]
+ prefix = if T.null textBody then Nothing else Just textBody
+ in userImage (imageToContent img) prefix
+
+systemHeader :: Signature i o -> Text
+systemHeader sig = getInstruction sig <> "\n\n"
+
+-- | A native-output guide: list the output fields with their descriptions.
+nativeOutputGuide :: Signature i o -> Text
+nativeOutputGuide sig =
+ "Reply with a JSON object containing these fields:\n"
+ <> T.unlines [describeField f | f <- outputFields sig]
+
+-- | A fallback-output guide: ask for one @[[ ## field ## ]]@ section per output
+-- field, then a final @[[ ## completed ## ]]@ marker (DSPy's convention).
+fallbackOutputGuide :: Signature i o -> Text
+fallbackOutputGuide sig =
+ "Reply using these sections, each marker on its own line:\n"
+ <> T.unlines [marker (fieldName f) <> describeSuffix f | f <- outputFields sig]
+ <> marker "completed"
+
+describeField :: FieldMeta -> Text
+describeField f = "- " <> fieldName f <> maybe "" (": " <>) (fieldDesc f)
+
+describeSuffix :: FieldMeta -> Text
+describeSuffix f = maybe "" (" -- " <>) (fieldDesc f)
+
+marker :: Text -> Text
+marker name = "[[ ## " <> name <> " ## ]]"
+
+-- | Render the demos as user/assistant message pairs.
+demoMessages :: (ToPrompt i, ToPrompt o) => Signature i o -> [Message]
+demoMessages sig = concatMap one (getDemos sig)
+ where
+ one (Demo i o) = [user (toPrompt i), assistant (renderOutputSections o)]
+
+-- | Render a demo output as @[[ ## field ## ]]@ sections (shared by both adapters
+-- for demo presentation).
+renderOutputSections :: (ToPrompt o) => o -> Text
+renderOutputSections o =
+ T.unlines [marker k <> "\n" <> v | (k, v) <- toPromptFields o] <> marker "completed"
+
+-- | An XML-output guide: ask for one @\…\@ element per output
+-- field. Mirrors 'fallbackOutputGuide' with XML tags instead of markers.
+xmlOutputGuide :: Signature i o -> Text
+xmlOutputGuide sig =
+ "Reply with each output field wrapped in an XML tag, on its own lines:\n"
+ <> T.unlines [openTag (fieldName f) <> "…" <> closeTag (fieldName f) <> describeSuffix f | f <- outputFields sig]
+
+-- | An XML open tag, @\@.
+openTag :: Text -> Text
+openTag name = "<" <> name <> ">"
+
+-- | An XML close tag, @\@.
+closeTag :: Text -> Text
+closeTag name = " name <> ">"
+
+-- | Render a demo output as @\…\@ elements (the XML adapter's demo
+-- shape, mirroring 'renderOutputSections' for the marker adapters).
+renderOutputXml :: (ToPrompt o) => o -> Text
+renderOutputXml o =
+ T.unlines [openTag k <> "\n" <> v <> "\n" <> closeTag k | (k, v) <- toPromptFields o]
+
+-- | Render the demos as user/assistant message pairs, the assistant turn shaped as
+-- XML so the model sees an example consistent with the XML adapter's wire format.
+xmlDemoMessages :: (ToPrompt i, ToPrompt o) => Signature i o -> [Message]
+xmlDemoMessages sig = concatMap one (getDemos sig)
+ where
+ one (Demo i o) = [user (toPrompt i), assistant (renderOutputXml o)]
+
+-- ---------------------------------------------------------------------------
+-- Parsing helpers (fallback path)
+-- ---------------------------------------------------------------------------
+
+-- | Concatenate the text of every assistant text block.
+responseText :: Response -> Text
+responseText resp =
+ T.concat [t | AssistantText (TextContent t) <- V.toList (flattenAssistantBlocks resp)]
+
+-- | Read the assistant text and parse it as a JSON value (native path).
+assistantJSON :: Response -> Either ShikumiError Value
+assistantJSON resp =
+ case eitherDecodeStrict (encodeUtf8 (responseText resp)) of
+ Left e -> Left (InvalidJSON (T.pack e))
+ Right v -> Right v
+
+-- | Split a @[[ ## name ## ]]@-delimited body into a map of section name to its
+-- (trimmed) text. The @completed@ marker carries no content.
+parseMarkers :: Text -> Map Text Text
+parseMarkers body = go (T.lines body) Nothing Map.empty
+ where
+ go [] cur acc = flush cur acc
+ go (l : ls) cur acc = case markerName l of
+ Just name -> go ls (Just (name, [])) (flush cur acc)
+ Nothing -> case cur of
+ Just (name, buf) -> go ls (Just (name, buf ++ [l])) acc
+ Nothing -> go ls Nothing acc
+ flush Nothing acc = acc
+ flush (Just (name, buf)) acc
+ | name == "completed" = acc
+ | otherwise = Map.insert name (T.strip (T.unlines buf)) acc
+
+-- | Extract @\…\@ sections into a name->text map. Only names that
+-- appear as output fields are kept (so stray tags are ignored, DSPy parity), and
+-- the first match per name wins. The content is whatever lies between the first
+-- @\@ and its next @\@ — a non-greedy match, the same as DSPy's
+-- @\<(?P\\\w+)\>(?P\.*?)\@ with DOTALL.
+parseXmlTags :: [Text] -> Text -> Map Text Text
+parseXmlTags names body =
+ Map.fromList [(nm, inner) | nm <- names, Just inner <- [extractTag nm body]]
+
+-- | The text between the first @\@ and its next @\@, trimmed;
+-- 'Nothing' if either tag is absent.
+extractTag :: Text -> Text -> Maybe Text
+extractTag nm body =
+ let open = openTag nm
+ close = closeTag nm
+ (_, afterOpen) = T.breakOn open body
+ in if T.null afterOpen
+ then Nothing
+ else
+ let rest = T.drop (T.length open) afterOpen
+ (inner, afterClose) = T.breakOn close rest
+ in if T.null afterClose then Nothing else Just (T.strip inner)
+
+-- | Recognize a @[[ ## name ## ]]@ marker line.
+markerName :: Text -> Maybe Text
+markerName line = do
+ a <- T.stripPrefix "[[ ## " (T.strip line)
+ b <- T.stripSuffix " ## ]]" a
+ pure (T.strip b)
+
+-- | Assemble a JSON object from marker sections, coercing each section to its
+-- output-schema type. Missing markers are simply absent (so a required field
+-- yields 'MissingField' downstream).
+sectionsToObject :: Value -> Map Text Text -> Value
+sectionsToObject schema sections =
+ Object (KM.fromList [(Key.fromText nm, sectionToValue (propSchema nm) raw) | (nm, raw) <- Map.toList sections])
+ where
+ props = schemaProps schema
+ propSchema nm = KM.lookup (Key.fromText nm) props
+
+-- | The @properties@ object of a record schema (empty if not a record schema).
+schemaProps :: Value -> Object
+schemaProps (Object o) = case KM.lookup "properties" o of
+ Just (Object p) -> p
+ _ -> KM.empty
+schemaProps _ = KM.empty
+
+-- | Coerce a raw section to a JSON value using its property schema as a guide:
+-- literal @null@ becomes JSON null; a string/enum-typed field stays a string;
+-- everything else is JSON-parsed (falling back to a string on parse failure).
+sectionToValue :: Maybe Value -> Text -> Value
+sectionToValue mschema raw
+ | T.strip raw == "null" = Null
+ | maybe False isStringLike mschema = String stripped
+ | otherwise = case eitherDecodeStrict (encodeUtf8 stripped) of
+ Right v -> v
+ Left _ -> String stripped
+ where
+ stripped = T.strip raw
+
+-- | Whether a property schema describes a string-like value (plain string or an
+-- enum, including a nullable string via @anyOf@).
+isStringLike :: Value -> Bool
+isStringLike (Object o) =
+ KM.lookup "type" o == Just (String "string")
+ || KM.member "enum" o
+ || case KM.lookup "anyOf" o of
+ Just (Array alts) -> any isStringLike (V.toList alts)
+ _ -> False
+isStringLike _ = False
diff --git a/src/Shikumi/Combinator.hs b/src/Shikumi/Combinator.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Combinator.hs
@@ -0,0 +1,169 @@
+-- | The ergonomic combinator surface (EP-5): operators and smart constructors
+-- that assemble small typed 'Program's into larger ones. The runtime behaviour
+-- lives in the GADT constructors owned by "Shikumi.Program" (so every node stays
+-- runnable, traversable by the optimizer, and serializable); this module is the
+-- thin, user-facing layer over them.
+--
+-- The combinators, by group:
+--
+-- * __Pipeline__ — '>>>' (binary sequential composition) and 'chain' (n-ary,
+-- same-type stages). Both are surface over "Shikumi.Program"'s @Compose@.
+-- * __Map__ — 'mapP' (bounded-concurrent) \/ 'mapSeqP' (sequential), running a
+-- per-element program across a list.
+-- * __Parallel__ — 'parallel2' (a pair on one input) and 'parallelN' (a
+-- homogeneous list on one input).
+-- * __Retry__ — 'retry' (any error) and 'retryWhen' (a chosen class of error).
+-- * __Validate__ — 'validate' (predicate + reason) and 'validateRetry'
+-- (re-run on rejection).
+-- * __MajorityVote__ — 'majorityVote' (modal under 'Eq') and 'majorityVoteBy'
+-- (a custom reducer).
+-- * __Ensemble__ — 'ensemble' (run distinct members, fold their results).
+--
+-- Concurrency is an /execution choice/, not part of a program's type: a program
+-- built here runs sequentially under 'Shikumi.Program.runProgram' and
+-- concurrently (honouring 'mapP' widths) under
+-- 'Shikumi.Program.runProgramConc'. See the plan's Decision Log for why
+-- 'Concurrent' is kept off @runProgram@'s constraint (MasterPlan integration
+-- point #4).
+module Shikumi.Combinator
+ ( -- * Pipeline
+ (>>>),
+ chain,
+
+ -- * Map
+ mapP,
+ mapSeqP,
+
+ -- * Parallel
+ parallel2,
+ parallelN,
+
+ -- * Retry
+ retry,
+ retryWhen,
+
+ -- * Validate
+ validate,
+ validateRetry,
+
+ -- * MajorityVote
+ majorityVote,
+ majorityVoteBy,
+ TempSchedule (..),
+
+ -- * Ensemble
+ ensemble,
+ )
+where
+
+import Data.Text (Text)
+import Shikumi.Error (ShikumiError)
+import Shikumi.Program
+ ( Program (Compose, Ensemble, MajorityVote, Map, Parallel, Retry, RetryWhen, Validate),
+ TempSchedule (..),
+ )
+
+-- ---------------------------------------------------------------------------
+-- Pipeline
+-- ---------------------------------------------------------------------------
+
+-- | Sequence two programs left-to-right: @p '>>>' q@ runs @p@, then feeds its
+-- output to @q@. Typechecks only when @p@'s output type equals @q@'s input type,
+-- so a mismatched pipeline is a compile error.
+(>>>) :: Program a b -> Program b c -> Program a c
+(>>>) = Compose
+
+infixr 1 >>>
+
+-- | Compose a non-empty list of same-type stages into one program, left-to-right.
+-- @'chain' [p, q, r]@ is @p '>>>' q '>>>' r@. (EP-4 owns the /binary/ typed
+-- @pipeline@\/@Compose@; this n-ary helper is restricted to a single carried type
+-- because the GADT has no identity node to seed an empty fold.)
+--
+-- Applying it to the empty list is a programmer error (there is no identity
+-- 'Program').
+chain :: [Program a a] -> Program a a
+chain [] = error "Shikumi.Combinator.chain: empty stage list"
+chain ps = foldr1 (>>>) ps
+
+-- ---------------------------------------------------------------------------
+-- Map
+-- ---------------------------------------------------------------------------
+
+-- | Apply a program to every element of a list, with bounded concurrency width
+-- @w@ honoured by 'Shikumi.Program.runProgramConc' (sequential under
+-- 'Shikumi.Program.runProgram' regardless). Output order matches input order.
+mapP :: Int -> Program a b -> Program [a] [b]
+mapP = Map
+
+-- | Apply a program to every element of a list, strictly sequentially (width 1).
+mapSeqP :: Program a b -> Program [a] [b]
+mapSeqP = Map 1
+
+-- ---------------------------------------------------------------------------
+-- Parallel
+-- ---------------------------------------------------------------------------
+
+-- | Run two programs on the /same/ input and pair their outputs. Concurrent
+-- under 'Shikumi.Program.runProgramConc'.
+parallel2 :: Program i a -> Program i b -> Program i (a, b)
+parallel2 = Parallel
+
+-- | Run a homogeneous list of programs on the same input and collect their
+-- outputs in order. Defined as an 'ensemble' with the identity reducer.
+parallelN :: [Program i o] -> Program i [o]
+parallelN ps = Ensemble ps id
+
+-- ---------------------------------------------------------------------------
+-- Retry
+-- ---------------------------------------------------------------------------
+
+-- | Re-run a program up to @n@ total attempts, returning the first success or
+-- the last error. Retries on any 'ShikumiError'.
+retry :: Int -> Program i o -> Program i o
+retry = Retry
+
+-- | Like 'retry', but only retries errors matching the predicate; a non-matching
+-- error propagates immediately (after a single attempt).
+retryWhen :: (ShikumiError -> Bool) -> Int -> Program i o -> Program i o
+retryWhen = RetryWhen
+
+-- ---------------------------------------------------------------------------
+-- Validate
+-- ---------------------------------------------------------------------------
+
+-- | Run a program, then check its output against a predicate; on rejection
+-- surface a 'Shikumi.Error.ValidationFailure' carrying @reason@. The accepted
+-- output is returned unchanged.
+validate :: (o -> Bool) -> Text -> Program i o -> Program i o
+validate ok reason = Validate (\o -> if ok o then Right o else Left reason)
+
+-- | 'validate' wrapped in a 'retry': a rejected output triggers re-running the
+-- inner program up to @n@ total attempts.
+validateRetry :: Int -> (o -> Bool) -> Text -> Program i o -> Program i o
+validateRetry n ok reason p = Retry n (validate ok reason p)
+
+-- ---------------------------------------------------------------------------
+-- MajorityVote
+-- ---------------------------------------------------------------------------
+
+-- | Sample a program @K@ times and return the modal output (most frequent under
+-- 'Eq', ties broken by first appearance). The 'TempSchedule' is carried and
+-- serialized but not yet applied to the wire — see 'TempSchedule'.
+majorityVote :: (Eq o) => Int -> TempSchedule -> Program i o -> Program i o
+majorityVote = MajorityVote
+
+-- | Sample a program @K@ times and fold the @K@ outputs with a custom reducer,
+-- for outputs that are not usefully 'Eq'. Defined as an 'ensemble' over @K@
+-- copies of the program.
+majorityVoteBy :: Int -> TempSchedule -> ([o] -> o) -> Program i o -> Program i o
+majorityVoteBy k _sched reducer p = Ensemble (replicate (max 1 k) p) reducer
+
+-- ---------------------------------------------------------------------------
+-- Ensemble
+-- ---------------------------------------------------------------------------
+
+-- | Run several programs on the same input, collect their (homogeneous) results,
+-- and fold them with a total reducer.
+ensemble :: [Program i r] -> ([r] -> o) -> Program i o
+ensemble = Ensemble
diff --git a/src/Shikumi/Effect/Time.hs b/src/Shikumi/Effect/Time.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Effect/Time.hs
@@ -0,0 +1,52 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# OPTIONS_GHC -Wno-redundant-constraints #-}
+
+-- | A minimal clock effect for shikumi, modeled on the @Clock@ effect from the
+-- third-party @time-effectful@ package. We own this module (rather than depend on
+-- @time-effectful@) because shikumi needs a monotonic-clock operation
+-- ('getMonotonicTimeNSec') that @time-effectful@'s @Clock@ does not provide, and
+-- because owning it keeps the dependency surface unchanged.
+--
+-- The effect is static-dispatch with side effects: each operation performs real
+-- 'IO' under the hood via 'unsafeEff_', so callers need only @Time :> es@ rather
+-- than the open-ended @IOE :> es@. The real 'IO' permission is required exactly
+-- once, at the discharge site 'runTime', which is where the program edge lives.
+module Shikumi.Effect.Time
+ ( Time,
+ runTime,
+ getCurrentTime,
+ getMonotonicTimeNSec,
+ UTCTime,
+ )
+where
+
+import Data.Time (UTCTime)
+import Data.Time qualified as Time
+import Data.Word (Word64)
+import Effectful
+import Effectful.Dispatch.Static
+import GHC.Clock qualified as Clock
+
+data Time :: Effect
+
+type instance DispatchOf Time = 'Static 'WithSideEffects
+
+data instance StaticRep Time = TimeRep
+
+-- | The current wall-clock time (UTC). Used for cache-entry metadata and trace
+-- span timestamps.
+getCurrentTime :: (Time :> es) => Eff es UTCTime
+getCurrentTime = unsafeEff_ Time.getCurrentTime
+
+-- | A monotonic nanosecond counter (never runs backwards), suitable for
+-- measuring elapsed time. Used by the evaluation framework to compute per-example
+-- latency. Wraps 'GHC.Clock.getMonotonicTimeNSec'.
+getMonotonicTimeNSec :: (Time :> es) => Eff es Word64
+getMonotonicTimeNSec = unsafeEff_ Clock.getMonotonicTimeNSec
+
+-- | Discharge the 'Time' effect against the real system clock.
+runTime :: (IOE :> es) => Eff (Time ': es) a -> Eff es a
+runTime = evalStaticRep TimeRep
diff --git a/src/Shikumi/Error.hs b/src/Shikumi/Error.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Error.hs
@@ -0,0 +1,62 @@
+-- | The single error vocabulary the whole shikumi framework shares, plus a total
+-- mapping from baikai's transport-level errors into it.
+--
+-- This module owns part of the MasterPlan's integration point #1: every later
+-- ExecPlan (signatures, caching, tracing, tools) MUST surface failures through
+-- 'ShikumiError' rather than inventing its own error type.
+module Shikumi.Error
+ ( ShikumiError (..),
+ fromBaikaiError,
+ isTransient,
+ )
+where
+
+import Baikai.Error (BaikaiError (..))
+import Data.Text (Text)
+import Data.Text qualified as T
+
+-- | The enumerated failure modes named by integration point #1. The decode /
+-- schema / validation constructors are produced by higher layers (structured
+-- output, program validation); the transport ones are produced here by
+-- 'fromBaikaiError'.
+data ShikumiError
+ = -- | provider returned text that is not valid JSON
+ InvalidJSON !Text
+ | -- | a required output field was absent (field name)
+ MissingField !Text
+ | -- | decoded JSON did not match the expected schema
+ SchemaMismatch !Text
+ | -- | a typed value failed a user/program validation rule
+ ValidationFailure !Text
+ | -- | the provider/transport failed (mapped from baikai)
+ ProviderFailure !Text
+ | -- | the call exceeded its time budget
+ Timeout !Text
+ | -- | the running cost ceiling was reached; the call was refused
+ BudgetExceeded !Text
+ deriving stock (Eq, Show)
+
+-- | Total mapping from baikai's transport-level errors into shikumi's
+-- vocabulary. baikai's constructors are 'ProviderError', 'RequestInvalid',
+-- 'DecodeError', and 'ProcessError'.
+--
+-- 'RequestInvalid' maps to 'SchemaMismatch' because in baikai a malformed request
+-- is almost always bad schema/parameters; 'DecodeError' maps to 'InvalidJSON'
+-- because baikai's decode failures are JSON parse failures of the provider
+-- response. These two judgments are the only non-mechanical choices.
+fromBaikaiError :: BaikaiError -> ShikumiError
+fromBaikaiError = \case
+ ProviderError t -> ProviderFailure t
+ RequestInvalid t -> SchemaMismatch ("invalid request: " <> t)
+ DecodeError t -> InvalidJSON t
+ ProcessError n t -> ProviderFailure ("process exited " <> T.pack (show n) <> ": " <> t)
+
+-- | Which errors are worth retrying. Provider/transport failures and timeouts are
+-- transient; decode, schema, validation, and budget errors are deterministic and
+-- retrying cannot fix them. Centralizing the policy here keeps it auditable
+-- (the resilience interpreter in "Shikumi.LLM" consults exactly this predicate).
+isTransient :: ShikumiError -> Bool
+isTransient = \case
+ ProviderFailure {} -> True
+ Timeout {} -> True
+ _ -> False
diff --git a/src/Shikumi/LLM.hs b/src/Shikumi/LLM.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/LLM.hs
@@ -0,0 +1,293 @@
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE TypeFamilies #-}
+
+-- | The provider-neutral @LLM@ effect — integration point #1 of the MasterPlan —
+-- and its interpreters, layered on the policy-free @Baikai@ transport effect from
+-- the @baikai-effectful@ package.
+--
+-- The effect exposes two operations ('Complete', 'Stream'). The bare interpreters
+-- 'runLLM' / 'runLLMWith' map baikai's 'BaikaiError' into 'ShikumiError' and do
+-- nothing else. The resilient interpreter 'runLLMResilient' adds the production
+-- features baikai deliberately omits: retries with exponential backoff, an
+-- in-flight rate limit, and a US-dollar budget ceiling.
+--
+-- Every interpreter is written /in terms of/ the @Baikai@ effect
+-- ('Baikai.Effectful.complete' / 'Baikai.Effectful.streamCollect') rather than
+-- calling baikai's 'IO' functions directly, so shikumi's framework code never
+-- carries 'IOE' — only the bottom @Baikai@ interpreter does. Later plans
+-- (caching, tracing) re-interpret the same @LLM@ operations, so they must sit
+-- above 'runLLMResilient' in the effect stack.
+module Shikumi.LLM
+ ( -- * The effect
+ LLM (..),
+ complete,
+ stream,
+
+ -- * Bare interpreters
+ runLLM,
+ runLLMWith,
+
+ -- * Resilience
+ RetryPolicy (..),
+ defaultRetryPolicy,
+ RateLimiter,
+ newRateLimiter,
+ LLMConfig (..),
+ defaultLLMConfig,
+ runLLMResilient,
+
+ -- * Re-exports of the baikai request/response vocabulary used at call sites
+ Model,
+ Context,
+ Options,
+ Response,
+ AssistantMessageEvent,
+ )
+where
+
+import Baikai
+ ( AssistantMessageEvent (..),
+ Context,
+ Message (..),
+ Model,
+ Options,
+ Response,
+ TerminalPayload,
+ )
+import Baikai.Effectful (Baikai, runBaikai, runBaikaiWith)
+import Baikai.Effectful qualified as BE
+import Baikai.Error (BaikaiError)
+import Baikai.Provider.Registry (ProviderRegistry)
+import Control.Concurrent.STM
+ ( TVar,
+ modifyTVar',
+ newTVarIO,
+ readTVar,
+ retry,
+ writeTVar,
+ )
+import Control.Lens ((^.))
+import Data.Generics.Labels ()
+import Data.Maybe (mapMaybe)
+import Effectful (Dispatch (Dynamic), DispatchOf, Eff, Effect, IOE, liftIO, (:>))
+import Effectful.Concurrent (Concurrent, threadDelay)
+import Effectful.Concurrent.STM (atomically)
+import Effectful.Dispatch.Dynamic (reinterpret_, send)
+import Effectful.Error.Static (Error, catchError, throwError)
+import Effectful.Exception (bracket_, try)
+import Shikumi.Error (ShikumiError (..), fromBaikaiError, isTransient)
+import Shikumi.LLM.Budget (Budget, recordCost, tryReserve)
+
+-- | The provider-neutral LM effect. 'Complete' is a blocking completion;
+-- 'Stream' returns the assembled list of typed events so callers that need
+-- deltas can fold them.
+data LLM :: Effect where
+ Complete :: Model -> Context -> Options -> LLM m Response
+ Stream :: Model -> Context -> Options -> LLM m [AssistantMessageEvent]
+
+type instance DispatchOf LLM = 'Dynamic
+
+-- | Issue a blocking completion. This is integration point #1 — the call every
+-- later plan makes. The argument order mirrors baikai's @completeRequest@.
+complete :: (LLM :> es) => Model -> Context -> Options -> Eff es Response
+complete m c o = send (Complete m c o)
+
+-- | Issue a streaming completion, returning the assembled event list.
+stream :: (LLM :> es) => Model -> Context -> Options -> Eff es [AssistantMessageEvent]
+stream m c o = send (Stream m c o)
+
+-- ---------------------------------------------------------------------------
+-- Bare interpreters
+-- ---------------------------------------------------------------------------
+
+-- | Bare interpreter over baikai's process-global registry. Maps 'BaikaiError'
+-- into 'ShikumiError' and adds no policy. Use 'runLLMWith' for an isolated
+-- registry (tests do this).
+runLLM ::
+ (IOE :> es, Error ShikumiError :> es) =>
+ Eff (LLM : es) a ->
+ Eff es a
+runLLM = reinterpret_ runBaikai bareHandler
+
+-- | Bare interpreter over an explicit registry.
+runLLMWith ::
+ (IOE :> es, Error ShikumiError :> es) =>
+ ProviderRegistry ->
+ Eff (LLM : es) a ->
+ Eff es a
+runLLMWith reg = reinterpret_ (runBaikaiWith reg) bareHandler
+
+-- | The bare handler, shared by both interpreters. It runs in the handler stack
+-- (@Baikai : es@), so it can call the @Baikai@ transport effect and throw
+-- through the @Error ShikumiError@ effect. The blocking path catches baikai's
+-- 'BaikaiError' (thrown by 'BE.complete') and remaps it; the streaming path
+-- surfaces failures in-band as a terminal @EventError@, so it does not catch.
+bareHandler ::
+ (Baikai :> es, Error ShikumiError :> es) =>
+ LLM (Eff localEs) a ->
+ Eff es a
+bareHandler = \case
+ Complete m c o -> do
+ res <- try @BaikaiError (BE.complete m c o)
+ either (throwError . fromBaikaiError) pure res
+ Stream m c o -> BE.streamCollect m c o
+
+-- ---------------------------------------------------------------------------
+-- Resilience: retries, rate limiting, budget
+-- ---------------------------------------------------------------------------
+
+-- | Exponential-backoff retry policy.
+data RetryPolicy = RetryPolicy
+ { -- | total tries including the first (>= 1)
+ maxAttempts :: !Int,
+ -- | first backoff delay, in milliseconds
+ baseDelayMs :: !Int,
+ -- | cap on any single backoff delay, in milliseconds
+ maxDelayMs :: !Int
+ }
+ deriving stock (Eq, Show)
+
+-- | A sensible default: up to three tries, 200ms base, capped at 5s.
+defaultRetryPolicy :: RetryPolicy
+defaultRetryPolicy = RetryPolicy {maxAttempts = 3, baseDelayMs = 200, maxDelayMs = 5000}
+
+-- | A simple in-flight rate limiter: a counter of available permits. Build it
+-- once with 'newRateLimiter' and store it in the 'LLMConfig' (not per call).
+newtype RateLimiter = RateLimiter (TVar Int)
+
+-- | Create a rate limiter that allows at most @n@ concurrent calls.
+newRateLimiter :: Int -> IO RateLimiter
+newRateLimiter n = RateLimiter <$> newTVarIO n
+
+-- | Interpretation-time policy for 'runLLMResilient'.
+data LLMConfig = LLMConfig
+ { retryPolicy :: !RetryPolicy,
+ -- | 'Nothing' = unlimited cost
+ budget :: !(Maybe Budget),
+ -- | 'Nothing' = unbounded concurrency
+ rateLimit :: !(Maybe RateLimiter),
+ -- | which baikai registry to dispatch against
+ registry :: !ProviderRegistry
+ }
+
+-- | A config with default retries, no budget, and no rate limit, dispatching
+-- against the given registry. Set 'budget' / 'rateLimit' to opt in.
+defaultLLMConfig :: ProviderRegistry -> LLMConfig
+defaultLLMConfig reg =
+ LLMConfig
+ { retryPolicy = defaultRetryPolicy,
+ budget = Nothing,
+ rateLimit = Nothing,
+ registry = reg
+ }
+
+-- | The resilient interpreter. Each operation is wrapped, outermost to
+-- innermost, by: budget check → rate-limit acquire → retry loop → the @Baikai@
+-- transport call. The budget is reserved once before the attempts and charged
+-- once after success; retries re-run only the transport call.
+runLLMResilient ::
+ (IOE :> es, Concurrent :> es, Error ShikumiError :> es) =>
+ LLMConfig ->
+ Eff (LLM : es) a ->
+ Eff es a
+runLLMResilient cfg = reinterpret_ (runBaikaiWith (registry cfg)) $ \case
+ Complete m c o ->
+ withBudget mb . withRateLimit mr . retrying rp $ do
+ res <- try @BaikaiError (BE.complete m c o)
+ case res of
+ Left be -> throwError (fromBaikaiError be)
+ Right resp -> do
+ liftIO (chargeBudget mb resp)
+ pure resp
+ Stream m c o ->
+ withBudget mb . withRateLimit mr . retrying rp $ do
+ evs <- BE.streamCollect m c o
+ liftIO (chargeBudgetFromEvents mb evs)
+ pure evs
+ where
+ mb = budget cfg
+ mr = rateLimit cfg
+ rp = retryPolicy cfg
+
+-- | Optimistic pre-call budget gate. Refuses the call with 'BudgetExceeded' when
+-- the running total has already reached the ceiling.
+withBudget ::
+ (IOE :> es, Error ShikumiError :> es) =>
+ Maybe Budget ->
+ Eff es a ->
+ Eff es a
+withBudget Nothing act = act
+withBudget (Just b) act = do
+ ok <- liftIO (tryReserve b)
+ if ok then act else throwError (BudgetExceeded "cost ceiling reached")
+
+-- | Bound concurrency to the limiter's permits, releasing on every exit path.
+withRateLimit ::
+ (Concurrent :> es) =>
+ Maybe RateLimiter ->
+ Eff es a ->
+ Eff es a
+withRateLimit Nothing act = act
+withRateLimit (Just (RateLimiter tv)) act =
+ bracket_ acquire release act
+ where
+ acquire = atomically $ do
+ n <- readTVar tv
+ if n <= 0 then retry else writeTVar tv (n - 1)
+ release = atomically (modifyTVar' tv (+ 1))
+
+-- | Retry a transient-failing action with exponential backoff. Non-transient
+-- errors (per 'isTransient') propagate immediately without consuming a retry.
+retrying ::
+ (Concurrent :> es, Error ShikumiError :> es) =>
+ RetryPolicy ->
+ Eff es a ->
+ Eff es a
+retrying pol act = go 1
+ where
+ go attempt =
+ act `catchError` \_cs e ->
+ if isTransient e && attempt < maxAttempts pol
+ then do
+ threadDelay (backoffMicros pol attempt)
+ go (attempt + 1)
+ else throwError e
+
+-- | Backoff for the @n@-th attempt (1-based), in microseconds, capped by the
+-- policy's 'maxDelayMs'.
+backoffMicros :: RetryPolicy -> Int -> Int
+backoffMicros pol attempt =
+ 1000 * min (maxDelayMs pol) (baseDelayMs pol * (2 ^ (attempt - 1)))
+
+-- | Charge a completed blocking call's cost against the budget, reading baikai's
+-- per-response @Usage.cost.usd@.
+chargeBudget :: Maybe Budget -> Response -> IO ()
+chargeBudget Nothing _ = pure ()
+chargeBudget (Just b) resp = recordCost b (responseCostUSD resp)
+
+-- | The US-dollar cost baikai computed for a response.
+responseCostUSD :: Response -> Rational
+responseCostUSD resp = resp ^. #message . #usage . #cost . #usd
+
+-- | Charge a completed streaming call's cost, read from the terminal event's
+-- assembled message.
+chargeBudgetFromEvents :: Maybe Budget -> [AssistantMessageEvent] -> IO ()
+chargeBudgetFromEvents Nothing _ = pure ()
+chargeBudgetFromEvents (Just b) evs =
+ case mapMaybe eventCostUSD evs of
+ (c : _) -> recordCost b c
+ [] -> pure ()
+
+-- | The cost carried by a terminal streaming event ('EventDone' / 'EventError'),
+-- if any.
+eventCostUSD :: AssistantMessageEvent -> Maybe Rational
+eventCostUSD = \case
+ EventDone tp -> terminalCostUSD tp
+ EventError tp -> terminalCostUSD tp
+ _ -> Nothing
+
+-- | Pull @usage.cost.usd@ out of a terminal payload's assembled message.
+terminalCostUSD :: TerminalPayload -> Maybe Rational
+terminalCostUSD tp = case tp ^. #message of
+ AssistantMessage payload -> Just (payload ^. #usage . #cost . #usd)
+ _ -> Nothing
diff --git a/src/Shikumi/LLM/Budget.hs b/src/Shikumi/LLM/Budget.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/LLM/Budget.hs
@@ -0,0 +1,55 @@
+-- | A running US-dollar cost ceiling for LM calls.
+--
+-- The budget is enforced /before/ a call by an optimistic 'tryReserve' check and
+-- updated /after/ each successful call with 'recordCost', reading baikai's
+-- per-response @Usage.cost.usd@ (a 'Rational', always populated). The semantics
+-- are deliberately simple: the budget refuses the /next/ call once the running
+-- total has reached the ceiling. A single in-flight call is never clipped
+-- mid-stream (baikai exposes no streaming cost preview), but once it pushes the
+-- total to/over the cap, every subsequent call fails fast.
+module Shikumi.LLM.Budget
+ ( Budget (..),
+ newBudget,
+ tryReserve,
+ recordCost,
+ spentUSD,
+ )
+where
+
+import Control.Concurrent.STM
+ ( TVar,
+ atomically,
+ modifyTVar',
+ newTVarIO,
+ readTVar,
+ readTVarIO,
+ )
+
+-- | A running cost ceiling in US dollars. @ceilingUSD == Nothing@ means
+-- "unlimited"; 'spentRef' accumulates the dollars charged so far.
+data Budget = Budget
+ { ceilingUSD :: !(Maybe Rational),
+ spentRef :: !(TVar Rational)
+ }
+
+-- | Build a fresh budget with zero spent. Pass 'Nothing' for an unlimited budget.
+newBudget :: Maybe Rational -> IO Budget
+newBudget c = Budget c <$> newTVarIO 0
+
+-- | Optimistic pre-call check: 'True' if a call is permitted (the running total
+-- has not yet reached the ceiling), 'False' once the ceiling is met/exceeded.
+-- An unlimited budget always permits.
+tryReserve :: Budget -> IO Bool
+tryReserve b = case ceilingUSD b of
+ Nothing -> pure True
+ Just cap -> atomically $ do
+ s <- readTVar (spentRef b)
+ pure (s < cap)
+
+-- | Add the actual cost of a completed call (read from baikai @Usage.cost.usd@).
+recordCost :: Budget -> Rational -> IO ()
+recordCost b c = atomically $ modifyTVar' (spentRef b) (+ c)
+
+-- | The dollars charged against this budget so far.
+spentUSD :: Budget -> IO Rational
+spentUSD b = readTVarIO (spentRef b)
diff --git a/src/Shikumi/Module.hs b/src/Shikumi/Module.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Module.hs
@@ -0,0 +1,216 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+-- | The two foundational /modules/ of shikumi (EP-4): @predict@ and
+-- @chainOfThought@. These are ordinary functions that build 'Program' values out
+-- of the three core constructors — not new constructors. This is the pattern the
+-- combinators in @docs/plans/5-module-combinators-and-control-flow.md@ follow.
+--
+-- @chainOfThought@ extends the output signature with a leading @reasoning@ field
+-- (DSPy's @ChainOfThought@), producing a @Program i (WithReasoning o)@, then
+-- projects the reasoning back out with 'FMap' so the caller's program type stays
+-- @Program i o@. The reasoning-augmented node is a perfectly ordinary 'Predict'
+-- node, so its instruction and demos are visible to @paramsTraversal@ like any
+-- other node — the optimizer tunes a chain-of-thought node with no special casing.
+--
+-- @WithReasoning@'s schema/decode/prompt instances are hand-written rather than
+-- @Generic@-derived: its @value@ field is polymorphic in @o@, and the schema
+-- classes' overlappable per-field instances cannot be resolved for an abstract
+-- type variable. EP-3 exposes no @withReasoningField@, so the (nested) augmentation
+-- lives here. See the plan's Decision Log.
+module Shikumi.Module
+ ( predict,
+ chainOfThought,
+ chainOfThoughtRaw,
+ twoStep,
+ WithReasoning (..),
+ )
+where
+
+import Baikai (user, _Context, _Model, _Options)
+import Control.Lens ((&), (.~))
+import Data.Aeson (Object, Value (Object))
+import Data.Aeson.Key qualified as Key
+import Data.Aeson.KeyMap qualified as KM
+import Data.Generics.Labels ()
+import Data.Proxy (Proxy (..))
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Vector qualified as V
+import Effectful.Error.Static (throwError)
+import GHC.Generics (Generic)
+import Shikumi.Adapter (Adapter (..), ToPrompt (..), fallbackAdapter, responseText)
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.LLM (complete)
+import Shikumi.Program (Program (FMap, Predict), embed, emptyParams)
+import Shikumi.Schema (FromModel (..), ToSchema (..), Validatable)
+import Shikumi.Schema.Types
+ ( FieldMeta (..),
+ FieldPath,
+ objectSchema,
+ pushField,
+ renderPath,
+ stringSchema,
+ withDescription,
+ )
+import Shikumi.Signature (Signature (..), getInstruction)
+
+-- | The basic predictor over a signature: a single 'Predict' node with default
+-- (empty) parameters — no instruction override, no demos. The constraints are
+-- exactly those the 'Predict' constructor captures.
+predict ::
+ (FromModel i, FromModel o, ToSchema o, Validatable o, ToPrompt i, ToPrompt o) =>
+ Signature i o ->
+ Program i o
+predict sig = Predict sig emptyParams
+
+-- ---------------------------------------------------------------------------
+-- Chain of thought
+-- ---------------------------------------------------------------------------
+
+-- | An output @o@ wrapped with a leading step-by-step @reasoning@ field. The model
+-- emits its reasoning first, then the structured answer nested under @value@ (the
+-- order matters: reason, then commit).
+data WithReasoning o = WithReasoning
+ { reasoning :: !Text,
+ value :: !o
+ }
+ deriving stock (Eq, Show, Generic)
+
+-- A nested object schema: @{reasoning : string, value : }@.
+instance (ToSchema o) => ToSchema (WithReasoning o) where
+ toSchema _ =
+ objectSchema
+ [ ("reasoning", withDescription "Step-by-step reasoning, written before the answer" stringSchema),
+ ("value", toSchema (Proxy @o))
+ ]
+ ["reasoning", "value"]
+
+instance (FromModel o) => FromModel (WithReasoning o) where
+ fromModelP path = \case
+ Object o ->
+ WithReasoning
+ <$> getField path "reasoning" o
+ <*> getField path "value" o
+ _ -> Left (SchemaMismatch (renderPath path <> ": expected object"))
+
+-- Render without leaning on @o@'s @Show@: reuse @o@'s own 'ToPrompt'.
+instance (ToPrompt o) => ToPrompt (WithReasoning o) where
+ toPromptFields wr = ("reasoning", reasoning wr) : toPromptFields (value wr)
+
+ -- @value@ is polymorphic in @o@, so the generic image walk cannot resolve; a
+ -- 'WithReasoning' is a text-only output wrapper and carries no image fields.
+ imageFields _ = []
+ imageFieldNames _ = []
+
+-- | Look up a required field in a JSON object, locating a miss precisely.
+getField :: (FromModel a) => FieldPath -> Text -> Object -> Either ShikumiError a
+getField path nm o = case KM.lookup (Key.fromText nm) o of
+ Nothing -> Left (MissingField (renderPath (pushField nm path)))
+ Just v -> fromModelP (pushField nm path) v
+
+-- | Chain-of-thought that yields the bare @o@: build the reasoning-augmented node,
+-- then 'FMap' out the answer.
+chainOfThought ::
+ (FromModel i, FromModel o, ToSchema o, ToPrompt i, ToPrompt o) =>
+ Signature i o ->
+ Program i o
+chainOfThought sig = FMap value (chainOfThoughtRaw sig)
+
+-- | Chain-of-thought that keeps the reasoning visible in the output.
+chainOfThoughtRaw ::
+ (FromModel i, FromModel o, ToSchema o, ToPrompt i, ToPrompt o) =>
+ Signature i o ->
+ Program i (WithReasoning o)
+chainOfThoughtRaw sig = Predict (withReasoningField sig) emptyParams
+
+-- | Augment a signature's output with a leading @reasoning@ field and amend its
+-- instruction to ask for step-by-step reasoning before the answer. The input
+-- field metadata is carried over from the source signature; the output metadata is
+-- the two fields of 'WithReasoning' (matching its hand-written schema).
+withReasoningField :: Signature i o -> Signature i (WithReasoning o)
+withReasoningField sig =
+ Signature
+ { instruction =
+ getInstruction sig
+ <> "\n\nThink step by step. First write your reasoning, then the answer.",
+ demos = [],
+ inputFields = inputFields sig,
+ outputFields =
+ [ FieldMeta {fieldName = "reasoning", fieldDesc = Just "step-by-step reasoning"},
+ FieldMeta {fieldName = "value", fieldDesc = Nothing}
+ ]
+ }
+
+-- ---------------------------------------------------------------------------
+-- Two-step extraction (EP-26)
+-- ---------------------------------------------------------------------------
+
+-- | A tiny internal input record holding a free-form answer, fed to the
+-- extraction call. Not exported.
+newtype ExtractIn = ExtractIn {text :: Text}
+ deriving stock (Generic, Show)
+
+instance ToPrompt ExtractIn
+
+-- | A two-call adapter expressed as a program combinator: ask the main model for a
+-- free-form answer (plain prose, no JSON/marker shape requested), then ask an
+-- extraction model to coerce that prose into the typed output @o@. Useful for
+-- strong reasoners that are weak at structured output.
+--
+-- It is an 'Shikumi.Program.embed' node, not an @Adapter@ value: an @Adapter@'s
+-- @parse@ is pure (@Response -> Either ShikumiError o@) and cannot issue the second
+-- model call, whereas an embedded body runs in 'Shikumi.Program.runProgram'\'s
+-- effect row and can. Because it carries no 'Shikumi.Program.Params', the
+-- parameter-count invariant (count == number of 'Predict' nodes) holds and the
+-- serializers/compilers pass it through unchanged.
+--
+-- Known limitation (matching DSPy's @TwoStepAdapter@): both calls go to the same
+-- ambient model under 'Shikumi.Program.runProgram'. A separate, smaller extraction
+-- model would be wired by running the extraction under a different interpreter;
+-- that is out of scope here.
+twoStep ::
+ (FromModel o, ToSchema o, Validatable o, ToPrompt i, ToPrompt o) =>
+ Signature i o ->
+ Program i o
+twoStep sig = embed $ \i -> do
+ -- 1. Free-form call: plain prose, no structured shape requested.
+ let ffCtx =
+ _Context
+ & #systemPrompt .~ Just (freeFormSystem sig)
+ & #messages .~ V.fromList [user (toPrompt i)]
+ ffResp <- complete _Model ffCtx _Options
+ -- 2. Extraction call: coerce the prose into the typed output via the marker
+ -- fallback adapter (a robust extraction target that already round-trips).
+ let exSig = extractSig sig
+ (exCtx, exOpts) = render fallbackAdapter exSig (ExtractIn (responseText ffResp))
+ exResp <- complete _Model exCtx exOpts
+ either throwError pure (parse fallbackAdapter exSig exResp)
+
+-- | The plain-prose system prompt for the free-form call (mirrors DSPy's
+-- @TwoStepAdapter.format_task_description@): the instruction, the input/output
+-- fields named in words, and a request to answer in detail.
+freeFormSystem :: Signature i o -> Text
+freeFormSystem sig =
+ "You are a helpful assistant. "
+ <> getInstruction sig
+ <> "\n\nAs input you will be provided with: "
+ <> namesOf (inputFields sig)
+ <> ".\nYour answer must contain: "
+ <> namesOf (outputFields sig)
+ <> ".\nLay out your answer in detail, in plain prose."
+ where
+ namesOf = T.intercalate ", " . map fieldName
+
+-- | The extraction signature @ExtractIn -> o@: a single @text@ input holding the
+-- free-form answer, reusing the original signature's output-field metadata (so it
+-- needs no @GFieldMetas (Rep o)@ constraint).
+extractSig :: Signature i o -> Signature ExtractIn o
+extractSig sig =
+ Signature
+ { instruction = "The text below contains the answer. Extract these fields verbatim.",
+ demos = [],
+ inputFields = [FieldMeta {fieldName = "text", fieldDesc = Nothing}],
+ outputFields = outputFields sig
+ }
diff --git a/src/Shikumi/Multimodal.hs b/src/Shikumi/Multimodal.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Multimodal.hs
@@ -0,0 +1,182 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DefaultSignatures #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE UndecidableInstances #-}
+
+-- | Multimodal /input/ field types for signatures (EP-24).
+--
+-- Shikumi's V1 is text-in, text-out: every input field renders to prompt text.
+-- This module adds a typed 'Image' field that lowers to baikai's native inline
+-- image block ('Baikai.Content.ImageContent' / 'Baikai.Content.UserImage'), so a
+-- provider actually /sees/ the picture instead of receiving base64 buried in the
+-- running prose.
+--
+-- 'Image' stores /decoded/ bytes plus a MIME type. base64 is purely a wire
+-- concern that baikai handles itself (its @ToJSON ImageContent@ base64-encodes
+-- under a @data@ key), so 'imageToContent' is a direct field copy — no encoding
+-- happens here.
+--
+-- Scope (honesty about provider limits): baikai's @UserContent@ models exactly
+-- @UserText@ and @UserImage@ today, so this module ships __image only__. Audio and
+-- document fields are upstream-gated on a new @Baikai.Content@ constructor; see the
+-- EP-24 plan's "Audio and document: upstream-gated future work" section. An 'Image'
+-- is __input-only__: it has no 'Shikumi.Schema.ToSchema' instance, so putting one in
+-- an /output/ record is a clean compile error (a model cannot emit raw image bytes
+-- through the structured-decode path).
+module Shikumi.Multimodal
+ ( -- * The image field type
+ Image (..),
+ imageFromBytes,
+ imageFromFile,
+ imageFromBase64,
+
+ -- * Lowering to baikai
+ imageToContent,
+
+ -- * Generic image-field discovery (drives "Shikumi.Adapter"\'s 'Shikumi.Adapter.ToPrompt' image methods)
+ GImageFields (..),
+ GImageFieldNames (..),
+ )
+where
+
+import Baikai (ImageContent (..))
+import Data.Bifunctor (first)
+import Data.ByteString (ByteString)
+import Data.ByteString qualified as BS
+import Data.ByteString.Base64 qualified as Base64
+import Data.Kind (Type)
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Text.Encoding (encodeUtf8)
+import GHC.Generics
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.Schema.Types (Field (..))
+import System.FilePath (takeExtension)
+
+-- | A typed image usable as a /signature input/ field. Stores decoded bytes and a
+-- MIME type; base64 is a wire detail handled by baikai when the image is sent.
+data Image = Image
+ { imageBytes :: !ByteString,
+ imageMime :: !Text
+ }
+ deriving stock (Eq, Show, Generic)
+
+-- | Build an image from already-decoded bytes and an explicit MIME type. The
+-- primitive constructor; the others normalise into it.
+imageFromBytes :: Text -> ByteString -> Image
+imageFromBytes mime bs = Image {imageBytes = bs, imageMime = mime}
+
+-- | Read an image file from disk, inferring the MIME type from the extension.
+-- Returns 'Left' a 'SchemaMismatch' if the extension is unrecognised (checked
+-- before any read, so an unsupported path never touches the filesystem).
+imageFromFile :: FilePath -> IO (Either ShikumiError Image)
+imageFromFile fp = case mimeForExtension (takeExtension fp) of
+ Nothing ->
+ pure (Left (SchemaMismatch ("unsupported image extension: " <> T.pack (takeExtension fp))))
+ Just mime -> do
+ bs <- BS.readFile fp
+ pure (Right (imageFromBytes mime bs))
+
+-- | Decode a base64 string into an image with the given MIME type. Returns 'Left'
+-- an 'InvalidJSON' decode error if the base64 is malformed.
+imageFromBase64 :: Text -> Text -> Either ShikumiError Image
+imageFromBase64 mime b64 =
+ imageFromBytes mime
+ <$> first (\e -> InvalidJSON ("image base64 decode: " <> T.pack e)) (Base64.decode (encodeUtf8 b64))
+
+-- | Lower an 'Image' into baikai's wire image block. Bytes pass through decoded;
+-- baikai base64-encodes them only when serialising to the wire.
+imageToContent :: Image -> ImageContent
+imageToContent img = ImageContent {imageData = imageBytes img, mimeType = imageMime img}
+
+-- | A tiny fixed extension-to-MIME table (case-insensitive). Deliberately not a
+-- MIME database: the supported image media types are few and stable.
+mimeForExtension :: String -> Maybe Text
+mimeForExtension ext = case T.toLower (T.pack ext) of
+ ".png" -> Just "image/png"
+ ".jpg" -> Just "image/jpeg"
+ ".jpeg" -> Just "image/jpeg"
+ ".gif" -> Just "image/gif"
+ ".webp" -> Just "image/webp"
+ _ -> Nothing
+
+-- ---------------------------------------------------------------------------
+-- Generic image-field discovery
+-- ---------------------------------------------------------------------------
+--
+-- These generic walks back the 'Shikumi.Adapter.ToPrompt' image methods
+-- ('Shikumi.Adapter.imageFields' / 'Shikumi.Adapter.imageFieldNames'): the former
+-- collects an input record's 'Image' values (so the adapter lowers them to baikai
+-- 'Baikai.Content.UserImage' blocks), the latter names those fields (so the text
+-- renderer drops them from the prompt body). A record with no image fields yields
+-- @[]@ from both, which is exactly the regression-safe text-only path. The methods
+-- live on 'Shikumi.Adapter.ToPrompt' (not a separate class) so that the single
+-- @ToPrompt i@ constraint already threaded through @predict@/@Predict@/@adapterFor@
+-- carries them with no new constraint anywhere; the generic defaults mean any
+-- @Generic@ input record gets both for free.
+
+-- | A leaf field's images: a bare 'Image' is one image, a 'Field'-wrapped 'Image'
+-- unwraps, and every other leaf type contributes none. This single check drives
+-- both the value collector and the field-name collector.
+class ImageLeaf t where
+ imageLeaf :: t -> [Image]
+
+instance {-# OVERLAPPABLE #-} ImageLeaf a where
+ imageLeaf _ = []
+
+instance ImageLeaf Image where
+ imageLeaf img = [img]
+
+instance (ImageLeaf a) => ImageLeaf (Field d a) where
+ imageLeaf (Field a) = imageLeaf a
+
+-- | Collect the 'Image' values of a record's generic representation, in field
+-- order.
+class GImageFields (f :: Type -> Type) where
+ gImageFields :: f p -> [Image]
+
+instance (GImageFields cs) => GImageFields (D1 d cs) where
+ gImageFields (M1 x) = gImageFields x
+
+instance (GImageFields cs) => GImageFields (C1 c cs) where
+ gImageFields (M1 x) = gImageFields x
+
+instance (GImageFields a, GImageFields b) => GImageFields (a :*: b) where
+ gImageFields (a :*: b) = gImageFields a ++ gImageFields b
+
+instance GImageFields U1 where
+ gImageFields _ = []
+
+instance (ImageLeaf t) => GImageFields (S1 s (K1 R t)) where
+ gImageFields (M1 (K1 v)) = imageLeaf v
+
+-- A sum type (an enum, or a tagged union) is never an image carrier: image inputs
+-- are records. Defined so the generic default is total for any 'Generic' type.
+instance GImageFields (l :+: r) where
+ gImageFields _ = []
+
+-- | Collect the /names/ of a record's image fields, in field order. Emits a
+-- selector's name iff its field type is an 'Image' (bare or 'Field'-wrapped).
+class GImageFieldNames (f :: Type -> Type) where
+ gImageFieldNames :: f p -> [Text]
+
+instance (GImageFieldNames cs) => GImageFieldNames (D1 d cs) where
+ gImageFieldNames (M1 x) = gImageFieldNames x
+
+instance (GImageFieldNames cs) => GImageFieldNames (C1 c cs) where
+ gImageFieldNames (M1 x) = gImageFieldNames x
+
+instance (GImageFieldNames a, GImageFieldNames b) => GImageFieldNames (a :*: b) where
+ gImageFieldNames (a :*: b) = gImageFieldNames a ++ gImageFieldNames b
+
+instance GImageFieldNames U1 where
+ gImageFieldNames _ = []
+
+instance (Selector s, ImageLeaf t) => GImageFieldNames (S1 s (K1 R t)) where
+ gImageFieldNames m@(M1 (K1 v)) = [T.pack (selName m) | not (null (imageLeaf v))]
+
+-- A sum type carries no image field names; see 'GImageFields' for the rationale.
+instance GImageFieldNames (l :+: r) where
+ gImageFieldNames _ = []
diff --git a/src/Shikumi/Prelude.hs b/src/Shikumi/Prelude.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Prelude.hs
@@ -0,0 +1,10 @@
+-- | A small shared prelude for shikumi. For now it re-exports the handful of
+-- common names every module reaches for; it grows in later ExecPlans.
+module Shikumi.Prelude
+ ( Text,
+ Vector,
+ )
+where
+
+import Data.Text (Text)
+import Data.Vector (Vector)
diff --git a/src/Shikumi/Program.hs b/src/Shikumi/Program.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Program.hs
@@ -0,0 +1,653 @@
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+-- | The keystone of shikumi (EP-4): a typed /deep embedding/ of an LM program as
+-- inspectable data. A 'Program' @i o@ is a tree of three constructors that can be
+-- done three different things with at once:
+--
+-- * __run__ as a typed function — 'runProgram' interprets the tree as an 'Eff'
+-- computation that issues @LLM@ calls and returns a typed @o@ (or throws a
+-- typed 'ShikumiError');
+-- * __rewritten as data__ — 'paramsTraversal' / 'foldParams' / 'mapParams' /
+-- 'mapParamsAt' read and replace each node's optimizable 'Params' (its
+-- instruction override and few-shot demos) without running the program and
+-- without runtime reflection, which is what the optimizer
+-- (@docs/plans/10-optimizer-framework.md@) needs;
+-- * __serialized__ — 'programShape' captures the closure-free structure and
+-- 'programParams' / 'setProgramParams' move the JSON-serializable parameter
+-- vector, so an optimized program's state can be saved and replayed.
+--
+-- The constructor set is deliberately minimal (three): richer modules
+-- (@chainOfThought@; the combinators in
+-- @docs/plans/5-module-combinators-and-control-flow.md@) are /derived/ functions
+-- that build these constructors, not new constructors.
+--
+-- This module consumes EP-1 (@Shikumi.LLM@, @Shikumi.Error@) and EP-3
+-- (@Shikumi.Signature@, @Shikumi.Adapter@, @Shikumi.Schema@). See the plan's
+-- Decision Log for the reconciliations with the delivered EP-3 surface (which
+-- exposes @render@/@parse@/@adapterFor@ rather than a single @runSignature@).
+module Shikumi.Program
+ ( -- * The representation
+ Program
+ ( Predict,
+ Compose,
+ FMap,
+ Map,
+ Parallel,
+ Retry,
+ RetryWhen,
+ Validate,
+ MajorityVote,
+ Ensemble,
+ Embed
+ ),
+ Params (..),
+ Demo (..),
+ emptyParams,
+ TempSchedule (..),
+
+ -- * Construction & execution
+ pipeline,
+ embed,
+ runProgram,
+ runProgramConc,
+
+ -- * Per-node field metadata (EP-16)
+ NodeFields (..),
+ nodeFieldsIndexed,
+
+ -- * Execution internals (reused by alternative executors, e.g. @runProgramTraced@)
+ retryWith,
+ acceptOrReject,
+ modal,
+ sampleTemps,
+ withSampleTemp,
+
+ -- * Parameter interface (the optimizer/compiler contract)
+ paramsTraversal,
+ foldParams,
+ mapParams,
+ mapParamsAt,
+
+ -- * Serialization (parameter state only — never closures)
+ ProgramShape (..),
+ ProgramShapeError (..),
+ programShape,
+ programParams,
+ setProgramParams,
+ )
+where
+
+import Baikai (Model, _Model)
+import Data.Aeson (FromJSON, ToJSON, Value)
+import Data.Functor.Const (Const (..))
+import Data.Functor.Identity (Identity (..))
+import Data.Maybe (fromMaybe)
+import Data.Text (Text)
+import Data.Text qualified as T
+import Effectful (Eff, (:>))
+import Effectful.Concurrent (Concurrent)
+import Effectful.Concurrent.Async (concurrently, mapConcurrently, pooledMapConcurrentlyN)
+import Effectful.Dispatch.Dynamic (interpose, passthrough)
+import Effectful.Error.Static (Error, catchError, throwError)
+import GHC.Generics (Generic)
+import Shikumi.Adapter
+ ( Adapter (..),
+ ToPrompt,
+ adapterFor,
+ attachSchema,
+ fallbackAdapter,
+ nativeAdapter,
+ stampTemperature,
+ )
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.LLM (LLM (..), Response, complete)
+import Shikumi.Schema (FromModel, ToSchema, Validatable, deriveSchema, fromModel)
+import Shikumi.Schema.Types (fieldName)
+import Shikumi.Signature (Signature, getInstruction, inputFields, outputFields, setDemos, setInstruction)
+import Shikumi.Signature qualified as Sig
+
+-- ---------------------------------------------------------------------------
+-- Optimizable node state
+-- ---------------------------------------------------------------------------
+
+-- | The optimizable overlay of a single node: an optional instruction override
+-- (@Nothing@ = use the signature's default) and an ordered list of few-shot
+-- demonstrations. This is the /uniform, serializable/ handle the compiler
+-- (@docs/plans/9-compiler-layer.md@) and optimizer
+-- (@docs/plans/10-optimizer-framework.md@) manipulate regardless of a node's
+-- @i@/@o@ — hence demos are stored as type-agnostic JSON (see 'Demo').
+data Params = Params
+ { instructionOverride :: !(Maybe Text),
+ demos :: ![Demo]
+ }
+ deriving stock (Eq, Show, Generic)
+
+instance ToJSON Params
+
+instance FromJSON Params
+
+-- | A worked input/output example, stored as JSON so it is uniform across nodes
+-- of differing types. At run time each demo is decoded back into the node's typed
+-- @Sig.Demo i o@ and spliced into the prompt by EP-3's adapter; a demo whose JSON
+-- does not decode surfaces as a 'ShikumiError'.
+data Demo = Demo
+ { input :: !Value,
+ output :: !Value
+ }
+ deriving stock (Eq, Show, Generic)
+
+instance ToJSON Demo
+
+instance FromJSON Demo
+
+-- | The default empty overlay: no instruction override, no demos.
+emptyParams :: Params
+emptyParams = Params Nothing []
+
+-- | How 'MajorityVote' varies sampling temperature across its @K@ samples.
+-- 'TempFixed' lists an explicit temperature per sample (cycled if shorter than
+-- @K@, empty = leave temperature at the provider default); 'TempSpread base spread'
+-- centres on @base@ and fans @K@ values out evenly across @[base-spread,
+-- base+spread]@.
+--
+-- Since EP-14 the schedule is /live/: each sample's temperature is stamped onto the
+-- private request-metadata channel ('Shikumi.Adapter.stampTemperature') and the
+-- router ("Shikumi.Routing".@routeLLM@) sets @Options.temperature@ from it, so the
+-- @K@ samples are sent with distinct temperatures on the wire. (Without a router
+-- installed the stamp is inert, as no interpreter realizes it — the hermetic stub
+-- path and the live path both install one.)
+data TempSchedule
+ = TempFixed ![Double]
+ | TempSpread !Double !Double
+ deriving stock (Eq, Show, Generic)
+
+instance ToJSON TempSchedule
+
+instance FromJSON TempSchedule
+
+-- ---------------------------------------------------------------------------
+-- The GADT
+-- ---------------------------------------------------------------------------
+
+-- | A typed LM program. 'Predict' is a single signature-backed LM call carrying
+-- its 'Params'; 'Compose' sequences two programs (its intermediate type is
+-- existential); 'FMap' applies a pure post-processing function (no LM call).
+--
+-- 'Predict' captures the adapter/decode dictionaries existentially so that
+-- 'runProgram' can recover them by pattern-matching — this is what lets a program
+-- be rewritten as data while staying type-checked.
+data Program i o where
+ Predict ::
+ (FromModel i, FromModel o, ToSchema o, Validatable o, ToPrompt i, ToPrompt o) =>
+ Signature i o ->
+ Params ->
+ Program i o
+ Compose :: Program a b -> Program b c -> Program a c
+ FMap :: (o -> o') -> Program i o -> Program i o'
+ -- | Apply a program to every element of a list. The 'Int' is the
+ -- concurrency width honoured by 'runProgramConc' (1 = sequential);
+ -- 'runProgram' always runs sequentially regardless of the width.
+ Map :: Int -> Program a b -> Program [a] [b]
+ -- | Run two programs on the /same/ input and pair their outputs.
+ Parallel :: Program i a -> Program i b -> Program i (a, b)
+ -- | Re-run a program up to @n@ total attempts on any 'ShikumiError'.
+ Retry :: Int -> Program i o -> Program i o
+ -- | Like 'Retry' but only retries errors satisfying the predicate; a
+ -- non-matching error propagates after a single attempt.
+ RetryWhen :: (ShikumiError -> Bool) -> Int -> Program i o -> Program i o
+ -- | Run a program, then check its output: 'Right' accepts (optionally
+ -- normalizing), 'Left' rejects with a reason surfaced as a
+ -- 'ValidationFailure'.
+ Validate :: (o -> Either Text o) -> Program i o -> Program i o
+ -- | Sample a program @K@ times and return the modal (most-frequent under
+ -- 'Eq', ties broken by first-seen) output. Each sample is run with its
+ -- 'TempSchedule' temperature applied to the wire (see 'TempSchedule').
+ MajorityVote :: (Eq o) => Int -> TempSchedule -> Program i o -> Program i o
+ -- | Run several programs on the same input, collect their (homogeneous)
+ -- results, and fold them with a total reducer.
+ Ensemble :: [Program i r] -> ([r] -> o) -> Program i o
+ -- | An opaque effectful step embedded as a program node. The body is
+ -- constrained to /exactly/ 'runProgram'\'s effect row
+ -- (@LLM@ + @Error ShikumiError@) so an 'Embed' node runs under the ordinary
+ -- 'runProgram'/'runProgramConc' without widening integration point #4\'s
+ -- constraint. This is the constructor multi-step agents (ReAct,
+ -- @docs/plans/11-typed-tools-and-react-agents.md@) are built on: the agent
+ -- loop is one 'Embed' node, so the agent is a real, composable 'Program' that
+ -- is runnable, structurally inspectable ('ShapeEmbed'), and serializable
+ -- (it carries no 'Params', like 'FMap'). The body is a closure, so it is
+ -- opaque to the parameter traversal — exactly as 'FMap'\'s function is.
+ Embed :: (forall es. (LLM :> es, Error ShikumiError :> es) => i -> Eff es o) -> Program i o
+
+-- | Sequence two programs, read left-to-right ("first @p@, then @q@"). Typechecks
+-- only when @p@'s output type equals @q@'s input type — an invalid pipeline is a
+-- compile error, which is the whole point.
+pipeline :: Program a b -> Program b c -> Program a c
+pipeline = Compose
+
+-- | Embed an opaque effectful step as a 'Program' node. The smart constructor for
+-- 'Embed': lift an @(i -> Eff es o)@ — runnable in 'runProgram'\'s effect row — into
+-- a @Program i o@. Used by ReAct (@docs/plans/11-typed-tools-and-react-agents.md@) to
+-- make a multi-step agent loop a first-class, composable program.
+embed :: (forall es. (LLM :> es, Error ShikumiError :> es) => i -> Eff es o) -> Program i o
+embed = Embed
+
+-- ---------------------------------------------------------------------------
+-- Execution
+-- ---------------------------------------------------------------------------
+
+-- | The inert placeholder model 'runPredict' passes to 'complete'. Since EP-14
+-- 'runProgram' is model-agnostic: it never chooses a real provider itself. The
+-- ambient model is supplied below the stack by "Shikumi.Routing".@runRouting@ and
+-- the router ("Shikumi.Routing".@routeLLM@) overwrites this placeholder with it on
+-- every outgoing call. With no router installed (the bare-stub test path) the call
+-- still carries '_Model', which 'adapterFor' maps to the prompt-fallback adapter —
+-- preserving EP-4's original behaviour for un-routed runs.
+placeholderModel :: Model
+placeholderModel = _Model
+
+-- | Interpret a program as a typed @Eff@ computation. A 'Predict' node overlays
+-- its 'Params' onto the signature (effective instruction + decoded demos), renders
+-- the request via EP-3's adapter, issues the 'LLM' call, and parses the response
+-- back into a typed @o@ — throwing a 'ShikumiError' on a parse or demo-decode
+-- failure. 'Compose' threads the intermediate value; 'FMap' maps the result purely.
+runProgram ::
+ (LLM :> es, Error ShikumiError :> es) =>
+ Program i o ->
+ i ->
+ Eff es o
+runProgram (Predict sig ps) i = runPredict sig ps i
+runProgram (Compose f g) i = runProgram f i >>= runProgram g
+runProgram (FMap k p) i = k <$> runProgram p i
+runProgram (Map _ p) xs = traverse (runProgram p) xs
+runProgram (Parallel pa pb) i = (,) <$> runProgram pa i <*> runProgram pb i
+runProgram (Retry n p) i = retryWith runProgram (const True) n p i
+runProgram (RetryWhen ok n p) i = retryWith runProgram ok n p i
+runProgram (Validate v p) i = runProgram p i >>= acceptOrReject v
+runProgram (MajorityVote k sched p) i =
+ modal <$> traverse (\mt -> withSampleTemp mt (runProgram p i)) (sampleTemps (max 1 k) sched)
+runProgram (Ensemble ps reduce) i = reduce <$> traverse (\p -> runProgram p i) ps
+runProgram (Embed f) i = f i
+
+-- | The concurrent executor: identical observable semantics to 'runProgram', but
+-- 'Map' (bounded by its width), 'Parallel', 'MajorityVote', and 'Ensemble' run
+-- their independent sub-programs concurrently via @effectful@'s 'Concurrent'
+-- effect. Offered as an additive opt-in so 'runProgram' can keep the exact
+-- @(LLM, Error ShikumiError)@ constraint that is the MasterPlan's integration
+-- point #4 — adding 'Concurrent' there would force it onto every consumer.
+runProgramConc ::
+ (LLM :> es, Error ShikumiError :> es, Concurrent :> es) =>
+ Program i o ->
+ i ->
+ Eff es o
+runProgramConc (Predict sig ps) i = runPredict sig ps i
+runProgramConc (Compose f g) i = runProgramConc f i >>= runProgramConc g
+runProgramConc (FMap k p) i = k <$> runProgramConc p i
+runProgramConc (Map w p) xs = pooledMapConcurrentlyN (max 1 w) (runProgramConc p) xs
+runProgramConc (Parallel pa pb) i = concurrently (runProgramConc pa i) (runProgramConc pb i)
+runProgramConc (Retry n p) i = retryWith runProgramConc (const True) n p i
+runProgramConc (RetryWhen ok n p) i = retryWith runProgramConc ok n p i
+runProgramConc (Validate v p) i = runProgramConc p i >>= acceptOrReject v
+runProgramConc (MajorityVote k sched p) i =
+ modal <$> mapConcurrently (\mt -> withSampleTemp mt (runProgramConc p i)) (sampleTemps (max 1 k) sched)
+runProgramConc (Ensemble ps reduce) i = reduce <$> mapConcurrently (\p -> runProgramConc p i) ps
+-- The body requires only @(LLM, Error ShikumiError)@, a subset of this row, so it
+-- runs unchanged; an embedded agent that wants concurrency uses 'runProgramConc'
+-- internally via its own 'Concurrent' handler.
+runProgramConc (Embed f) i = f i
+
+-- | Interpret a single 'Predict' node: overlay its 'Params', render via EP-3's
+-- adapter, issue the 'LLM' call, parse back to a typed @o@. Shared by both
+-- executors so the wire behaviour is defined once.
+runPredict ::
+ forall i o es.
+ (FromModel i, FromModel o, ToSchema o, ToPrompt i, ToPrompt o) =>
+ (LLM :> es, Error ShikumiError :> es) =>
+ Signature i o ->
+ Params ->
+ i ->
+ Eff es o
+runPredict sig ps i = do
+ sig' <- effectiveSignature sig ps
+ let adapter = adapterFor placeholderModel
+ (ctx, opts0) = render adapter sig' i
+ -- Stamp the derived schema onto the metadata channel regardless of which
+ -- adapter rendered the prompt; the router attaches a real @responseFormat@
+ -- from it for native-capable models and strips it otherwise.
+ opts = attachSchema (deriveSchema @o) opts0
+ resp <- complete placeholderModel ctx opts
+ either throwError pure (parseResponse sig' resp)
+
+-- | Decode a response back into the typed output, tolerant of /either/ wire shape:
+-- a native model (whose @responseFormat@ the router enforced) replies with a JSON
+-- object, while a fallback model replies with @[[ ## field ## ]]@ sections. Because
+-- a model-agnostic 'runPredict' cannot know at render time which the real model is,
+-- we try the native JSON parse first and fall back to the section parser, reporting
+-- the fallback parser's error when both fail (so un-routed runs keep their exact
+-- prior error behaviour).
+parseResponse ::
+ forall i o.
+ (FromModel o, ToSchema o, Validatable o, ToPrompt i, ToPrompt o) =>
+ Signature i o ->
+ Response ->
+ Either ShikumiError o
+parseResponse sig' resp =
+ case parse (nativeAdapter @i @o) sig' resp of
+ Right o -> Right o
+ Left _ -> parse (fallbackAdapter @i @o) sig' resp
+
+-- | The per-sample temperatures for a 'MajorityVote' of @k@ samples. 'Nothing'
+-- leaves the provider default in place; 'Just t' is stamped onto the sample's
+-- requests. The multiset is identical between 'runProgram' and 'runProgramConc'.
+sampleTemps :: Int -> TempSchedule -> [Maybe Double]
+sampleTemps k (TempFixed xs)
+ | null xs = replicate k Nothing
+ | otherwise = map Just (take k (cycle xs))
+sampleTemps k (TempSpread base spread) = map Just (spreadTemps k base spread)
+
+-- | @k@ temperatures fanned evenly across @[base-spread, base+spread]@; a single
+-- sample sits exactly on @base@.
+spreadTemps :: Int -> Double -> Double -> [Double]
+spreadTemps k base spread
+ | k <= 1 = [base]
+ | otherwise = [base - spread + fromIntegral i * step | i <- [0 .. k - 1]]
+ where
+ step = (2 * spread) / fromIntegral (k - 1)
+
+-- | Run an action with one sample's temperature stamped onto every outgoing
+-- 'Complete' (via the private metadata channel the router realizes). 'Nothing' runs
+-- the action unchanged. Scoped to this sample only via 'interpose', so sibling
+-- samples are unaffected and nested votes compose.
+withSampleTemp :: (LLM :> es) => Maybe Double -> Eff es a -> Eff es a
+withSampleTemp Nothing act = act
+withSampleTemp (Just t) act =
+ interpose
+ ( \env -> \case
+ Complete m c o -> complete m c (stampTemperature t o)
+ other -> passthrough env other
+ )
+ act
+
+-- | Apply a validator to a program's output, surfacing a rejection as a
+-- 'ValidationFailure'.
+acceptOrReject ::
+ (Error ShikumiError :> es) => (o -> Either Text o) -> o -> Eff es o
+acceptOrReject v o = either (throwError . ValidationFailure) pure (v o)
+
+-- | The retry loop, parameterized over the executor so 'runProgram' and
+-- 'runProgramConc' share it. Runs @p@ on @i@; on a matching error with attempts
+-- remaining, re-runs; otherwise surfaces the last error.
+retryWith ::
+ (Error ShikumiError :> es) =>
+ (Program i o -> i -> Eff es o) ->
+ (ShikumiError -> Bool) ->
+ Int ->
+ Program i o ->
+ i ->
+ Eff es o
+retryWith run ok n p i = go (max 1 n)
+ where
+ go left =
+ run p i `catchError` \_cs e ->
+ if ok e && left > 1 then go (left - 1) else throwError e
+
+-- | The modal value of a non-empty list: most frequent under 'Eq', ties broken
+-- by first appearance. Tallies in first-seen order, then keeps the first entry
+-- whose count is strictly greatest.
+modal :: (Eq o) => [o] -> o
+modal = pickBest . foldl' tally []
+ where
+ tally acc x = case break ((== x) . fst) acc of
+ (pre, (y, c) : post) -> pre ++ (y, c + 1) : post
+ (pre, []) -> pre ++ [(x, 1 :: Int)]
+ pickBest [] = error "Shikumi.Program.modal: empty sample list"
+ pickBest (z : zs) =
+ fst (foldl' (\best cur -> if snd cur > snd best then cur else best) z zs)
+
+-- | Overlay a node's 'Params' onto its signature: substitute the instruction
+-- override (when present) and decode the JSON demos into the signature's typed
+-- demo channel. A demo whose JSON does not decode is reported as the located
+-- 'ShikumiError' from 'fromModel'.
+effectiveSignature ::
+ (FromModel i, FromModel o, Error ShikumiError :> es) =>
+ Signature i o ->
+ Params ->
+ Eff es (Signature i o)
+effectiveSignature sig ps = do
+ typed <- either throwError pure (traverse decodeDemo (demos ps))
+ pure (setDemos typed (setInstruction instr sig))
+ where
+ instr = fromMaybe (getInstruction sig) (instructionOverride ps)
+ decodeDemo (Demo inJ outJ) = Sig.Demo <$> fromModel inJ <*> fromModel outJ
+
+-- ---------------------------------------------------------------------------
+-- Parameter interface
+-- ---------------------------------------------------------------------------
+
+-- | The source-of-truth traversal: focuses every 'Params' in a program in
+-- /left-to-right depth-first/ order (for @Compose f g@ all of @f@'s come before
+-- @g@'s). Composite nodes ('Compose', 'FMap') carry no 'Params' of their own — so
+-- a program's parameter count equals its number of 'Predict' nodes. Obeys the
+-- @lens@ @Traversal'@ laws; use it directly with @toListOf@/@over@/@set@.
+paramsTraversal :: (Applicative f) => (Params -> f Params) -> Program i o -> f (Program i o)
+paramsTraversal h (Predict sig ps) = Predict sig <$> h ps
+paramsTraversal h (Compose f g) = Compose <$> paramsTraversal h f <*> paramsTraversal h g
+paramsTraversal h (FMap k p) = FMap k <$> paramsTraversal h p
+paramsTraversal h (Map w p) = Map w <$> paramsTraversal h p
+paramsTraversal h (Parallel pa pb) = Parallel <$> paramsTraversal h pa <*> paramsTraversal h pb
+paramsTraversal h (Retry n p) = Retry n <$> paramsTraversal h p
+paramsTraversal h (RetryWhen ok n p) = RetryWhen ok n <$> paramsTraversal h p
+paramsTraversal h (Validate v p) = Validate v <$> paramsTraversal h p
+paramsTraversal h (MajorityVote k sched p) = MajorityVote k sched <$> paramsTraversal h p
+paramsTraversal h (Ensemble ps reduce) = Ensemble <$> traverse (paramsTraversal h) ps <*> pure reduce
+-- 'Embed' carries no 'Params' (its body is an opaque closure, like 'FMap'\'s
+-- function), so it is a traversal leaf — preserved untouched.
+paramsTraversal _ (Embed f) = pure (Embed f)
+
+-- | Read every node's 'Params', in traversal order.
+foldParams :: Program i o -> [Params]
+foldParams = getConst . paramsTraversal (\ps -> Const [ps])
+
+-- | The input- and output-field names of a single 'Predict' node, recovered
+-- structurally. A 'Predict' hides its @i@\/@o@ types existentially, so a typed
+-- @Signature@ cannot escape the GADT — but the field /names/ are plain 'Text' and
+-- can. This is what the EP-16 node-correlated trace and the grounded instruction
+-- proposer (@docs/plans/19-grounded-instruction-proposer.md@) consume.
+data NodeFields = NodeFields
+ { inputFieldNames :: ![Text],
+ outputFieldNames :: ![Text]
+ }
+ deriving stock (Eq, Show, Generic)
+
+-- | For every 'Predict' node, in 'foldParams' order, its index-aligned field
+-- metadata. So @nodeFieldsIndexed p !! n@ describes the same node @mapParamsAt n@
+-- edits and @foldParams p !! n@ parameterizes — the integration-point-#3 ordering
+-- law, shared with @Shikumi.Trace.Node.programNodePaths@.
+nodeFieldsIndexed :: Program i o -> [NodeFields]
+nodeFieldsIndexed = go
+ where
+ go :: forall x y. Program x y -> [NodeFields]
+ go (Predict sig _) =
+ [NodeFields (map fieldName (inputFields sig)) (map fieldName (outputFields sig))]
+ go (Compose a b) = go a ++ go b
+ go (FMap _ p) = go p
+ go (Map _ p) = go p
+ go (Parallel a b) = go a ++ go b
+ go (Retry _ p) = go p
+ go (RetryWhen _ _ p) = go p
+ go (Validate _ p) = go p
+ go (MajorityVote _ _ p) = go p
+ go (Ensemble ps _) = concatMap go ps
+ go (Embed _) = []
+
+-- | Apply a function to every node's 'Params', preserving structure and types.
+mapParams :: (Params -> Params) -> Program i o -> Program i o
+mapParams f = runIdentity . paramsTraversal (Identity . f)
+
+-- | Apply a function to the 'Params' at a single 0-based index in traversal order;
+-- an out-of-range index leaves the program unchanged. The optimizer's primary edit
+-- primitive: "replace node @n@'s instruction/demos". The index it addresses is the
+-- same index 'foldParams' produces (the ordering law).
+mapParamsAt :: Int -> (Params -> Params) -> Program i o -> Program i o
+mapParamsAt n f = fst . go 0
+ where
+ go :: forall x y. Int -> Program x y -> (Program x y, Int)
+ go idx (Predict sig ps) = (Predict sig (if idx == n then f ps else ps), idx + 1)
+ go idx (Compose a b) =
+ let (a', idx') = go idx a
+ (b', idx'') = go idx' b
+ in (Compose a' b', idx'')
+ go idx (FMap k p) =
+ let (p', idx') = go idx p
+ in (FMap k p', idx')
+ go idx (Map w p) =
+ let (p', idx') = go idx p
+ in (Map w p', idx')
+ go idx (Parallel a b) =
+ let (a', idx') = go idx a
+ (b', idx'') = go idx' b
+ in (Parallel a' b', idx'')
+ go idx (Retry m p) =
+ let (p', idx') = go idx p
+ in (Retry m p', idx')
+ go idx (RetryWhen ok m p) =
+ let (p', idx') = go idx p
+ in (RetryWhen ok m p', idx')
+ go idx (Validate v p) =
+ let (p', idx') = go idx p
+ in (Validate v p', idx')
+ go idx (MajorityVote k sched p) =
+ let (p', idx') = go idx p
+ in (MajorityVote k sched p', idx')
+ go idx (Ensemble ps reduce) =
+ let (ps', idx') = goList idx ps
+ in (Ensemble ps' reduce, idx')
+ go idx (Embed body) = (Embed body, idx) -- leaf, no 'Params'
+ -- Thread the running index left-to-right across a member list.
+ goList :: forall x y. Int -> [Program x y] -> ([Program x y], Int)
+ goList idx [] = ([], idx)
+ goList idx (q : qs) =
+ let (q', idx') = go idx q
+ (qs', idx'') = goList idx' qs
+ in (q' : qs', idx'')
+
+-- ---------------------------------------------------------------------------
+-- Serialization (parameter state only, never closures)
+-- ---------------------------------------------------------------------------
+
+-- | A closure-free description of a program's structure, paired with a saved
+-- parameter vector to verify it loads onto the program it was saved from. It
+-- records the constructor tree and a per-'Predict' label; an 'FMap' node's mapped
+-- function is intentionally omitted (opaque, unserializable).
+data ProgramShape
+ = -- | a 'Predict' node, labeled by its joined output-field names
+ ShapePredict !Text
+ | ShapeCompose !ProgramShape !ProgramShape
+ | -- | an 'FMap' node; the function is opaque and omitted
+ ShapeFMap !ProgramShape
+ | -- | a 'Map' node, carrying its concurrency width
+ ShapeMap !Int !ProgramShape
+ | ShapeParallel !ProgramShape !ProgramShape
+ | ShapeRetry !Int !ProgramShape
+ | -- | a 'RetryWhen' node; the predicate is opaque and omitted
+ ShapeRetryWhen !Int !ProgramShape
+ | -- | a 'Validate' node; the validator is opaque and omitted
+ ShapeValidate !ProgramShape
+ | ShapeMajorityVote !Int !TempSchedule !ProgramShape
+ | -- | an 'Ensemble' node; the reducer is opaque and omitted
+ ShapeEnsemble ![ProgramShape]
+ | -- | an 'Embed' node; the embedded effectful body is opaque and omitted
+ ShapeEmbed
+ deriving stock (Eq, Show, Generic)
+
+instance ToJSON ProgramShape
+
+instance FromJSON ProgramShape
+
+-- | Why a parameter vector could not be applied to a program.
+newtype ProgramShapeError
+ = -- | the vector's length did not match the program's node count (expected, got)
+ ParamCountMismatch (Int, Int)
+ deriving stock (Eq, Show, Generic)
+
+instance ToJSON ProgramShapeError
+
+instance FromJSON ProgramShapeError
+
+-- | Extract a program's closure-free structural shape. Stable across parameter
+-- changes (parameters do not affect shape).
+programShape :: Program i o -> ProgramShape
+programShape (Predict sig _) = ShapePredict (sigLabel sig)
+programShape (Compose a b) = ShapeCompose (programShape a) (programShape b)
+programShape (FMap _ p) = ShapeFMap (programShape p)
+programShape (Map w p) = ShapeMap w (programShape p)
+programShape (Parallel a b) = ShapeParallel (programShape a) (programShape b)
+programShape (Retry n p) = ShapeRetry n (programShape p)
+programShape (RetryWhen _ n p) = ShapeRetryWhen n (programShape p)
+programShape (Validate _ p) = ShapeValidate (programShape p)
+programShape (MajorityVote k sched p) = ShapeMajorityVote k sched (programShape p)
+programShape (Ensemble ps _) = ShapeEnsemble (map programShape ps)
+programShape (Embed _) = ShapeEmbed
+
+-- | A stable, parameter-independent label for a 'Predict' node: its output-field
+-- names joined. (EP-3 exposes no @signatureName@; the field names are the stable
+-- structural identity available.)
+sigLabel :: Signature i o -> Text
+sigLabel sig = T.intercalate "," (map fieldName (outputFields sig))
+
+-- | The ordered parameter vector, in 'foldParams' order — JSON-serializable
+-- because 'Params'/'Demo' are. Saving an optimized program = write
+-- @(programShape p, programParams p)@; loading = read the @[Params]@, reconstruct
+-- @p@ in code, then 'setProgramParams'.
+programParams :: Program i o -> [Params]
+programParams = foldParams
+
+-- | Apply a saved parameter vector onto a program of the matching shape, replacing
+-- each node's 'Params' in 'foldParams' order. The vector must have exactly one
+-- entry per 'Predict' node; a length mismatch is a 'ParamCountMismatch' 'Left'.
+setProgramParams :: [Params] -> Program i o -> Either ProgramShapeError (Program i o)
+setProgramParams ps prog
+ | length ps /= n = Left (ParamCountMismatch (n, length ps))
+ | otherwise = Right (fst (go ps prog))
+ where
+ n = length (foldParams prog)
+ go :: forall x y. [Params] -> Program x y -> (Program x y, [Params])
+ go (q : qs) (Predict sig _) = (Predict sig q, qs)
+ go qs (Predict sig old) = (Predict sig old, qs) -- unreachable after the length guard
+ go qs (Compose a b) =
+ let (a', qs') = go qs a
+ (b', qs'') = go qs' b
+ in (Compose a' b', qs'')
+ go qs (FMap k p) =
+ let (p', qs') = go qs p
+ in (FMap k p', qs')
+ go qs (Map w p) =
+ let (p', qs') = go qs p
+ in (Map w p', qs')
+ go qs (Parallel a b) =
+ let (a', qs') = go qs a
+ (b', qs'') = go qs' b
+ in (Parallel a' b', qs'')
+ go qs (Retry m p) =
+ let (p', qs') = go qs p
+ in (Retry m p', qs')
+ go qs (RetryWhen ok m p) =
+ let (p', qs') = go qs p
+ in (RetryWhen ok m p', qs')
+ go qs (Validate v p) =
+ let (p', qs') = go qs p
+ in (Validate v p', qs')
+ go qs (MajorityVote k sched p) =
+ let (p', qs') = go qs p
+ in (MajorityVote k sched p', qs')
+ go qs (Ensemble ms reduce) =
+ let (ms', qs') = goList qs ms
+ in (Ensemble ms' reduce, qs')
+ go qs (Embed f) = (Embed f, qs) -- leaf, consumes no 'Params'
+ goList :: forall x y. [Params] -> [Program x y] -> ([Program x y], [Params])
+ goList qs [] = ([], qs)
+ goList qs (m : ms) =
+ let (m', qs') = go qs m
+ (ms', qs'') = goList qs' ms
+ in (m' : ms', qs'')
diff --git a/src/Shikumi/Refine.hs b/src/Shikumi/Refine.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Refine.hs
@@ -0,0 +1,385 @@
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+-- | Reward-driven self-refinement modules (EP-18): three ways to wrap an existing
+-- 'Program' so its re-runs are /steered by how good the answer is/.
+--
+-- * 'bestOfN' — run the inner program @N@ times at spread temperatures, score
+-- each output with a 'Reward', and keep the best (short-circuiting on a pass).
+-- * 'refine' — run the inner program; on a sub-threshold output, ask an LM to
+-- write a textual critique ("advice") and feed it into the next attempt.
+-- * 'multiChainComparison' — run @M@ reasoning attempts, then make one final LM
+-- call shown all @M@ candidates and asked to synthesize a consensus answer.
+--
+-- Each module is itself an ordinary 'Program' built from the existing 'embed'
+-- ('Shikumi.Program.Embed') leaf — /no new GADT constructor/ — so it runs under the
+-- unchanged 'runProgram'/'runProgramConc', composes with every combinator, holds no
+-- optimizable 'Shikumi.Program.Params' of its own (it serializes as
+-- 'Shikumi.Program.ShapeEmbed', exactly like @FMap@), and needs no edit to
+-- @paramsTraversal@/@programShape@/@setProgramParams@. This is the same pattern V1's
+-- ReAct agent uses. The cost, accepted: the inner program's per-node params are not
+-- visible to an optimizer /through/ the wrapper (the @Embed@ body is opaque) — fine
+-- for inference-time selection\/retry modules.
+--
+-- Per-sample temperature is requested through the same @TempSchedule@ /
+-- 'Shikumi.Program.withSampleTemp' channel EP-14 makes live: each attempt stamps its
+-- temperature onto the private request-metadata channel and the router
+-- ("Shikumi.Routing".@routeLLM@) sets @Options.temperature@ from it. With no router
+-- installed the stamp is inert (every attempt sampled identically — correct, just
+-- less diverse); the hermetic stub honours the stamped temperature directly.
+module Shikumi.Refine
+ ( -- * The reward vocabulary (re-exported from "Shikumi.Reward")
+ Reward (..),
+ mkReward,
+ boolReward,
+
+ -- * Self-refinement modules
+ bestOfN,
+ bestOfNWith,
+ refine,
+ refineWith,
+ multiChainComparison,
+ MultiChainInput (..),
+ multiChainSig,
+
+ -- * Defaults
+ defaultSpread,
+ )
+where
+
+import Control.Lens ((&), (.~), (^.))
+import Data.Aeson (Object, Value (Object))
+import Data.Aeson.Key qualified as Key
+import Data.Aeson.KeyMap qualified as KM
+import Data.Generics.Labels ()
+import Data.Maybe (fromMaybe)
+import Data.Text (Text)
+import Data.Text qualified as T
+import Effectful (Eff, (:>))
+import Effectful.Dispatch.Dynamic (interpose, passthrough)
+import Effectful.Error.Static (Error, catchError, throwError)
+import GHC.Generics (Generic, Rep)
+import Shikumi.Adapter (ToPrompt (..), toPrompt)
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.LLM (Context, LLM (..), complete)
+import Shikumi.Module (WithReasoning (..), predict)
+import Shikumi.Program
+ ( Program,
+ TempSchedule (..),
+ embed,
+ runProgram,
+ sampleTemps,
+ withSampleTemp,
+ )
+import Shikumi.Reward (Reward (..), boolReward, mkReward)
+import Shikumi.Schema
+ ( FromModel (..),
+ GFieldMetas,
+ ToSchema,
+ Validatable,
+ fieldMetasOf,
+ )
+import Shikumi.Schema.Types (FieldPath, pushField, renderPath)
+import Shikumi.Signature (Signature (..), mkSignature)
+
+-- ---------------------------------------------------------------------------
+-- Shared building blocks
+-- ---------------------------------------------------------------------------
+
+-- | The default per-attempt temperature spread: centred at 0.7, fanning out by
+-- 0.6 — the same convention @majorityVote@ uses.
+defaultSpread :: TempSchedule
+defaultSpread = TempSpread 0.7 0.6
+
+-- | Run an effect that may throw a 'ShikumiError', recovering the error as a
+-- 'Left' instead of propagating it (so a single failed attempt does not abort the
+-- whole loop).
+tryShikumi :: (Error ShikumiError :> es) => Eff es a -> Eff es (Either ShikumiError a)
+tryShikumi act = (Right <$> act) `catchError` \_cs e -> pure (Left e)
+
+-- | Keep the higher-reward of the running best and a fresh @(output, reward)@.
+-- Ties keep the incumbent (first-seen wins).
+keepBest :: Maybe (o, Double) -> o -> Double -> Maybe (o, Double)
+keepBest Nothing o r = Just (o, r)
+keepBest (Just (bo, br)) o r = if r > br then Just (o, r) else Just (bo, br)
+
+-- | Render a value as 'Text' via 'Show'.
+tshow :: (Show a) => a -> Text
+tshow = T.pack . show
+
+-- ---------------------------------------------------------------------------
+-- bestOfN
+-- ---------------------------------------------------------------------------
+
+-- | Run @inner@ up to @N@ times at spread temperatures, score each output, and
+-- return the highest-scoring one — short-circuiting as soon as an output's reward
+-- reaches @threshold@. Defaults: 'defaultSpread' temperatures and a @failCount@
+-- (tolerated inner-program errors) of @N@.
+bestOfN ::
+ -- | @N@ attempts (clamped to @>= 1@)
+ Int ->
+ -- | pass threshold (reward @>=@ this short-circuits)
+ Double ->
+ Reward o ->
+ Program i o ->
+ Program i o
+bestOfN n = bestOfNWith n (max 1 n) defaultSpread
+
+-- | 'bestOfN' with an explicit @failCount@ budget and 'TempSchedule'.
+bestOfNWith ::
+ -- | @N@ attempts (clamped to @>= 1@)
+ Int ->
+ -- | @failCount@: tolerated inner-program errors before the error propagates
+ Int ->
+ TempSchedule ->
+ -- | pass threshold
+ Double ->
+ Reward o ->
+ Program i o ->
+ Program i o
+bestOfNWith n failCount sched threshold reward inner = embed $ \i ->
+ let temps = sampleTemps (max 1 n) sched
+ go [] best lastErr _ = finish best lastErr
+ go (t : ts) best lastErr budget = do
+ res <- tryShikumi (withSampleTemp t (runProgram inner i))
+ case res of
+ Left e ->
+ let budget' = budget - 1
+ in if budget' < 0 then throwError e else go ts best (Just e) budget'
+ Right o ->
+ let r = runReward reward o
+ best' = keepBest best o r
+ in if r >= threshold then pure o else go ts best' lastErr budget
+ finish (Just (o, _)) _ = pure o
+ finish Nothing (Just e) = throwError e
+ finish Nothing Nothing = throwError (ProviderFailure "Shikumi.Refine.bestOfN: no attempt produced an output")
+ in go temps Nothing Nothing failCount
+
+-- ---------------------------------------------------------------------------
+-- refine
+-- ---------------------------------------------------------------------------
+
+-- | The input to the advice generator: the achieved reward and the target
+-- threshold (rendered as text, since the inner program's input is opaque to the
+-- wrapper). Mirrors DSPy's @OfferFeedback@ predictor, scoped to the whole inner
+-- program.
+data AdviceIn = AdviceIn
+ { achievedReward :: Text,
+ targetThreshold :: Text
+ }
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema AdviceIn
+
+instance FromModel AdviceIn
+
+instance ToPrompt AdviceIn
+
+instance Validatable AdviceIn
+
+-- | The advice generator's output: one concrete piece of advice for the next
+-- attempt.
+newtype AdviceOut = AdviceOut {advice :: Text}
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema AdviceOut
+
+instance FromModel AdviceOut
+
+instance ToPrompt AdviceOut
+
+instance Validatable AdviceOut
+
+adviceSig :: Signature AdviceIn AdviceOut
+adviceSig =
+ mkSignature
+ "You are reviewing a previous attempt that scored below its target. \
+ \Look at the achieved reward and the target threshold, then write one concrete \
+ \piece of advice that would help the next attempt score higher."
+
+-- | Like 'bestOfN', but on a sub-threshold output it derives textual advice from
+-- the reward (via an LM call) and threads it into the next attempt. Returns the
+-- best output seen across all attempts. Defaults: @failCount = N@.
+refine ::
+ -- | @N@ attempts (clamped to @>= 1@)
+ Int ->
+ -- | pass threshold
+ Double ->
+ Reward o ->
+ Program i o ->
+ Program i o
+refine n = refineWith n (max 1 n)
+
+-- | 'refine' with an explicit @failCount@ budget.
+refineWith ::
+ -- | @N@ attempts (clamped to @>= 1@)
+ Int ->
+ -- | @failCount@: tolerated inner-program errors before the error propagates
+ Int ->
+ -- | pass threshold
+ Double ->
+ Reward o ->
+ Program i o ->
+ Program i o
+refineWith n failCount threshold reward inner = embed $ \i ->
+ let cnt = max 1 n
+ go k best lastErr mAdvice budget
+ | k >= cnt = finish best lastErr
+ | otherwise = do
+ res <- tryShikumi (withAdvice mAdvice (runProgram inner i))
+ case res of
+ Left e ->
+ let budget' = budget - 1
+ in if budget' < 0
+ then throwError e
+ else go (k + 1) best (Just e) mAdvice budget'
+ Right o ->
+ let r = runReward reward o
+ best' = keepBest best o r
+ in if r >= threshold
+ then pure o
+ else
+ if k == cnt - 1
+ then finish best' lastErr
+ else do
+ adv <- generateAdvice r threshold
+ go (k + 1) best' lastErr (Just adv) budget
+ finish (Just (o, _)) _ = pure o
+ finish Nothing (Just e) = throwError e
+ finish Nothing Nothing = throwError (ProviderFailure "Shikumi.Refine.refine: no attempt produced an output")
+ in go 0 Nothing Nothing Nothing failCount
+
+-- | Ask the advice generator for one piece of advice, given the achieved reward
+-- and the target threshold.
+generateAdvice ::
+ (LLM :> es, Error ShikumiError :> es) =>
+ Double ->
+ Double ->
+ Eff es Text
+generateAdvice r threshold = do
+ out <- runProgram (predict adviceSig) (AdviceIn (tshow r) (tshow threshold))
+ pure (advice out)
+
+-- | Run an action with a previous attempt's advice injected into every outgoing
+-- request: the advice is appended to the rendered system prompt as a "Hint from a
+-- previous attempt" line, so the model (and the hermetic stub) sees it. 'Nothing'
+-- runs the action unchanged. Scoped via 'interpose' so it only affects this
+-- attempt's calls.
+withAdvice :: (LLM :> es) => Maybe Text -> Eff es a -> Eff es a
+withAdvice Nothing act = act
+withAdvice (Just adv) act =
+ interpose
+ ( \env -> \case
+ Complete m ctx o -> complete m (appendHint adv ctx) o
+ other -> passthrough env other
+ )
+ act
+
+-- | Append an advice hint to a request's system prompt.
+appendHint :: Text -> Context -> Context
+appendHint adv ctx =
+ ctx & #systemPrompt .~ Just (cur <> "\n\nHint from a previous attempt: " <> adv)
+ where
+ cur = fromMaybe "" (ctx ^. #systemPrompt)
+
+-- ---------------------------------------------------------------------------
+-- multiChainComparison
+-- ---------------------------------------------------------------------------
+
+-- | The synthesis input: the original input plus the @M@ candidate
+-- @(reasoning, answer)@ attempts. Its instances are hand-written (not
+-- @Generic@-derived) because @attempts@ is polymorphic in @o@ — the same reason
+-- 'WithReasoning'\'s instances are hand-written.
+data MultiChainInput i o = MultiChainInput
+ { original :: i,
+ -- | length @M@; each rendered as "Student Attempt #k"
+ attempts :: [WithReasoning o]
+ }
+
+-- | Render the original input, then each attempt as a "Student Attempt #k" line:
+-- @"I tried to ; my answer is "@ (mirroring DSPy's rendering).
+instance (ToPrompt i, ToPrompt o) => ToPrompt (MultiChainInput i o) where
+ toPromptFields mci =
+ toPromptFields (original mci)
+ ++ [ ("Student Attempt #" <> tshow k, renderAttempt a)
+ | (k, a) <- zip [1 :: Int ..] (attempts mci)
+ ]
+ where
+ renderAttempt a = "I tried to " <> reasoning a <> "; my answer is " <> toPrompt (value a)
+
+ -- Fields are polymorphic in @i@/@o@, so the generic image walk cannot resolve; a
+ -- synthesis input is text-only and carries no image fields.
+ imageFields _ = []
+ imageFieldNames _ = []
+
+-- | A 'MultiChainInput' is only ever an /input/ to the synthesis 'predict' node,
+-- so this decoder exists to satisfy the 'Shikumi.Program.Predict' constraint (used
+-- when decoding demos, of which the synthesis node has none).
+instance (FromModel i, FromModel o) => FromModel (MultiChainInput i o) where
+ fromModelP path = \case
+ Object obj ->
+ MultiChainInput
+ <$> getField path "original" obj
+ <*> getField path "attempts" obj
+ _ -> Left (SchemaMismatch (renderPath path <> ": expected object"))
+
+instance Validatable (MultiChainInput i o)
+
+-- | Look up a required field in a JSON object, locating a miss precisely.
+getField :: (FromModel a) => FieldPath -> Text -> Object -> Either ShikumiError a
+getField path nm o = case KM.lookup (Key.fromText nm) o of
+ Nothing -> Left (MissingField (renderPath (pushField nm path)))
+ Just v -> fromModelP (pushField nm path) v
+
+-- | Build a synthesis signature for 'multiChainComparison'. The input is rendered
+-- via 'MultiChainInput'\'s 'ToPrompt' (so the input-field metadata is empty); the
+-- output-field metadata is derived from @o2@ so the output guide and demo
+-- rendering name the right fields.
+multiChainSig ::
+ forall i o o2.
+ (GFieldMetas (Rep o2)) =>
+ -- | the synthesis instruction
+ Text ->
+ Signature (MultiChainInput i o) o2
+multiChainSig instr =
+ Signature
+ { instruction = instr,
+ demos = [],
+ inputFields = [],
+ outputFields = fieldMetasOf @o2
+ }
+
+-- | Run @M@ independent reasoning attempts (each a @(reasoning, answer)@ via a
+-- @chainOfThoughtRaw@-shaped program) at spread temperatures, then make one final
+-- synthesis 'predict' call shown all the candidates and asked to produce a single
+-- corrected answer. Attempts that error are dropped; synthesis proceeds with the
+-- survivors (at least one — if every attempt errors, the last error propagates).
+multiChainComparison ::
+ ( FromModel i,
+ ToPrompt i,
+ FromModel o,
+ ToPrompt o,
+ FromModel o2,
+ ToSchema o2,
+ Validatable o2,
+ ToPrompt o2
+ ) =>
+ -- | @M@ reasoning attempts (clamped to @>= 1@)
+ Int ->
+ -- | the reasoning program (e.g. @chainOfThoughtRaw sig@)
+ Program i (WithReasoning o) ->
+ -- | the synthesis signature
+ Signature (MultiChainInput i o) o2 ->
+ Program i o2
+multiChainComparison m reasoner synthSig = embed $ \i -> do
+ let temps = sampleTemps (max 1 m) defaultSpread
+ results <- traverse (\t -> tryShikumi (withSampleTemp t (runProgram reasoner i))) temps
+ let oks = [o | Right o <- results]
+ errs = [e | Left e <- results]
+ case oks of
+ [] -> case errs of
+ (e : _) -> throwError e
+ [] -> throwError (ProviderFailure "Shikumi.Refine.multiChainComparison: no reasoning attempts")
+ _ -> runProgram (predict synthSig) (MultiChainInput {original = i, attempts = oks})
diff --git a/src/Shikumi/Reward.hs b/src/Shikumi/Reward.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Reward.hs
@@ -0,0 +1,37 @@
+-- | The reward vocabulary owned by EP-18 (MasterPlan integration point #1): a
+-- function that scores a single program output, plus the smart constructors used
+-- to build one.
+--
+-- A 'Reward' returns a plain 'Double' (higher is better; by convention in @[0, 1]@
+-- but not clamped — the self-refinement modules only ever /compare/ rewards, so any
+-- totally-ordered range works). It lives in the base @shikumi@ package — not in
+-- @shikumi-eval@ alongside @Score@/@Metric@/@Prediction@ — because the modules that
+-- consume it ('Shikumi.Refine') are first-class 'Shikumi.Program.Program' values
+-- and @Program@ lives here; reusing @shikumi-eval@'s @Metric@ would invert the
+-- one-way @shikumi-eval -> shikumi@ dependency. A caller who already has a
+-- @shikumi-eval@ @Metric@ derives a 'Reward' from it with the one-line
+-- @rewardFromMetric@ adapter that lives in @shikumi-eval@ (where both types are in
+-- scope), not here.
+--
+-- A reward at inference time only ever sees the /one/ output it is scoring (not a
+-- held-out expected value and not a multi-sample prediction), so the simple
+-- @o -> Double@ shape is the honest one. EP-22 (GEPA,
+-- @docs/plans/22-gepa-reflective-optimizer.md@) reuses this same vocabulary for its
+-- reflective feedback and must not define a parallel reward type.
+module Shikumi.Reward
+ ( Reward (..),
+ mkReward,
+ boolReward,
+ )
+where
+
+-- | Score a single output. Higher is better.
+newtype Reward o = Reward {runReward :: o -> Double}
+
+-- | Build a reward from a scoring function.
+mkReward :: (o -> Double) -> Reward o
+mkReward = Reward
+
+-- | Build a reward from a pass/fail predicate: @True -> 1.0@, @False -> 0.0@.
+boolReward :: (o -> Bool) -> Reward o
+boolReward p = Reward (\o -> if p o then 1.0 else 0.0)
diff --git a/src/Shikumi/Routing.hs b/src/Shikumi/Routing.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Routing.hs
@@ -0,0 +1,113 @@
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE TypeFamilies #-}
+
+-- | Ambient model routing (EP-14) — the keystone that lets
+-- 'Shikumi.Program.runProgram' dispatch a /named/ model without changing its
+-- pinned signature (MasterPlan integration point #4:
+-- @runProgram :: (LLM :> es, Error ShikumiError :> es) => Program i o -> i -> Eff es o@).
+--
+-- The mechanism is approach (a) from the parent MasterPlan: a re-interpreting
+-- router installed /below/ 'Shikumi.Program.runProgram' on the existing @LLM@
+-- effect — exactly the seam @cachedLLM@ ("Shikumi.Cache") and @tracedLLM@
+-- ("Shikumi.Trace") already use. 'runProgram' renders each 'Predict' node
+-- model-agnostically against the inert placeholder model and stamps its intentions
+-- (the derived JSON schema, and any per-sample temperature) onto the private
+-- request-metadata channel ('Shikumi.Adapter.attachSchema' /
+-- 'Shikumi.Adapter.stampTemperature'). 'routeLLM' then reads the /ambient/ model
+-- supplied by 'runRouting', overwrites the placeholder with it, translates the
+-- stamped metadata into @Options.responseFormat@ (for native-capable models only)
+-- and @Options.temperature@, strips the private keys, and forwards the call to the
+-- real @LLM@ interpreter beneath it.
+--
+-- Install order (mirroring @runTrace . runKeyedLLM . tracedLLM@): 'runRouting' is
+-- /outer/ of the real @LLM@ interpreter, which is /outer/ of 'routeLLM' (the
+-- innermost, closest to 'runProgram'). For example:
+--
+-- @
+-- runEff
+-- . runErrorNoCallStack \@ShikumiError
+-- . runRouting openai_gpt_4o_mini -- supplies the ambient model
+-- . runLLMResilient cfg -- the real LLM interpreter (or a stub)
+-- . routeLLM -- rewrites outgoing LLM calls
+-- $ runProgram program input
+-- @
+module Shikumi.Routing
+ ( -- * The effect
+ Routing (..),
+ currentModel,
+ runRouting,
+
+ -- * The router
+ routeLLM,
+ )
+where
+
+import Baikai (Model, Options, ResponseFormat (..))
+import Control.Lens ((&), (.~), (^.))
+import Data.Aeson (Result (..), Value, fromJSON)
+import Data.Generics.Labels ()
+import Data.Map.Strict qualified as Map
+import Effectful (Dispatch (Dynamic), DispatchOf, Eff, Effect, (:>))
+import Effectful.Dispatch.Dynamic (interpose, interpret, passthrough, send)
+import Shikumi.Adapter
+ ( ModelCapability (..),
+ capabilityFor,
+ metaResponseSchemaKey,
+ metaTemperatureKey,
+ )
+import Shikumi.LLM (LLM (..), complete)
+
+-- | The ambient model-routing effect. Its single operation reads the model every
+-- 'Predict' node should dispatch against. It is supplied by an interpreter at the
+-- bottom of the stack ('runRouting'), exactly as @LLM@ and @Error@ already are, so
+-- it never appears in 'Shikumi.Program.runProgram'\'s constraint row.
+data Routing :: Effect where
+ CurrentModel :: Routing m Model
+
+type instance DispatchOf Routing = 'Dynamic
+
+-- | Read the ambient model.
+currentModel :: (Routing :> es) => Eff es Model
+currentModel = send CurrentModel
+
+-- | Supply a constant ambient model to everything above. A pure reader: every
+-- 'currentModel' resolves to @m@.
+runRouting :: Model -> Eff (Routing : es) a -> Eff es a
+runRouting m = interpret (\_ CurrentModel -> pure m)
+
+-- | The router. Re-interpret the @LLM@ effect so each outgoing 'Complete' is
+-- dispatched against the ambient model (overwriting the placeholder model the
+-- model-agnostic 'Shikumi.Program.runProgram' passes), with the private metadata
+-- channel translated to real wire options and stripped. 'Stream' is forwarded
+-- unchanged ('runProgram' never streams).
+routeLLM :: (Routing :> es, LLM :> es) => Eff es a -> Eff es a
+routeLLM = interpose $ \env -> \case
+ Complete _placeholder ctx opts -> do
+ m <- currentModel
+ complete m ctx (translateForWire m opts)
+ other -> passthrough env other
+
+-- | Realize the private request-metadata channel against the real model: attach a
+-- native @responseFormat@ when the schema was stamped /and/ the model is
+-- native-capable; set @temperature@ when one was stamped; then strip both private
+-- keys so nothing private reaches the transport.
+translateForWire :: Model -> Options -> Options
+translateForWire m opts =
+ let md = opts ^. #metadata
+ mSchema = Map.lookup metaResponseSchemaKey md
+ mTemp = Map.lookup metaTemperatureKey md >>= valueToDouble
+ stripped = Map.delete metaResponseSchemaKey (Map.delete metaTemperatureKey md)
+ withSchema o = case (mSchema, capabilityFor m) of
+ (Just s, NativeSchema) ->
+ o & #responseFormat .~ Just (JsonSchema {name = "output", schema = s, strict = True})
+ _ -> o
+ withTemp o = case mTemp of
+ Just t -> o & #temperature .~ Just t
+ Nothing -> o
+ in withTemp (withSchema (opts & #metadata .~ stripped))
+
+-- | Read a JSON number back into a 'Double' (the per-sample temperature).
+valueToDouble :: Value -> Maybe Double
+valueToDouble v = case fromJSON v of
+ Success d -> Just d
+ Error _ -> Nothing
diff --git a/src/Shikumi/Schema.hs b/src/Shikumi/Schema.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Schema.hs
@@ -0,0 +1,458 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DefaultSignatures #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE UndecidableInstances #-}
+
+-- | The schema engine: derive a provider JSON Schema from a Haskell record
+-- ('ToSchema' / 'deriveSchema'), and totally decode a provider JSON reply back
+-- into a record ('FromModel' / 'parseOutput'), where every failure becomes a
+-- precise, located 'Shikumi.Error.ShikumiError'.
+--
+-- Both directions are hand-rolled on 'GHC.Generics' (the schema we emit is the
+-- small subset OpenAI/Anthropic structured-output endpoints accept). Per-field
+-- descriptions flow from the 'Field' wrapper's type-level 'Symbol'.
+module Shikumi.Schema
+ ( -- * Forward: record -> JSON Schema
+ ToSchema (..),
+ deriveSchema,
+
+ -- * Inverse: JSON -> record (total)
+ FromModel (..),
+ fromModel,
+ fromModelChecked,
+ parseOutput,
+
+ -- * Validation hook
+ Validatable (..),
+
+ -- * Declarative field constraints (EP-26)
+ ReflectConstraints (..),
+
+ -- * Field-metadata traversal (shared with "Shikumi.Signature")
+ GFieldMetas (..),
+ fieldMetasOf,
+
+ -- * Re-exports
+ module Shikumi.Schema.Types,
+ )
+where
+
+import Control.Applicative ((<|>))
+import Data.Aeson (Object, Value (..), eitherDecodeStrict)
+import Data.Aeson.Key qualified as Key
+import Data.Aeson.KeyMap qualified as KM
+import Data.Bifunctor (first)
+import Data.Kind (Type)
+import Data.Proxy (Proxy (..))
+import Data.Scientific qualified as Sci
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Text.Encoding (encodeUtf8)
+import Data.Vector (Vector)
+import Data.Vector qualified as V
+import GHC.Generics
+import GHC.TypeLits (KnownNat, KnownSymbol, Symbol, natVal, symbolVal)
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.Schema.Types
+import Text.Read (readMaybe)
+
+-- ---------------------------------------------------------------------------
+-- Forward direction: ToSchema
+-- ---------------------------------------------------------------------------
+
+-- | Types that can describe themselves as a provider JSON Schema. Records get a
+-- generic default; leaves and containers have explicit instances.
+class ToSchema a where
+ toSchema :: Proxy a -> Value
+ default toSchema :: (GToSchema (Rep a)) => Proxy a -> Value
+ toSchema _ = gToSchema (Proxy @(Rep a))
+
+-- | The JSON Schema for @a@.
+deriveSchema :: forall a. (ToSchema a) => Value
+deriveSchema = toSchema (Proxy @a)
+
+instance ToSchema Text where toSchema _ = stringSchema
+
+instance ToSchema Int where toSchema _ = integerSchema
+
+instance ToSchema Integer where toSchema _ = integerSchema
+
+instance ToSchema Double where toSchema _ = numberSchema
+
+instance ToSchema Bool where toSchema _ = boolSchema
+
+instance (ToSchema a) => ToSchema [a] where toSchema _ = arraySchema (toSchema (Proxy @a))
+
+instance (ToSchema a) => ToSchema (Vector a) where toSchema _ = arraySchema (toSchema (Proxy @a))
+
+instance (ToSchema a) => ToSchema (Maybe a) where toSchema _ = nullableSchema (toSchema (Proxy @a))
+
+-- | The description is attached at the record-field layer ('FieldSchema'), not
+-- here, so this instance stays composable (e.g. a nested @Field@).
+instance (ToSchema a) => ToSchema (Field d a) where toSchema _ = toSchema (Proxy @a)
+
+-- | Generic schema walk. A single record constructor becomes an object schema; a
+-- sum of nullary constructors becomes a string @enum@.
+class GToSchema (f :: Type -> Type) where
+ gToSchema :: Proxy f -> Value
+
+instance (GToSchema cs) => GToSchema (D1 d cs) where
+ gToSchema _ = gToSchema (Proxy @cs)
+
+instance (GRecordFields fs) => GToSchema (C1 ('MetaCons n fx 'True) fs) where
+ gToSchema _ = uncurry objectSchema (gRecordFields (Proxy @fs))
+
+instance (GEnumNames (l :+: r)) => GToSchema (l :+: r) where
+ gToSchema _ = enumSchema (gEnumNames (Proxy @(l :+: r)))
+
+-- | Collect @(name, schema)@ properties and the required-name list for a record.
+class GRecordFields (f :: Type -> Type) where
+ gRecordFields :: Proxy f -> ([(Text, Value)], [Text])
+
+instance (GRecordFields a, GRecordFields b) => GRecordFields (a :*: b) where
+ gRecordFields _ = gRecordFields (Proxy @a) <> gRecordFields (Proxy @b)
+
+instance (Selector s, FieldSchema t) => GRecordFields (S1 s (K1 R t)) where
+ gRecordFields _ =
+ let nm = T.pack (selName (undefined :: S1 s (K1 R t) ()))
+ (sch, req) = fieldSchema (Proxy @t)
+ in ([(nm, sch)], [nm | req])
+
+-- | Per-field schema and whether the field is required. The general
+-- (overlappable) case is "required, no description"; 'Maybe' is optional/nullable;
+-- 'Field' attaches the description and preserves the inner required-ness.
+class FieldSchema t where
+ fieldSchema :: Proxy t -> (Value, Bool)
+
+instance {-# OVERLAPPABLE #-} (ToSchema a) => FieldSchema a where
+ fieldSchema _ = (toSchema (Proxy @a), True)
+
+instance (ToSchema a) => FieldSchema (Maybe a) where
+ fieldSchema _ = (nullableSchema (toSchema (Proxy @a)), False)
+
+instance (KnownSymbol d, FieldSchema a) => FieldSchema (Field d a) where
+ fieldSchema _ =
+ let (s, req) = fieldSchema (Proxy @a)
+ in (withDescription (T.pack (symbolVal (Proxy @d))) s, req)
+
+-- | A 'Constrained' field emits its constraints' JSON-Schema keywords onto the
+-- inner type's schema, preserving the inner required-ness (EP-26).
+instance forall cs a. (ReflectConstraints cs a, FieldSchema a) => FieldSchema (Constrained cs a) where
+ fieldSchema _ =
+ let (s, req) = fieldSchema (Proxy @a)
+ in (constraintSchema @cs @a Proxy s, req)
+
+-- | Constructor names of a sum of nullary constructors.
+class GEnumNames (f :: Type -> Type) where
+ gEnumNames :: Proxy f -> [Text]
+
+instance (GEnumNames a, GEnumNames b) => GEnumNames (a :+: b) where
+ gEnumNames _ = gEnumNames (Proxy @a) ++ gEnumNames (Proxy @b)
+
+instance (Constructor c) => GEnumNames (C1 c U1) where
+ gEnumNames _ = [T.pack (conName (undefined :: C1 c U1 ()))]
+
+-- ---------------------------------------------------------------------------
+-- Inverse direction: FromModel (total decode)
+-- ---------------------------------------------------------------------------
+
+-- | Types decodable from provider JSON, totally. The path-carrying 'fromModelP'
+-- threads a breadcrumb so every error is located; 'fromModel' starts at the root.
+class FromModel a where
+ fromModelP :: FieldPath -> Value -> Either ShikumiError a
+ default fromModelP :: (Generic a, GFromModel (Rep a)) => FieldPath -> Value -> Either ShikumiError a
+ fromModelP path v = to <$> gFromModel path v
+
+-- | Decode from a JSON 'Value' (no validation).
+fromModel :: (FromModel a) => Value -> Either ShikumiError a
+fromModel = fromModelP []
+
+-- | Decode then run the type's 'Validatable' rule, locating a failure as a
+-- 'ValidationFailure'.
+fromModelChecked :: (FromModel a, Validatable a) => Value -> Either ShikumiError a
+fromModelChecked v = do
+ a <- fromModel v
+ first ValidationFailure (validate a)
+
+-- | Parse a raw response body to JSON (mapping a parse failure to 'InvalidJSON'),
+-- decode it, and validate.
+parseOutput :: forall a. (FromModel a, Validatable a) => Text -> Either ShikumiError a
+parseOutput body = case eitherDecodeStrict (encodeUtf8 body) of
+ Left e -> Left (InvalidJSON (T.pack e))
+ Right v -> fromModelChecked v
+
+instance FromModel Text where
+ fromModelP path = \case String t -> Right t; v -> mismatch path "string" v
+
+instance FromModel Int where
+ fromModelP path = \case
+ Number n -> maybe (mismatch path "integer" (Number n)) Right (Sci.toBoundedInteger n)
+ v -> mismatch path "integer" v
+
+instance FromModel Integer where
+ fromModelP path = \case
+ Number n -> case Sci.floatingOrInteger n :: Either Double Integer of
+ Right i -> Right i
+ Left _ -> mismatch path "integer" (Number n)
+ v -> mismatch path "integer" v
+
+instance FromModel Double where
+ fromModelP path = \case Number n -> Right (Sci.toRealFloat n); v -> mismatch path "number" v
+
+instance FromModel Bool where
+ fromModelP path = \case Bool b -> Right b; v -> mismatch path "boolean" v
+
+instance (FromModel a) => FromModel [a] where
+ fromModelP path = \case
+ Array arr -> traverse (\(i, x) -> fromModelP (pushIndex i path) x) (zip [0 ..] (V.toList arr))
+ v -> mismatch path "array" v
+
+instance (FromModel a) => FromModel (Vector a) where
+ fromModelP path v = V.fromList <$> (fromModelP path v :: Either ShikumiError [a])
+
+instance (FromModel a) => FromModel (Maybe a) where
+ fromModelP _ Null = Right Nothing
+ fromModelP path v = Just <$> fromModelP path v
+
+instance (FromModel a) => FromModel (Field d a) where
+ fromModelP path v = Field <$> fromModelP path v
+
+-- | A 'Constrained' field decodes its inner value, then enforces the constraints
+-- (EP-26). A violation becomes a 'ValidationFailure' located at the field, so the
+-- check composes with any record-level 'Validatable' rule with no change to
+-- 'fromModelChecked'. (@FromField (Constrained cs a)@ is covered by the
+-- overlappable @FromField@ instance via this @FromModel@ instance.)
+instance (FromModel a, ReflectConstraints cs a) => FromModel (Constrained cs a) where
+ fromModelP path v = do
+ a <- fromModelP path v
+ case checkConstraints (Proxy @cs) a of
+ Right ok -> Right (Constrained ok)
+ Left msg -> Left (ValidationFailure (renderPath path <> ": " <> msg))
+
+-- | Generic decode walk, mirroring 'GToSchema'.
+class GFromModel (f :: Type -> Type) where
+ gFromModel :: FieldPath -> Value -> Either ShikumiError (f p)
+
+instance (GFromModel cs) => GFromModel (D1 d cs) where
+ gFromModel path v = M1 <$> gFromModel path v
+
+instance (GFromRecord fs) => GFromModel (C1 ('MetaCons n fx 'True) fs) where
+ gFromModel path = \case
+ Object o -> M1 <$> gFromRecord path o
+ v -> mismatch path "object" v
+
+instance (GEnumParse (l :+: r)) => GFromModel (l :+: r) where
+ gFromModel path = \case
+ String t -> case gEnumParse t of
+ Just x -> Right x
+ Nothing ->
+ Left
+ ( SchemaMismatch
+ (renderPath path <> ": expected one of [" <> T.intercalate ", " (gEnumNamesP (Proxy @(l :+: r))) <> "]")
+ )
+ v -> mismatch path "string" v
+
+-- | Decode a record object field by field, threading the path.
+class GFromRecord (f :: Type -> Type) where
+ gFromRecord :: FieldPath -> Object -> Either ShikumiError (f p)
+
+instance (GFromRecord a, GFromRecord b) => GFromRecord (a :*: b) where
+ gFromRecord path o = (:*:) <$> gFromRecord path o <*> gFromRecord path o
+
+instance (Selector s, FromField t) => GFromRecord (S1 s (K1 R t)) where
+ gFromRecord path o =
+ let nm = T.pack (selName (undefined :: S1 s (K1 R t) ()))
+ in M1 . K1 <$> fromField (pushField nm path) nm o
+
+-- | Look a field up in the JSON object. The general (overlappable) case requires
+-- the field; 'Maybe' and a missing key succeed as 'Nothing'; 'Field' delegates.
+class FromField t where
+ fromField :: FieldPath -> Text -> Object -> Either ShikumiError t
+
+instance {-# OVERLAPPABLE #-} (FromModel t) => FromField t where
+ fromField path nm o = case KM.lookup (Key.fromText nm) o of
+ Nothing -> Left (MissingField (renderPath path))
+ Just v -> fromModelP path v
+
+instance (FromModel a) => FromField (Maybe a) where
+ fromField path nm o = case KM.lookup (Key.fromText nm) o of
+ Nothing -> Right Nothing
+ Just Null -> Right Nothing
+ Just v -> Just <$> fromModelP path v
+
+instance (FromField a) => FromField (Field d a) where
+ fromField path nm o = Field <$> fromField path nm o
+
+-- | Parse a string against the constructor names of an enum-like sum.
+class GEnumParse (f :: Type -> Type) where
+ gEnumParse :: Text -> Maybe (f p)
+ gEnumNamesP :: Proxy f -> [Text]
+
+instance (GEnumParse a, GEnumParse b) => GEnumParse (a :+: b) where
+ gEnumParse t = (L1 <$> gEnumParse t) <|> (R1 <$> gEnumParse t)
+ gEnumNamesP _ = gEnumNamesP (Proxy @a) ++ gEnumNamesP (Proxy @b)
+
+instance (Constructor c) => GEnumParse (C1 c U1) where
+ gEnumParse t
+ | t == nm = Just (M1 U1)
+ | otherwise = Nothing
+ where
+ nm = T.pack (conName (undefined :: C1 c U1 ()))
+ gEnumNamesP _ = [T.pack (conName (undefined :: C1 c U1 ()))]
+
+-- ---------------------------------------------------------------------------
+-- Validation hook
+-- ---------------------------------------------------------------------------
+
+-- | A domain rule applied after a successful decode. Defaults to "always valid";
+-- override per output type to enforce rules (e.g. "bullets has 3-5 elements").
+class Validatable a where
+ validate :: a -> Either Text a
+ validate = Right
+
+-- | Default "always valid" for any type without an explicit rule.
+instance {-# OVERLAPPABLE #-} Validatable a
+
+-- ---------------------------------------------------------------------------
+-- Declarative field constraints (EP-26)
+-- ---------------------------------------------------------------------------
+
+-- | Reflect a type-level constraint list (the @cs@ of a
+-- 'Shikumi.Schema.Types.Constrained') to (a) a schema-decorating function that
+-- inserts the matching JSON-Schema keywords and (b) a runtime validator over the
+-- decoded leaf value. The vocabulary is a small, closed set
+-- ('Shikumi.Schema.Types.Constraint'); instances cover the five constructors over
+-- 'Text' and the numeric leaves.
+class ReflectConstraints (cs :: [Constraint]) a where
+ -- | Compose all keyword inserts onto the inner type's schema.
+ constraintSchema :: Proxy cs -> Value -> Value
+
+ -- | Check the decoded value against every constraint; 'Left' names the first
+ -- rule that failed.
+ checkConstraints :: Proxy cs -> a -> Either Text a
+
+instance ReflectConstraints '[] a where
+ constraintSchema _ = id
+ checkConstraints _ = Right
+
+instance (KnownNat n, ReflectConstraints cs Text) => ReflectConstraints ('MinLen n ': cs) Text where
+ constraintSchema _ = constraintSchema @cs @Text Proxy . withMinLength (fromIntegral (natVal (Proxy @n)))
+ checkConstraints _ t
+ | T.length t < fromIntegral (natVal (Proxy @n)) = Left ("minLength " <> tshow (natVal (Proxy @n)) <> " violated")
+ | otherwise = checkConstraints (Proxy @cs) t
+
+instance (KnownNat n, ReflectConstraints cs Text) => ReflectConstraints ('MaxLen n ': cs) Text where
+ constraintSchema _ = constraintSchema @cs @Text Proxy . withMaxLength (fromIntegral (natVal (Proxy @n)))
+ checkConstraints _ t
+ | T.length t > fromIntegral (natVal (Proxy @n)) = Left ("maxLength " <> tshow (natVal (Proxy @n)) <> " violated")
+ | otherwise = checkConstraints (Proxy @cs) t
+
+instance forall s cs a. (KnownSymbol s, ToScientificLeaf a, ReflectConstraints cs a) => ReflectConstraints ('MinVal s ': cs) a where
+ constraintSchema _ = constraintSchema @cs @a Proxy . withMinimum (parseBound (Proxy @s))
+ checkConstraints _ x
+ | toScientificLeaf x < parseBound (Proxy @s) = Left ("minimum " <> T.pack (symbolVal (Proxy @s)) <> " violated")
+ | otherwise = checkConstraints (Proxy @cs) x
+
+instance forall s cs a. (KnownSymbol s, ToScientificLeaf a, ReflectConstraints cs a) => ReflectConstraints ('MaxVal s ': cs) a where
+ constraintSchema _ = constraintSchema @cs @a Proxy . withMaximum (parseBound (Proxy @s))
+ checkConstraints _ x
+ | toScientificLeaf x > parseBound (Proxy @s) = Left ("maximum " <> T.pack (symbolVal (Proxy @s)) <> " violated")
+ | otherwise = checkConstraints (Proxy @cs) x
+
+instance (KnownSymbols ss, ReflectConstraints cs Text) => ReflectConstraints ('EnumOneOf ss ': cs) Text where
+ constraintSchema _ = constraintSchema @cs @Text Proxy . withEnum (symbolTexts (Proxy @ss))
+ checkConstraints _ t
+ | t `elem` symbolTexts (Proxy @ss) = checkConstraints (Proxy @cs) t
+ | otherwise = Left ("enum membership violated (allowed: " <> T.intercalate ", " (symbolTexts (Proxy @ss)) <> ")")
+
+-- | A leaf type comparable against a 'Sci.Scientific' bound. Closed to the numeric
+-- leaves shikumi decodes.
+class ToScientificLeaf a where
+ toScientificLeaf :: a -> Sci.Scientific
+
+instance ToScientificLeaf Int where toScientificLeaf = fromIntegral
+
+instance ToScientificLeaf Integer where toScientificLeaf = fromIntegral
+
+instance ToScientificLeaf Double where toScientificLeaf = Sci.fromFloatDigits
+
+-- | Reflect a type-level list of 'Symbol's to their 'Text's (for @EnumOneOf@).
+class KnownSymbols (ss :: [Symbol]) where
+ symbolTexts :: Proxy ss -> [Text]
+
+instance KnownSymbols '[] where symbolTexts _ = []
+
+instance (KnownSymbol s, KnownSymbols ss) => KnownSymbols (s ': ss) where
+ symbolTexts _ = T.pack (symbolVal (Proxy @s)) : symbolTexts (Proxy @ss)
+
+-- | Parse a numeric-bound 'Symbol' (e.g. @"100"@, @"-3.5"@) to a 'Sci.Scientific'.
+parseBound :: forall s. (KnownSymbol s) => Proxy s -> Sci.Scientific
+parseBound _ = case readMaybe (symbolVal (Proxy @s)) of
+ Just v -> v
+ Nothing -> error ("Shikumi.Schema: invalid numeric bound symbol " <> symbolVal (Proxy @s))
+
+tshow :: (Show a) => a -> Text
+tshow = T.pack . show
+
+-- ---------------------------------------------------------------------------
+-- Shared field-metadata traversal (also used by Shikumi.Signature)
+-- ---------------------------------------------------------------------------
+
+-- | Collect 'FieldMeta' (name + optional description) for every record field of a
+-- type, recovering descriptions from 'Field' wrappers.
+fieldMetasOf :: forall a. (GFieldMetas (Rep a)) => [FieldMeta]
+fieldMetasOf = gFieldMetas (Proxy @(Rep a))
+
+-- | Generic field-metadata walk.
+class GFieldMetas (f :: Type -> Type) where
+ gFieldMetas :: Proxy f -> [FieldMeta]
+
+instance (GFieldMetas cs) => GFieldMetas (D1 d cs) where
+ gFieldMetas _ = gFieldMetas (Proxy @cs)
+
+instance (GFieldMetas fs) => GFieldMetas (C1 c fs) where
+ gFieldMetas _ = gFieldMetas (Proxy @fs)
+
+instance (GFieldMetas a, GFieldMetas b) => GFieldMetas (a :*: b) where
+ gFieldMetas _ = gFieldMetas (Proxy @a) ++ gFieldMetas (Proxy @b)
+
+instance GFieldMetas U1 where
+ gFieldMetas _ = []
+
+instance (Selector s, FieldDoc t) => GFieldMetas (S1 s (K1 R t)) where
+ gFieldMetas _ =
+ [ FieldMeta
+ { fieldName = T.pack (selName (undefined :: S1 s (K1 R t) ())),
+ fieldDesc = fieldDoc (Proxy @t)
+ }
+ ]
+
+-- | The description of a field's type, if it is a 'Field' wrapper.
+class FieldDoc t where
+ fieldDoc :: Proxy t -> Maybe Text
+
+instance {-# OVERLAPPABLE #-} FieldDoc a where
+ fieldDoc _ = Nothing
+
+instance (KnownSymbol d) => FieldDoc (Field d a) where
+ fieldDoc _ = Just (T.pack (symbolVal (Proxy @d)))
+
+-- ---------------------------------------------------------------------------
+-- Internal helpers
+-- ---------------------------------------------------------------------------
+
+mismatch :: FieldPath -> Text -> Value -> Either ShikumiError a
+mismatch path expected got =
+ Left (SchemaMismatch (renderPath path <> ": expected " <> expected <> ", got " <> jsonTypeOf got))
+
+jsonTypeOf :: Value -> Text
+jsonTypeOf = \case
+ String {} -> "string"
+ Number {} -> "number"
+ Bool {} -> "boolean"
+ Null -> "null"
+ Array {} -> "array"
+ Object {} -> "object"
diff --git a/src/Shikumi/Schema/Types.hs b/src/Shikumi/Schema/Types.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Schema/Types.hs
@@ -0,0 +1,192 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE KindSignatures #-}
+
+-- | Shared vocabulary for the schema engine: the per-field description wrapper
+-- 'Field', the derived 'FieldMeta' record, JSON-Schema smart constructors (the
+-- small subset OpenAI/Anthropic structured-output endpoints accept), and the
+-- field-path breadcrumb used to locate decode errors.
+--
+-- This module is owned by EP-3 (integration point #3). The user-facing error
+-- type is 'Shikumi.Error.ShikumiError' (owned by EP-1); decode errors render a
+-- 'FieldPath' breadcrumb into that type's @Text@ payloads.
+module Shikumi.Schema.Types
+ ( -- * Field description wrapper
+ Field (..),
+ field,
+
+ -- * Declarative field constraints (EP-26)
+ Constraint (..),
+ Constrained (..),
+ constrained,
+
+ -- * Derived field metadata
+ FieldMeta (..),
+
+ -- * Field-path breadcrumbs
+ FieldPath,
+ pushField,
+ pushIndex,
+ renderPath,
+
+ -- * JSON-Schema smart constructors
+ objectSchema,
+ stringSchema,
+ integerSchema,
+ numberSchema,
+ boolSchema,
+ arraySchema,
+ enumSchema,
+ nullableSchema,
+ withDescription,
+ withMinLength,
+ withMaxLength,
+ withMinimum,
+ withMaximum,
+ withEnum,
+ )
+where
+
+import Data.Aeson (Value (..), object, toJSON, (.=))
+import Data.Aeson.Key qualified as Key
+import Data.Aeson.KeyMap qualified as KM
+import Data.Scientific (Scientific)
+import Data.Text (Text)
+import Data.Text qualified as T
+import GHC.Generics (Generic)
+import GHC.TypeLits (Nat, Symbol)
+
+-- | A field carrying a compile-time description as a type-level 'Symbol'. A
+-- single 'GHC.Generics' traversal recovers both the field name (from the record
+-- selector) and its description (from @desc@), so the two cannot drift. The
+-- wrapper is opt-in per field: a bare @Text@ field simply has no description.
+newtype Field (desc :: Symbol) a = Field {unField :: a}
+ deriving stock (Eq, Show)
+
+-- | Smart constructor mirroring 'unField'.
+field :: a -> Field desc a
+field = Field
+
+-- | A type-level vocabulary of field constraints (EP-26). Each constructor names
+-- one JSON-Schema-expressible rule. 'Nat' bounds are non-negative (fine for the
+-- string-length rules); the signed/decimal numeric bounds 'MinVal'\/'MaxVal' carry
+-- the bound as a 'Symbol' (e.g. @MinVal "0"@, @MaxVal "100"@, @MinVal "-3.5"@) so
+-- both the schema emitter and the validator parse it to a 'Scientific'. Used
+-- promoted, as the @cs@ of 'Constrained'.
+data Constraint
+ = -- | string minimum length -> @"minLength"@
+ MinLen Nat
+ | -- | string maximum length -> @"maxLength"@
+ MaxLen Nat
+ | -- | numeric lower bound -> @"minimum"@
+ MinVal Symbol
+ | -- | numeric upper bound -> @"maximum"@
+ MaxVal Symbol
+ | -- | allowed string values -> @"enum"@
+ EnumOneOf [Symbol]
+
+-- | A field value carrying a compile-time list of constraints (EP-26). Parallel to
+-- 'Field'; the two compose: @Field "desc" (Constrained '[MinLen 10] Text)@. The
+-- reflection of @cs@ to schema keywords and a post-decode validator lives in
+-- "Shikumi.Schema" ('Shikumi.Schema.ReflectConstraints').
+newtype Constrained (cs :: [Constraint]) a = Constrained {unConstrained :: a}
+ deriving stock (Eq, Show)
+
+-- | Smart constructor mirroring 'unConstrained'.
+constrained :: a -> Constrained cs a
+constrained = Constrained
+
+-- | Per-field metadata recovered by the generic walk (used by 'Shikumi.Signature'
+-- and the prompt adapters).
+data FieldMeta = FieldMeta
+ { fieldName :: !Text,
+ fieldDesc :: !(Maybe Text)
+ }
+ deriving stock (Eq, Show, Generic)
+
+-- | A breadcrumb trail to a location inside a decoded value, e.g.
+-- @["bullets", "[2]"]@. Rendered into 'Shikumi.Error.ShikumiError' messages.
+type FieldPath = [Text]
+
+-- | Extend a path with a named record field.
+pushField :: Text -> FieldPath -> FieldPath
+pushField name p = p ++ [name]
+
+-- | Extend a path with an array index.
+pushIndex :: Int -> FieldPath -> FieldPath
+pushIndex i p = p ++ ["[" <> T.pack (show i) <> "]"]
+
+-- | Render a path to a single dotted breadcrumb (@"bullets.[2]"@). The empty path
+-- renders as @"(root)"@ so a top-level error is still legible.
+renderPath :: FieldPath -> Text
+renderPath [] = "(root)"
+renderPath p = T.intercalate "." p
+
+-- | An object schema from @(name, schema)@ properties and a required-name list.
+objectSchema :: [(Text, Value)] -> [Text] -> Value
+objectSchema props required =
+ object
+ [ "type" .= ("object" :: Text),
+ "properties" .= object [Key.fromText k .= v | (k, v) <- props],
+ "required" .= required,
+ "additionalProperties" .= False
+ ]
+
+-- | @{"type":"string"}@.
+stringSchema :: Value
+stringSchema = object ["type" .= ("string" :: Text)]
+
+-- | @{"type":"integer"}@.
+integerSchema :: Value
+integerSchema = object ["type" .= ("integer" :: Text)]
+
+-- | @{"type":"number"}@.
+numberSchema :: Value
+numberSchema = object ["type" .= ("number" :: Text)]
+
+-- | @{"type":"boolean"}@.
+boolSchema :: Value
+boolSchema = object ["type" .= ("boolean" :: Text)]
+
+-- | @{"type":"array","items":}@.
+arraySchema :: Value -> Value
+arraySchema items = object ["type" .= ("array" :: Text), "items" .= items]
+
+-- | @{"enum":[...]}@ for an enum-like sum of nullary constructors.
+enumSchema :: [Text] -> Value
+enumSchema names = object ["enum" .= names]
+
+-- | Permit JSON @null@ in addition to the wrapped schema. The @anyOf@ form is
+-- used (more portable across providers than @"type":["T","null"]@).
+nullableSchema :: Value -> Value
+nullableSchema s = object ["anyOf" .= [s, object ["type" .= ("null" :: Text)]]]
+
+-- | Attach a @"description"@ key to a schema object (no-op on a non-object).
+withDescription :: Text -> Value -> Value
+withDescription desc (Object o) = Object (KM.insert "description" (String desc) o)
+withDescription _ v = v
+
+-- | Attach a @"minLength"@ keyword to a string schema object (no-op otherwise).
+withMinLength :: Int -> Value -> Value
+withMinLength n (Object o) = Object (KM.insert "minLength" (toJSON n) o)
+withMinLength _ v = v
+
+-- | Attach a @"maxLength"@ keyword to a string schema object (no-op otherwise).
+withMaxLength :: Int -> Value -> Value
+withMaxLength n (Object o) = Object (KM.insert "maxLength" (toJSON n) o)
+withMaxLength _ v = v
+
+-- | Attach a @"minimum"@ keyword to a numeric schema object (no-op otherwise).
+withMinimum :: Scientific -> Value -> Value
+withMinimum n (Object o) = Object (KM.insert "minimum" (Number n) o)
+withMinimum _ v = v
+
+-- | Attach a @"maximum"@ keyword to a numeric schema object (no-op otherwise).
+withMaximum :: Scientific -> Value -> Value
+withMaximum n (Object o) = Object (KM.insert "maximum" (Number n) o)
+withMaximum _ v = v
+
+-- | Attach an @"enum"@ keyword (the allowed string values) to a schema object
+-- (no-op otherwise).
+withEnum :: [Text] -> Value -> Value
+withEnum names (Object o) = Object (KM.insert "enum" (toJSON names) o)
+withEnum _ v = v
diff --git a/src/Shikumi/Signature.hs b/src/Shikumi/Signature.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Signature.hs
@@ -0,0 +1,79 @@
+-- | A 'Signature' bundles a task's natural-language instruction, its worked
+-- demonstrations, and the input/output field metadata derived from the @i@/@o@
+-- record types.
+--
+-- The @instruction@ and @demos@ are the framework's /optimizable parameters/
+-- (integration point #3): they are first-class, replaceable values that the
+-- compiler (EP-9) and optimizer (EP-10) rewrite to improve quality. The
+-- @inputFields@/@outputFields@ are /derived metadata/ (not optimizable). Because
+-- 'Signature' derives 'Generic', the @#instruction@ / @#demos@ @generic-lens@
+-- optics are available to downstream plans for free.
+module Shikumi.Signature
+ ( Demo (..),
+ Signature (..),
+ mkSignature,
+ getInstruction,
+ setInstruction,
+ getDemos,
+ setDemos,
+ )
+where
+
+import Data.Text (Text)
+import GHC.Generics (Generic, Rep)
+import Shikumi.Schema (GFieldMetas, fieldMetasOf)
+import Shikumi.Schema.Types (FieldMeta)
+
+-- | A worked input -> output example shown to the model. Optimizers select and
+-- rewrite these.
+data Demo i o = Demo
+ { input :: i,
+ output :: o
+ }
+ deriving stock (Generic, Show, Eq)
+
+-- | A typed task signature. @i@ and @o@ are the input/output record types; they
+-- are kept in the type so a program can hold a @Signature Article Summary@ and
+-- the compiler enforces the input/output types.
+data Signature i o = Signature
+ { -- | optimizable: the task description
+ instruction :: Text,
+ -- | optimizable: worked examples
+ demos :: [Demo i o],
+ -- | derived metadata: the input record's fields
+ inputFields :: [FieldMeta],
+ -- | derived metadata: the output record's fields
+ outputFields :: [FieldMeta]
+ }
+ deriving stock (Generic, Show)
+
+-- | Build a signature from an instruction; demos start empty and the field
+-- metadata is derived from @i@ and @o@ by the shared 'fieldMetasOf' traversal.
+mkSignature ::
+ forall i o.
+ (GFieldMetas (Rep i), GFieldMetas (Rep o)) =>
+ Text ->
+ Signature i o
+mkSignature instr =
+ Signature
+ { instruction = instr,
+ demos = [],
+ inputFields = fieldMetasOf @i,
+ outputFields = fieldMetasOf @o
+ }
+
+-- | Read the (optimizable) instruction.
+getInstruction :: Signature i o -> Text
+getInstruction = instruction
+
+-- | Replace the (optimizable) instruction.
+setInstruction :: Text -> Signature i o -> Signature i o
+setInstruction t sig = sig {instruction = t}
+
+-- | Read the (optimizable) demonstrations.
+getDemos :: Signature i o -> [Demo i o]
+getDemos = demos
+
+-- | Replace the (optimizable) demonstrations.
+setDemos :: [Demo i o] -> Signature i o -> Signature i o
+setDemos ds sig = sig {demos = ds}
diff --git a/src/Shikumi/Stream.hs b/src/Shikumi/Stream.hs
new file mode 100644
--- /dev/null
+++ b/src/Shikumi/Stream.hs
@@ -0,0 +1,288 @@
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
+-- | Program-level streaming (EP-25). A new, additive entry point 'streamProgram'
+-- runs a 'Program' and delivers a sequence of 'StreamEvent's to a caller-supplied
+-- callback as execution proceeds, while still returning the fully-decoded typed
+-- output — identical to what 'Shikumi.Program.runProgram' would return for the same
+-- input. 'runProgram'\'s blocking contract (MasterPlan integration point #4) is
+-- untouched; this module adds surface, it changes nothing.
+--
+-- Two kinds of event reach the callback:
+--
+-- * a __field chunk__ — a piece of an output field's text as the model writes it.
+-- Honestly scoped to a single 'Shikumi.Program.Predict' node (and chains of
+-- them) on the prompt-fallback / raw-text path — the path exercised in the
+-- hermetic stubs, since 'Shikumi.Program.runProgram'\'s placeholder model maps to
+-- the fallback adapter. Native whole-JSON field chunking is /not/ promised (a
+-- JSON blob is only parseable once whole).
+-- * a __status message__ — a human-readable signal that a phase started or ended
+-- ("LM call started/finished", "tool started/finished", "node started/finished").
+-- Status brackets every node, including composites; aggregating combinators
+-- ('Shikumi.Program.Map', 'Shikumi.Program.Parallel',
+-- 'Shikumi.Program.MajorityVote', 'Shikumi.Program.Ensemble') and opaque
+-- 'Shikumi.Program.Embed' bodies stream status only, not field chunks.
+--
+-- The chunk text is the raw provider delta (marker boilerplate included on the
+-- fallback path); trimming markers out of the stream is left to a future refinement.
+module Shikumi.Stream
+ ( -- * Event types
+ FieldChunk (..),
+ StatusPhase (..),
+ Status (..),
+ StreamEvent (..),
+
+ -- * Entry points
+ streamComplete,
+ streamProgram,
+ )
+where
+
+import Baikai
+ ( AssistantContent (..),
+ AssistantMessageEvent (..),
+ BlockEndPayload (..),
+ Context,
+ DeltaPayload (..),
+ Message (..),
+ Model,
+ Options,
+ Response,
+ TerminalPayload (..),
+ _Model,
+ _Response,
+ _TextContent,
+ )
+import Control.Lens ((&), (.~))
+import Data.Generics.Labels ()
+import Data.Maybe (fromMaybe, listToMaybe)
+import Data.Text (Text)
+import Data.Vector qualified as V
+import Effectful (Eff, (:>))
+import Effectful.Error.Static (Error, throwError)
+import GHC.Generics (Generic)
+import Shikumi.Adapter (Adapter (..), ToPrompt, adapterFor)
+import Shikumi.Error (ShikumiError)
+import Shikumi.LLM (LLM, stream)
+import Shikumi.Program
+ ( Demo (..),
+ Params (..),
+ Program (Compose, Embed, FMap, MajorityVote, Map, Parallel, Predict, Retry, RetryWhen, Validate),
+ acceptOrReject,
+ retryWith,
+ runProgram,
+ )
+import Shikumi.Schema (FromModel, ToSchema, fromModel)
+import Shikumi.Schema.Types qualified as ST
+import Shikumi.Signature (Signature, getInstruction, outputFields, setDemos, setInstruction)
+import Shikumi.Signature qualified as Sig
+
+-- ---------------------------------------------------------------------------
+-- Event types
+-- ---------------------------------------------------------------------------
+
+-- | A piece of one output field's text, as the model writes it. Mirrors DSPy's
+-- @StreamResponse { predict_name, signature_field_name, chunk, is_last_chunk }@.
+data FieldChunk = FieldChunk
+ { -- | which output field this chunk belongs to (@""@ if not attributed)
+ fieldName :: !Text,
+ -- | the text delta
+ chunk :: !Text,
+ -- | 'True' on the final chunk of this field
+ isLast :: !Bool
+ }
+ deriving stock (Eq, Show, Generic)
+
+-- | The phase a status message marks. Borrowed in spirit from the trace SpanKind.
+data StatusPhase
+ = -- | an LM call is about to begin
+ LmStart
+ | -- | an LM call finished
+ LmEnd
+ | -- | a tool call is about to begin
+ ToolStart
+ | -- | a tool call finished
+ ToolEnd
+ | -- | a (sub-)program node started
+ NodeStart
+ | -- | a (sub-)program node finished
+ NodeEnd
+ deriving stock (Eq, Show, Generic)
+
+-- | A human-readable status signal. Mirrors DSPy's @StatusMessage@, plus the phase.
+data Status = Status
+ { phase :: !StatusPhase,
+ message :: !Text
+ }
+ deriving stock (Eq, Show, Generic)
+
+-- | One streaming event handed to the caller's callback, in order.
+data StreamEvent
+ = StreamFieldChunk !FieldChunk
+ | StreamStatus !Status
+ deriving stock (Eq, Show, Generic)
+
+-- ---------------------------------------------------------------------------
+-- M1: driving one streaming LM call
+-- ---------------------------------------------------------------------------
+
+-- | Drive one streaming LM call, folding baikai events into 'StreamEvent's
+-- delivered to the callback in order, and return the fully assembled 'Response'
+-- (reassembled from the terminal event) so the caller can 'parse' it into a typed
+-- value exactly as the blocking path would.
+--
+-- Each @TextDelta@ becomes a 'StreamFieldChunk' attributed to @fieldNm@; the final
+-- delta carries @isLast = True@. If the stream carried no deltas (e.g. a cached
+-- whole response arrived as a single @TextEnd@), one synthetic chunk carrying the
+-- @TextEnd@ content is emitted so the caller still sees the field.
+streamComplete ::
+ (LLM :> es) =>
+ -- | the output field name to attribute chunks to
+ Text ->
+ Model ->
+ Context ->
+ Options ->
+ -- | per-event callback, runs in the caller's @Eff@
+ (StreamEvent -> Eff es ()) ->
+ Eff es Response
+streamComplete fieldNm model ctx opts cb = do
+ evs <- stream model ctx opts
+ emitChunks fieldNm [d | TextDelta (DeltaPayload _ d) <- evs] (firstTextEnd evs) cb
+ pure (reassemble evs)
+
+-- | Emit the field chunks for one block: every delta in order, the last marked
+-- 'isLast'. With no deltas, synthesize a single final chunk from the @TextEnd@
+-- content (if any), so the field is still observed.
+emitChunks ::
+ Text ->
+ [Text] ->
+ Maybe Text ->
+ (StreamEvent -> Eff es ()) ->
+ Eff es ()
+emitChunks fieldNm deltas mEnd cb = case deltas of
+ [] -> maybe (pure ()) (\c -> cb (mkChunk c True)) mEnd
+ _ ->
+ let n = length deltas
+ in mapM_ (\(k, d) -> cb (mkChunk d (k == n))) (zip [1 :: Int ..] deltas)
+ where
+ mkChunk t l = StreamFieldChunk (FieldChunk fieldNm t l)
+
+-- | The content of the first @TextEnd@ block, if any.
+firstTextEnd :: [AssistantMessageEvent] -> Maybe Text
+firstTextEnd evs = listToMaybe [c | TextEnd (BlockEndPayload _ c) <- evs]
+
+-- | Reassemble a 'Response' from the stream's terminal event (its payload's
+-- fully-assembled assistant message). 'EventError' surfaces the assembled
+-- (possibly-partial) response just like the blocking path, so the caller's 'parse'
+-- fails on a bad response identically. Falls back to a response synthesized from
+-- the @TextEnd@ content if no terminal assistant message is present.
+reassemble :: [AssistantMessageEvent] -> Response
+reassemble evs = case terminalPayloads of
+ (p : _) -> _Response & #message .~ p
+ [] -> synthResponse (fromMaybe "" (firstTextEnd evs))
+ where
+ terminalPayloads =
+ [p | EventDone (TerminalPayload _ (AssistantMessage p)) <- evs]
+ ++ [p | EventError (TerminalPayload _ (AssistantMessage p)) <- evs]
+
+-- | A response carrying @t@ as its single assistant text block.
+synthResponse :: Text -> Response
+synthResponse t =
+ _Response & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t))
+
+-- ---------------------------------------------------------------------------
+-- M2/M3: streaming a whole program
+-- ---------------------------------------------------------------------------
+
+-- | Run a program, delivering 'StreamEvent's to the callback as it executes, and
+-- return the typed output equal to 'Shikumi.Program.runProgram'\'s for the same
+-- input. The effect row is exactly 'runProgram'\'s plus the callback, so a streamed
+-- program runs under the same interpreters as a blocking one.
+--
+-- A leaf 'Predict' streams its first output field's chunks bracketed by
+-- @LmStart@/@LmEnd@. Composites are bracketed by @NodeStart@/@NodeEnd@:
+-- 'Compose'/'FMap' and the single-result control nodes
+-- ('Retry'/'RetryWhen'/'Validate') recurse so their leaf predicts still stream;
+-- the aggregating combinators ('Map'/'Parallel'/'MajorityVote'/'Ensemble') and the
+-- opaque 'Embed' body run blocking and emit node-boundary status only (field chunks
+-- come from leaf LM calls on the single-result path).
+streamProgram ::
+ forall i o es.
+ (LLM :> es, Error ShikumiError :> es) =>
+ Program i o ->
+ i ->
+ (StreamEvent -> Eff es ()) ->
+ Eff es o
+streamProgram prog i cb = case prog of
+ Predict sig ps -> streamPredict sig ps i cb
+ Compose f g -> bracketNode cb $ do
+ b <- streamProgram f i cb
+ streamProgram g b cb
+ FMap k p -> bracketNode cb (k <$> streamProgram p i cb)
+ Retry n p -> bracketNode cb (retryWith (\pp ii -> streamProgram pp ii cb) (const True) n p i)
+ RetryWhen ok n p -> bracketNode cb (retryWith (\pp ii -> streamProgram pp ii cb) ok n p i)
+ Validate v p -> bracketNode cb (streamProgram p i cb >>= acceptOrReject v)
+ -- Aggregating / opaque nodes: bracket with node status, run blocking (no field
+ -- chunks where per-field streaming is ambiguous or the body is opaque).
+ Map _ _ -> blockingNode prog i cb
+ Parallel _ _ -> blockingNode prog i cb
+ MajorityVote _ _ _ -> blockingNode prog i cb
+ Embed _ -> blockingNode prog i cb
+ _ -> blockingNode prog i cb
+
+-- | Stream a single 'Predict': overlay its 'Params', render via the (fallback)
+-- adapter, stream the first output field's chunks bracketed by @LmStart@/@LmEnd@,
+-- then run the identical 'parse' on the reassembled response so the typed result
+-- equals 'runProgram'\'s.
+streamPredict ::
+ forall i o es.
+ (FromModel i, FromModel o, ToSchema o, ToPrompt i, ToPrompt o) =>
+ (LLM :> es, Error ShikumiError :> es) =>
+ Signature i o ->
+ Params ->
+ i ->
+ (StreamEvent -> Eff es ()) ->
+ Eff es o
+streamPredict sig ps i cb = do
+ sig' <- effectiveSig sig ps
+ let adapter = adapterFor _Model
+ (ctx, opts) = render adapter sig' i
+ fieldNm = case map ST.fieldName (outputFields sig') of
+ (x : _) -> x
+ [] -> ""
+ cb (StreamStatus (Status LmStart "LM call started"))
+ resp <- streamComplete fieldNm _Model ctx opts cb
+ cb (StreamStatus (Status LmEnd "LM call finished"))
+ either throwError pure (parse adapter sig' resp)
+
+-- | Bracket an action with @NodeStart@/@NodeEnd@ status messages.
+bracketNode :: (StreamEvent -> Eff es ()) -> Eff es a -> Eff es a
+bracketNode cb act = do
+ cb (StreamStatus (Status NodeStart "node started"))
+ r <- act
+ cb (StreamStatus (Status NodeEnd "node finished"))
+ pure r
+
+-- | Run a node blocking (no field chunks) bracketed by node-boundary status.
+blockingNode ::
+ (LLM :> es, Error ShikumiError :> es) =>
+ Program i o ->
+ i ->
+ (StreamEvent -> Eff es ()) ->
+ Eff es o
+blockingNode prog i cb = bracketNode cb (runProgram prog i)
+
+-- | Overlay a node's 'Params' onto its signature (effective instruction + decoded
+-- demos), replicating 'Shikumi.Program.runProgram'\'s internal @effectiveSignature@
+-- so a streamed 'Predict' renders byte-for-byte as the blocking one would.
+effectiveSig ::
+ (FromModel i, FromModel o, Error ShikumiError :> es) =>
+ Signature i o ->
+ Params ->
+ Eff es (Signature i o)
+effectiveSig sig (Params {instructionOverride = ovr, demos = ds}) = do
+ typed <- either throwError pure (traverse decodeDemo ds)
+ pure (setDemos typed (setInstruction (fromMaybe (getInstruction sig) ovr) sig))
+ where
+ decodeDemo (Demo inJ outJ) = Sig.Demo <$> fromModel inJ <*> fromModel outJ
diff --git a/test/AdapterSpec.hs b/test/AdapterSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/AdapterSpec.hs
@@ -0,0 +1,115 @@
+-- | M5: 'Shikumi.Adapter'. Capability selection is a pure function of the model;
+-- the fallback adapter renders @[[ ## field ## ]]@ sections and re-parses them
+-- into a typed output, with a missing marker producing 'MissingField'. The native
+-- path's schema attachment is pending EP-2.
+module AdapterSpec (tests) where
+
+import Baikai
+ ( Api (..),
+ AssistantContent (..),
+ Model,
+ Response,
+ _Model,
+ _Response,
+ _TextContent,
+ )
+import Control.Lens ((&), (.~), (^.))
+import Data.Generics.Labels ()
+import Data.Maybe (fromMaybe)
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Vector qualified as V
+import Fixtures (Article, Author (..), Sentiment (..), Summary (..), sampleArticle, sampleSummary)
+import Shikumi.Adapter
+ ( Adapter (..),
+ ModelCapability (..),
+ capabilityFor,
+ fallbackAdapter,
+ nativeAdapter,
+ )
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.Schema.Types (field)
+import Shikumi.Signature (Demo (..), Signature, mkSignature, setDemos)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (testCase, (@?=))
+
+anthropicModel :: Model
+anthropicModel = _Model & #provider .~ "anthropic" & #api .~ AnthropicMessages
+
+ollamaModel :: Model
+ollamaModel = _Model & #provider .~ "ollama" & #api .~ Custom "ollama"
+
+sig :: Signature Article Summary
+sig = setDemos [Demo sampleArticle sampleSummary] (mkSignature "Summarize the article")
+
+-- | Build a response whose single assistant text block carries the given body.
+mkResponse :: Text -> Response
+mkResponse t =
+ _Response & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t))
+
+markerBody :: Text
+markerBody =
+ T.intercalate
+ "\n"
+ [ "[[ ## headline ## ]]",
+ "Shikumi types LM programs",
+ "[[ ## bullets ## ]]",
+ "[\"records in\", \"records out\", \"errors are typed\"]",
+ "[[ ## author ## ]]",
+ "{\"name\": \"Ada\"}",
+ "[[ ## sentiment ## ]]",
+ "Positive",
+ "[[ ## note ## ]]",
+ "null",
+ "[[ ## completed ## ]]"
+ ]
+
+markerBodyNoBullets :: Text
+markerBodyNoBullets =
+ T.intercalate
+ "\n"
+ [ "[[ ## headline ## ]]",
+ "Shikumi types LM programs",
+ "[[ ## author ## ]]",
+ "{\"name\": \"Ada\"}",
+ "[[ ## sentiment ## ]]",
+ "Positive",
+ "[[ ## note ## ]]",
+ "null",
+ "[[ ## completed ## ]]"
+ ]
+
+expectedSummary :: Summary
+expectedSummary =
+ Summary
+ { headline = field "Shikumi types LM programs",
+ bullets = field ["records in", "records out", "errors are typed"],
+ author = Author {name = field "Ada"},
+ sentiment = Positive,
+ note = Nothing
+ }
+
+sysOf :: Adapter Article Summary -> Text
+sysOf adapter = fromMaybe "" (fst (render adapter sig sampleArticle) ^. #systemPrompt)
+
+tests :: TestTree
+tests =
+ testGroup
+ "AdapterSpec"
+ [ testCase "capabilityFor: Anthropic messages -> NativeSchema" $
+ capabilityFor anthropicModel @?= NativeSchema,
+ testCase "capabilityFor: Custom (ollama) host -> PromptFallback" $
+ capabilityFor ollamaModel @?= PromptFallback,
+ testCase "fallback render: system prompt has the instruction and field markers" $ do
+ T.isInfixOf "Summarize the article" (sysOf fallbackAdapter) @?= True
+ T.isInfixOf "[[ ## headline ## ]]" (sysOf fallbackAdapter) @?= True,
+ testCase "fallback render: a demo becomes a user/assistant message pair" $
+ -- one demo (2 messages) + the actual input (1 message)
+ V.length (fst (render fallbackAdapter sig sampleArticle) ^. #messages) @?= 3,
+ testCase "fallback parse: marker body decodes to the expected Summary" $
+ parse fallbackAdapter sig (mkResponse markerBody) @?= Right expectedSummary,
+ testCase "fallback parse: a missing marker -> MissingField (located)" $
+ parse fallbackAdapter sig (mkResponse markerBodyNoBullets) @?= Left (MissingField "bullets"),
+ testCase "native render: system prompt carries the instruction (schema attach pending EP-2)" $
+ T.isInfixOf "Summarize the article" (sysOf nativeAdapter) @?= True
+ ]
diff --git a/test/CombinatorSpec.hs b/test/CombinatorSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/CombinatorSpec.hs
@@ -0,0 +1,330 @@
+-- | EP-5 acceptance: every combinator demonstrated against a deterministic mock
+-- LM, asserting observable aggregate behaviour. Order-sensitive cases run under
+-- the sequential 'runProgram'; 'runProgramConc' is exercised with
+-- order-independent (identical) replies so the concurrent pop order cannot affect
+-- the assertion. The final group proves integration point #4's cross-cutting
+-- properties (the parameter traversal reaches every nested leaf; the structural
+-- shape and parameter vector round-trip).
+module CombinatorSpec (tests) where
+
+import Data.IORef (newIORef, readIORef)
+import Data.Text (Text)
+import Data.Text qualified as T
+import Effectful (runEff)
+import Effectful.Concurrent (runConcurrent)
+import Effectful.Error.Static (runErrorNoCallStack)
+import ProgramFixtures
+ ( Cell (..),
+ Draft (..),
+ Topic (..),
+ cellSig,
+ draftResponse,
+ markerBody,
+ mkResponse,
+ outlineResponse,
+ outlineToDraft,
+ runRecordingLLM,
+ runScriptedLLM,
+ topicToOutline,
+ )
+import Shikumi.Combinator
+ ( TempSchedule (..),
+ chain,
+ ensemble,
+ majorityVote,
+ majorityVoteBy,
+ mapP,
+ mapSeqP,
+ parallel2,
+ parallelN,
+ retry,
+ retryWhen,
+ validate,
+ validateRetry,
+ (>>>),
+ )
+import Shikumi.Error (ShikumiError (..), isTransient)
+import Shikumi.LLM (Response)
+import Shikumi.LLM.Mock (MockReply (..), runMockLLM, runMockLLMCounting)
+import Shikumi.Module (predict)
+import Shikumi.Program
+ ( Params (..),
+ Program,
+ ProgramShape (..),
+ ProgramShapeError (..),
+ emptyParams,
+ foldParams,
+ mapParams,
+ programParams,
+ programShape,
+ runProgram,
+ runProgramConc,
+ setProgramParams,
+ )
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (assertBool, testCase, (@?=))
+
+-- ---------------------------------------------------------------------------
+-- Run harnesses
+-- ---------------------------------------------------------------------------
+
+-- | Run a program sequentially against a list of canned success 'Response's.
+runSeq :: Program i o -> [Response] -> i -> IO (Either ShikumiError o)
+runSeq prog rs i = do
+ ref <- newIORef rs
+ runEff . runErrorNoCallStack @ShikumiError . runScriptedLLM ref $ runProgram prog i
+
+-- | Run a program sequentially against scripted replies (successes /and/
+-- failures).
+runMockSeq :: Program i o -> [MockReply] -> i -> IO (Either ShikumiError o)
+runMockSeq prog rs i =
+ runEff . runErrorNoCallStack @ShikumiError . runMockLLM rs $ runProgram prog i
+
+-- | Run a program and report both the result and the number of completions
+-- issued. The counter is read after the run resolves, so it survives a failure
+-- short-circuit.
+runCounted :: Program i o -> [MockReply] -> i -> IO (Either ShikumiError o, Int)
+runCounted prog rs i = do
+ cnt <- newIORef 0
+ res <- runEff . runErrorNoCallStack @ShikumiError . runMockLLMCounting cnt rs $ runProgram prog i
+ n <- readIORef cnt
+ pure (res, n)
+
+-- | Run a program through the concurrent executor. Replies are popped from a
+-- shared queue, so callers pass identical replies when the assertion is
+-- order-sensitive.
+runConc ::
+ Program i o -> [Response] -> i -> IO (Either ShikumiError o)
+runConc prog rs i = do
+ ref <- newIORef rs
+ runEff . runErrorNoCallStack @ShikumiError . runConcurrent . runScriptedLLM ref $
+ runProgramConc prog i
+
+-- | Run sequentially while recording each request's rendered system prompt.
+runRec :: Program i o -> [Response] -> i -> IO [Text]
+runRec prog rs i = do
+ cap <- newIORef []
+ ref <- newIORef rs
+ _ <- runEff . runErrorNoCallStack @ShikumiError . runRecordingLLM cap ref $ runProgram prog i
+ readIORef cap
+
+-- ---------------------------------------------------------------------------
+-- Fixtures
+-- ---------------------------------------------------------------------------
+
+-- | A @Cell -> Cell@ predictor — the simplest typed leaf to compose.
+cellP :: Program Cell Cell
+cellP = predict cellSig
+
+-- | A canned response decoding to @Cell {cell = v}@.
+cellResp :: Text -> Response
+cellResp v = mkResponse (markerBody [("cell", v)])
+
+-- | A transient (retryable) failure.
+transientFail :: MockReply
+transientFail = MockFail (ProviderFailure "stub transient")
+
+-- | Concatenate the cells (a reducer for outputs that are not usefully modal).
+concatCells :: [Cell] -> Cell
+concatCells cs = Cell (T.concat (map cell cs))
+
+-- | A fixed single-temperature schedule (the schedule is inert on the wire; see
+-- 'Shikumi.Program.TempSchedule').
+sched1 :: TempSchedule
+sched1 = TempFixed [0.0]
+
+-- ---------------------------------------------------------------------------
+-- Tests
+-- ---------------------------------------------------------------------------
+
+tests :: TestTree
+tests =
+ testGroup
+ "CombinatorSpec"
+ [ smokeTests,
+ pipelineTests,
+ mapTests,
+ parallelTests,
+ retryTests,
+ validateTests,
+ majorityVoteTests,
+ ensembleTests,
+ crossCuttingTests
+ ]
+
+smokeTests :: TestTree
+smokeTests =
+ testGroup
+ "smoke"
+ [ testCase "predict through the mock LLM returns the decoded value" $ do
+ r <- runSeq cellP [cellResp "echoed"] (Cell "in")
+ r @?= Right (Cell "echoed")
+ ]
+
+pipelineTests :: TestTree
+pipelineTests =
+ testGroup
+ "Pipeline"
+ [ testCase "(>>>) threads an intermediate type absent from the composite signature" $ do
+ let pipe = predict topicToOutline >>> predict outlineToDraft :: Program Topic Draft
+ r <- runSeq pipe [outlineResponse, draftResponse] (Topic "haskell")
+ r @?= Right (Draft "A fine essay about the outline."),
+ testCase "chain composes n same-type stages, returning the last stage's output" $ do
+ let pipe = chain [cellP, cellP, cellP]
+ r <- runSeq pipe [cellResp "a", cellResp "b", cellResp "c"] (Cell "in")
+ r @?= Right (Cell "c")
+ ]
+
+mapTests :: TestTree
+mapTests =
+ testGroup
+ "Map"
+ [ testCase "mapSeqP runs once per element, preserving input order" $ do
+ let prog = mapSeqP cellP
+ r <- runSeq prog [cellResp "1", cellResp "2", cellResp "3", cellResp "4", cellResp "5"] (replicate 5 (Cell "in"))
+ r @?= Right [Cell "1", Cell "2", Cell "3", Cell "4", Cell "5"],
+ testCase "mapP runs under the concurrent executor and yields one output per input" $ do
+ let prog = mapP 2 cellP
+ r <- runConc prog (replicate 3 (cellResp "x")) (replicate 3 (Cell "in"))
+ r @?= Right [Cell "x", Cell "x", Cell "x"]
+ ]
+
+parallelTests :: TestTree
+parallelTests =
+ testGroup
+ "Parallel"
+ [ testCase "parallel2 runs both programs on the same input and pairs the outputs" $ do
+ let prog = parallel2 cellP cellP
+ r <- runSeq prog [cellResp "L", cellResp "R"] (Cell "in")
+ r @?= Right (Cell "L", Cell "R"),
+ testCase "parallelN collects every program's output in order" $ do
+ let prog = parallelN [cellP, cellP, cellP]
+ r <- runSeq prog [cellResp "1", cellResp "2", cellResp "3"] (Cell "in")
+ r @?= Right [Cell "1", Cell "2", Cell "3"],
+ testCase "parallel2 runs concurrently under runProgramConc" $ do
+ let prog = parallel2 cellP cellP
+ r <- runConc prog [cellResp "z", cellResp "z"] (Cell "in")
+ r @?= Right (Cell "z", Cell "z")
+ ]
+
+retryTests :: TestTree
+retryTests =
+ testGroup
+ "Retry"
+ [ testCase "fail-then-succeed returns the success after exactly two attempts" $ do
+ (r, n) <- runCounted (retry 3 cellP) [transientFail, MockOk (cellResp "ok")] (Cell "in")
+ r @?= Right (Cell "ok")
+ n @?= 2,
+ testCase "all-fail surfaces the last error after exactly n attempts" $ do
+ (r, n) <- runCounted (retry 3 cellP) [transientFail, transientFail, transientFail] (Cell "in")
+ r @?= Left (ProviderFailure "stub transient")
+ n @?= 3,
+ testCase "retryWhen skips a non-matching error after a single attempt" $ do
+ (r, n) <-
+ runCounted
+ (retryWhen isTransient 3 cellP)
+ [MockFail (SchemaMismatch "nope"), MockOk (cellResp "ok")]
+ (Cell "in")
+ r @?= Left (SchemaMismatch "nope")
+ n @?= 1
+ ]
+
+validateTests :: TestTree
+validateTests =
+ testGroup
+ "Validate"
+ [ testCase "a rejected output surfaces ValidationFailure with the reason" $ do
+ let prog = validate (\c -> cell c == "good") "cell must be good" cellP
+ r <- runMockSeq prog [MockOk (cellResp "bad")] (Cell "in")
+ r @?= Left (ValidationFailure "cell must be good"),
+ testCase "an accepted output passes through unchanged" $ do
+ let prog = validate (\c -> cell c == "good") "cell must be good" cellP
+ r <- runMockSeq prog [MockOk (cellResp "good")] (Cell "in")
+ r @?= Right (Cell "good"),
+ testCase "validateRetry recovers when a later attempt passes" $ do
+ let prog = validateRetry 2 (\c -> cell c == "good") "cell must be good" cellP
+ (r, n) <- runCounted prog [MockOk (cellResp "bad"), MockOk (cellResp "good")] (Cell "in")
+ r @?= Right (Cell "good")
+ n @?= 2
+ ]
+
+majorityVoteTests :: TestTree
+majorityVoteTests =
+ testGroup
+ "MajorityVote"
+ [ testCase "five samples A,B,A,A,B return the modal A" $ do
+ let prog = majorityVote 5 sched1 cellP
+ r <-
+ runSeq
+ prog
+ [cellResp "A", cellResp "B", cellResp "A", cellResp "A", cellResp "B"]
+ (Cell "in")
+ r @?= Right (Cell "A"),
+ testCase "a tie resolves to the first-seen value" $ do
+ let prog = majorityVote 2 sched1 cellP
+ r <- runSeq prog [cellResp "A", cellResp "B"] (Cell "in")
+ r @?= Right (Cell "A"),
+ testCase "majorityVoteBy applies the custom reducer" $ do
+ let prog = majorityVoteBy 3 sched1 concatCells cellP
+ r <- runSeq prog [cellResp "x", cellResp "y", cellResp "z"] (Cell "in")
+ r @?= Right (Cell "xyz")
+ ]
+
+ensembleTests :: TestTree
+ensembleTests =
+ testGroup
+ "Ensemble"
+ [ testCase "ensemble folds every member's output with the reducer" $ do
+ let prog = ensemble [cellP, cellP, cellP] concatCells
+ r <- runSeq prog [cellResp "1", cellResp "2", cellResp "3"] (Cell "in")
+ r @?= Right (Cell "123"),
+ testCase "the parameter traversal reaches a leaf inside every member" $
+ length (foldParams (ensemble [cellP, cellP, cellP] concatCells)) @?= 3
+ ]
+
+-- A deliberately deep program: a Pipeline of (a Retry of a MajorityVote) then a
+-- Validate — every combinator wrapper nests a single Predict leaf, so the
+-- traversal must reach two leaves total.
+deepProg :: Program Cell Cell
+deepProg =
+ chain
+ [ retry 2 (majorityVote 3 sched1 cellP),
+ validate (const True) "always ok" cellP
+ ]
+
+crossCuttingTests :: TestTree
+crossCuttingTests =
+ testGroup
+ "cross-cutting (traversal + serialization)"
+ [ testCase "the traversal reaches every nested leaf" $
+ length (foldParams deepProg) @?= 2,
+ testCase "an instruction written through the traversal reaches the deepest leaf at run time" $ do
+ let sentinel = "SENTINEL-INSTRUCTION"
+ rewritten = mapParams (\p -> p {instructionOverride = Just sentinel}) deepProg
+ -- stage 1 samples the MajorityVote leaf 3×, stage 2 the Validate leaf 1×.
+ prompts <-
+ runRec
+ rewritten
+ (replicate 4 (cellResp "v"))
+ (Cell "in")
+ assertBool
+ ("expected the sentinel in every captured prompt; got: " <> show prompts)
+ (length prompts == 4 && all (T.isInfixOf sentinel) prompts),
+ testCase "the structural shape is parameter-independent and round-trips the parameter vector" $ do
+ let p1 = emptyParams {instructionOverride = Just "one"}
+ p2 = emptyParams {instructionOverride = Just "two"}
+ programParams deepProg @?= [emptyParams, emptyParams]
+ case setProgramParams [p1, p2] deepProg of
+ Left e -> assertBool ("unexpected setProgramParams failure: " <> show e) False
+ Right prog' -> do
+ programParams prog' @?= [p1, p2]
+ programShape prog' @?= programShape deepProg,
+ testCase "a parameter-vector length mismatch is a typed error, not a crash" $
+ case setProgramParams [emptyParams] deepProg of
+ Left (ParamCountMismatch (expected, got)) -> (expected, got) @?= (2, 1)
+ Right _ -> assertBool "expected a ParamCountMismatch" False,
+ testCase "a function-carrying node serializes its structure (closure omitted) without crashing" $
+ case programShape (ensemble [cellP] concatCells) of
+ ShapeEnsemble [ShapePredict _] -> pure ()
+ other -> assertBool ("expected ShapeEnsemble [ShapePredict _]; got: " <> show other) False
+ ]
diff --git a/test/ConstraintSpec.hs b/test/ConstraintSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/ConstraintSpec.hs
@@ -0,0 +1,90 @@
+{-# LANGUAGE DataKinds #-}
+
+-- | EP-26 M2: declarative field constraints. A 'Constrained' field carries a
+-- type-level constraint list that (a) emits the matching JSON-Schema keywords in
+-- the derived schema and (b) is enforced after decode, turning a violation into a
+-- located 'ValidationFailure'. The contrast record 'BioUnchecked' (same fields, no
+-- constraint) shows the constraint is what rejects an out-of-bounds value.
+module ConstraintSpec (tests) where
+
+import Data.Aeson (Value (..))
+import Data.Aeson.Key qualified as Key
+import Data.Aeson.KeyMap qualified as KM
+import Data.Text (Text)
+import GHC.Generics (Generic)
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.Schema (FromModel, ToSchema, Validatable, deriveSchema, fromModel)
+import Shikumi.Schema.Types (Constrained, Constraint (..))
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (assertFailure, testCase, (@?=))
+
+-- A small constrained record: a tagline of at least 10 characters and a score in
+-- the inclusive range [0, 100].
+data Bio = Bio
+ { tagline :: Constrained '[ 'MinLen 10] Text,
+ score :: Constrained '[ 'MinVal "0", 'MaxVal "100"] Int
+ }
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema Bio
+
+instance FromModel Bio
+
+instance Validatable Bio
+
+-- The same shape with no constraints, to prove the constraint does the rejecting.
+data BioUnchecked = BioUnchecked
+ { tagline :: Text,
+ score :: Int
+ }
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema BioUnchecked
+
+instance FromModel BioUnchecked
+
+instance Validatable BioUnchecked
+
+-- | Drill into @properties.@ of a record schema 'Value'.
+propOf :: Text -> Value -> Maybe Value
+propOf name (Object o) = do
+ Object props <- KM.lookup "properties" o
+ KM.lookup (Key.fromText name) props
+propOf _ _ = Nothing
+
+-- | Look a keyword up inside a property schema object.
+keyword :: Text -> Value -> Maybe Value
+keyword k (Object o) = KM.lookup (Key.fromText k) o
+keyword _ _ = Nothing
+
+tests :: TestTree
+tests =
+ testGroup
+ "ConstraintSpec"
+ [ testCase "schema: tagline carries minLength 10" $
+ (propOf "tagline" (deriveSchema @Bio) >>= keyword "minLength") @?= Just (Number 10),
+ testCase "schema: score carries minimum 0 and maximum 100" $ do
+ (propOf "score" (deriveSchema @Bio) >>= keyword "minimum") @?= Just (Number 0)
+ (propOf "score" (deriveSchema @Bio) >>= keyword "maximum") @?= Just (Number 100),
+ testCase "decode: short tagline -> ValidationFailure" $
+ fromModel @Bio (obj "short" 50)
+ @?= Left (ValidationFailure "tagline: minLength 10 violated"),
+ testCase "decode: conforming value -> Right" $
+ case fromModel @Bio (obj "long enough tagline" 50) of
+ Right _ -> pure ()
+ Left e -> assertFailure ("expected Right, got " <> show e),
+ testCase "decode: out-of-range score -> ValidationFailure" $
+ fromModel @Bio (obj "long enough tagline" 150)
+ @?= Left (ValidationFailure "score: maximum 100 violated"),
+ testCase "contrast: the unconstrained record accepts the short value" $
+ fromModel @BioUnchecked (obj "short" 50)
+ @?= Right (BioUnchecked {tagline = "short", score = 50})
+ ]
+ where
+ obj t s =
+ Object
+ ( KM.fromList
+ [ ("tagline", String t),
+ ("score", Number (fromIntegral (s :: Int)))
+ ]
+ )
diff --git a/test/EndToEndSpec.hs b/test/EndToEndSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/EndToEndSpec.hs
@@ -0,0 +1,106 @@
+-- | M6: the full path tied together against a fake 'LLM' interpreter (no
+-- network): input record -> rendered request -> canned response -> parsed output
+-- record, plus the error path. The real @predict@ belongs to EP-4; this defines a
+-- tiny local @runSig@ driver.
+module EndToEndSpec (tests) where
+
+import Baikai
+ ( AssistantContent (..),
+ Model,
+ Response,
+ _Model,
+ _Response,
+ _TextContent,
+ )
+import Control.Lens ((&), (.~))
+import Data.Generics.Labels ()
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Vector qualified as V
+import Effectful (Eff, runEff, type (:>))
+import Effectful.Dispatch.Dynamic (interpret)
+import Fixtures (Article, Author (..), Sentiment (..), Summary (..), sampleArticle, sampleSummary)
+import Shikumi.Adapter (Adapter (..), fallbackAdapter)
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.LLM (LLM (..), complete)
+import Shikumi.Schema.Types (field)
+import Shikumi.Signature (Demo (..), Signature, mkSignature, setDemos)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (testCase, (@?=))
+
+-- | Interpret the 'LLM' effect by returning a fixed canned response, ignoring the
+-- request entirely.
+runFakeLLM :: Response -> Eff (LLM : es) a -> Eff es a
+runFakeLLM canned = interpret $ \_ -> \case
+ Complete _ _ _ -> pure canned
+ Stream _ _ _ -> pure []
+
+-- | A minimal @predict@-like driver: render, call the (fake) model, parse.
+runSig :: (LLM :> es) => Adapter i o -> Signature i o -> i -> Eff es (Either ShikumiError o)
+runSig adapter signature i = do
+ let (ctx, opts) = render adapter signature i
+ resp <- complete (_Model :: Model) ctx opts
+ pure (parse adapter signature resp)
+
+mkResponse :: Text -> Response
+mkResponse t =
+ _Response & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t))
+
+sig :: Signature Article Summary
+sig = setDemos [Demo sampleArticle sampleSummary] (mkSignature "Summarize the article")
+
+goodResponse :: Response
+goodResponse =
+ mkResponse $
+ T.intercalate
+ "\n"
+ [ "[[ ## headline ## ]]",
+ "Shikumi types LM programs",
+ "[[ ## bullets ## ]]",
+ "[\"records in\", \"records out\", \"errors are typed\"]",
+ "[[ ## author ## ]]",
+ "{\"name\": \"Ada\"}",
+ "[[ ## sentiment ## ]]",
+ "Positive",
+ "[[ ## note ## ]]",
+ "null",
+ "[[ ## completed ## ]]"
+ ]
+
+malformedResponse :: Response
+malformedResponse =
+ mkResponse $
+ T.intercalate
+ "\n"
+ [ "[[ ## headline ## ]]",
+ "Shikumi types LM programs",
+ "[[ ## author ## ]]",
+ "{\"name\": \"Ada\"}",
+ "[[ ## sentiment ## ]]",
+ "Positive",
+ "[[ ## note ## ]]",
+ "null",
+ "[[ ## completed ## ]]"
+ ]
+
+expectedSummary :: Summary
+expectedSummary =
+ Summary
+ { headline = field "Shikumi types LM programs",
+ bullets = field ["records in", "records out", "errors are typed"],
+ author = Author {name = field "Ada"},
+ sentiment = Positive,
+ note = Nothing
+ }
+
+tests :: TestTree
+tests =
+ testGroup
+ "EndToEndSpec"
+ [ testCase "good response decodes to a Summary through the fake LLM" $ do
+ out <- runEff . runFakeLLM goodResponse $ runSig fallbackAdapter sig sampleArticle
+ out @?= Right expectedSummary,
+ testCase "malformed response -> MissingField bullets" $ do
+ out <- runEff . runFakeLLM malformedResponse $ runSig fallbackAdapter sig sampleArticle
+ out @?= Left (MissingField "bullets")
+ ]
diff --git a/test/ErrorSpec.hs b/test/ErrorSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/ErrorSpec.hs
@@ -0,0 +1,35 @@
+-- | Pins the 'BaikaiError' -> 'ShikumiError' mapping and the 'isTransient'
+-- classification. Deliberately breaking the mapping (e.g. mapping 'DecodeError'
+-- to 'ProviderFailure') makes this group fail.
+module ErrorSpec (tests) where
+
+import Baikai.Error (BaikaiError (..))
+import Shikumi.Error
+ ( ShikumiError (..),
+ fromBaikaiError,
+ isTransient,
+ )
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (testCase, (@?=))
+
+tests :: TestTree
+tests =
+ testGroup
+ "ErrorSpec"
+ [ testCase "maps ProviderError -> ProviderFailure" $
+ fromBaikaiError (ProviderError "x") @?= ProviderFailure "x",
+ testCase "maps DecodeError -> InvalidJSON" $
+ fromBaikaiError (DecodeError "y") @?= InvalidJSON "y",
+ testCase "maps RequestInvalid -> SchemaMismatch" $
+ fromBaikaiError (RequestInvalid "z") @?= SchemaMismatch "invalid request: z",
+ testCase "maps ProcessError -> ProviderFailure" $
+ fromBaikaiError (ProcessError 2 "boom") @?= ProviderFailure "process exited 2: boom",
+ testCase "isTransient classification" $ do
+ isTransient (ProviderFailure "") @?= True
+ isTransient (Timeout "") @?= True
+ isTransient (BudgetExceeded "") @?= False
+ isTransient (SchemaMismatch "") @?= False
+ isTransient (InvalidJSON "") @?= False
+ isTransient (MissingField "") @?= False
+ isTransient (ValidationFailure "") @?= False
+ ]
diff --git a/test/Fixtures.hs b/test/Fixtures.hs
new file mode 100644
--- /dev/null
+++ b/test/Fixtures.hs
@@ -0,0 +1,88 @@
+{-# LANGUAGE DataKinds #-}
+
+-- | Shared record fixtures for the EP-3 specs (schema, signature, adapter,
+-- end-to-end). An @Article@ input and a @Summary@ output exercise every supported
+-- shape: described 'Field's, a nested record, a list, an enum-like sum, and an
+-- optional field, plus a 'Validatable' rule on @Summary@.
+module Fixtures
+ ( Sentiment (..),
+ Author (..),
+ Article (..),
+ Summary (..),
+ sampleArticle,
+ sampleSummary,
+ )
+where
+
+import Data.Text (Text)
+import GHC.Generics (Generic)
+import Shikumi.Adapter (ToPrompt)
+import Shikumi.Schema (FromModel, ToSchema, Validatable (..))
+import Shikumi.Schema.Types (Field (..), field)
+
+data Sentiment = Positive | Neutral | Negative
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema Sentiment
+
+instance FromModel Sentiment
+
+newtype Author = Author
+ {name :: Field "Author full name" Text}
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema Author
+
+instance FromModel Author
+
+data Article = Article
+ { title :: Field "The article's headline" Text,
+ body :: Field "The full article text" Text
+ }
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema Article
+
+instance FromModel Article
+
+instance ToPrompt Article
+
+data Summary = Summary
+ { headline :: Field "A one-line summary" Text,
+ bullets :: Field "Three to five key points" [Text],
+ author :: Author,
+ sentiment :: Sentiment,
+ note :: Maybe Text
+ }
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema Summary
+
+instance FromModel Summary
+
+instance ToPrompt Summary
+
+-- | A domain rule: a summary must have three to five bullet points.
+instance Validatable Summary where
+ validate s
+ | n < 3 || n > 5 = Left "bullets: must have 3 to 5 items"
+ | otherwise = Right s
+ where
+ n = length (unField (bullets s))
+
+sampleArticle :: Article
+sampleArticle =
+ Article
+ { title = field "Typed LM programs",
+ body = field "Shikumi makes LM calls behave like ordinary typed software."
+ }
+
+sampleSummary :: Summary
+sampleSummary =
+ Summary
+ { headline = field "Shikumi types LM programs",
+ bullets = field ["records in", "records out", "errors are typed"],
+ author = Author {name = field "Ada"},
+ sentiment = Positive,
+ note = Nothing
+ }
diff --git a/test/LLMSpec.hs b/test/LLMSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/LLMSpec.hs
@@ -0,0 +1,42 @@
+-- | Hermetic end-to-end test of the bare 'LLM' interpreter: a value flows
+-- @call site -> send -> interpret -> Baikai effect -> baikai (isolated registry)
+-- -> Response -> Eff -> caller@, and a genuine baikai dispatch failure surfaces
+-- through the 'BaikaiError' -> 'ShikumiError' mapping.
+module LLMSpec (tests) where
+
+import Baikai (flattenAssistantBlocks, newProviderRegistry)
+import Effectful (runEff)
+import Effectful.Error.Static (runErrorNoCallStack)
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.LLM (complete, runLLMWith)
+import StubProvider
+ ( flattenAssistantText,
+ stubContext,
+ stubModel,
+ stubOptions,
+ stubRegistry,
+ )
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (assertFailure, testCase, (@?=))
+
+tests :: TestTree
+tests =
+ testGroup
+ "LLMSpec"
+ [ testCase "stub end-to-end returns text" $ do
+ reg <- stubRegistry "hello from stub"
+ res <-
+ runEff . runErrorNoCallStack @ShikumiError . runLLMWith reg $ do
+ r <- complete stubModel stubContext stubOptions
+ pure (flattenAssistantText (flattenAssistantBlocks r))
+ res @?= Right "hello from stub",
+ testCase "unregistered tag -> ProviderFailure" $ do
+ reg <- newProviderRegistry -- empty registry: no handler for the stub tag
+ res <-
+ runEff . runErrorNoCallStack @ShikumiError . runLLMWith reg $ do
+ r <- complete stubModel stubContext stubOptions
+ pure (flattenAssistantText (flattenAssistantBlocks r))
+ case res of
+ Left (ProviderFailure _) -> pure ()
+ other -> assertFailure ("expected Left (ProviderFailure ...), got " <> show other)
+ ]
diff --git a/test/LiveSpec.hs b/test/LiveSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/LiveSpec.hs
@@ -0,0 +1,48 @@
+-- | A live, end-to-end demo proving the same 'LLM' effect drives a real provider
+-- through the full resilience stack. Gated at runtime on @SHIKUMI_LIVE@: with
+-- @SHIKUMI_LIVE=1@ (and a valid @OPENAI_API_KEY@) it registers the OpenAI
+-- provider, runs 'complete' through 'runLLMResilient', and asserts a non-empty
+-- reply. Without the variable it prints a skip line and stays green, so the
+-- default test run remains hermetic.
+module LiveSpec (tests) where
+
+import Baikai hiding (complete)
+import Baikai.Models.Generated (openai_gpt_4o_mini)
+import Baikai.Prelude
+import Baikai.Provider.OpenAI.Api qualified as OpenAI
+import Data.Text qualified as T
+import Data.Vector qualified as V
+import Effectful (runEff)
+import Effectful.Concurrent (runConcurrent)
+import Effectful.Error.Static (runErrorNoCallStack)
+import Shikumi.Error (ShikumiError)
+import Shikumi.LLM (complete, defaultLLMConfig, runLLMResilient)
+import StubProvider (flattenAssistantText)
+import System.Environment (lookupEnv)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (assertBool, assertFailure, testCase)
+
+tests :: TestTree
+tests =
+ testGroup
+ "LiveSpec"
+ [ testCase "live provider call returns text" $ do
+ live <- lookupEnv "SHIKUMI_LIVE"
+ case live of
+ Just "1" -> do
+ OpenAI.register
+ let ctx = _Context & #messages .~ V.singleton (user "Reply with a single word.")
+ opts = _Options & #maxTokens .~ Just 16
+ cfg = defaultLLMConfig globalProviderRegistry
+ out <-
+ runEff . runConcurrent . runErrorNoCallStack @ShikumiError . runLLMResilient cfg $ do
+ r <- complete openai_gpt_4o_mini ctx opts
+ pure (flattenAssistantText (flattenAssistantBlocks r))
+ case out of
+ Right t -> do
+ putStrLn ("LIVE: " <> T.unpack t)
+ assertBool "non-empty reply" (not (T.null t))
+ Left e -> assertFailure ("live call failed: " <> show e)
+ _ ->
+ putStrLn "SHIKUMI_LIVE not set; skipping live test"
+ ]
diff --git a/test/Main.hs b/test/Main.hs
new file mode 100644
--- /dev/null
+++ b/test/Main.hs
@@ -0,0 +1,58 @@
+-- | Test entry point for shikumi. Drives the hermetic specs through one tasty
+-- 'defaultMain'; the live spec (M5) is gated at runtime on @SHIKUMI_LIVE@.
+module Main (main) where
+
+import AdapterSpec qualified
+import CombinatorSpec qualified
+import ConstraintSpec qualified
+import EndToEndSpec qualified
+import ErrorSpec qualified
+import LLMSpec qualified
+import LiveSpec qualified
+import ModuleSpec qualified
+import MultimodalAdapterSpec qualified
+import MultimodalEndToEndSpec qualified
+import MultimodalSpec qualified
+import ProgramAcceptanceSpec qualified
+import ProgramSpec qualified
+import RefineSpec qualified
+import ResilienceSpec qualified
+import RoutingSpec qualified
+import SchemaSpec qualified
+import SerializeSpec qualified
+import Shikumi.Effect.TimeSpec qualified
+import SignatureSpec qualified
+import StreamSpec qualified
+import Test.Tasty (defaultMain, testGroup)
+import TwoStepSpec qualified
+import XmlAdapterSpec qualified
+
+main :: IO ()
+main =
+ defaultMain $
+ testGroup
+ "shikumi"
+ [ ErrorSpec.tests,
+ SchemaSpec.tests,
+ SignatureSpec.tests,
+ AdapterSpec.tests,
+ EndToEndSpec.tests,
+ LLMSpec.tests,
+ ResilienceSpec.tests,
+ ProgramSpec.tests,
+ RoutingSpec.tests,
+ SerializeSpec.tests,
+ ModuleSpec.tests,
+ ProgramAcceptanceSpec.tests,
+ CombinatorSpec.tests,
+ RefineSpec.tests,
+ MultimodalSpec.tests,
+ MultimodalAdapterSpec.tests,
+ MultimodalEndToEndSpec.tests,
+ XmlAdapterSpec.tests,
+ ConstraintSpec.tests,
+ TwoStepSpec.tests,
+ StreamSpec.tests,
+ Shikumi.Effect.TimeSpec.tests,
+ LiveSpec.tests
+ ]
diff --git a/test/ModuleSpec.hs b/test/ModuleSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/ModuleSpec.hs
@@ -0,0 +1,63 @@
+-- | EP-4 M4: the @predict@ and @chainOfThought@ modules. @predict@ builds a single
+-- default-parameter node; @chainOfThought@ extends the output with a leading
+-- @reasoning@ field and (in the bare form) projects it back out, while the
+-- underlying node's parameters stay reachable via @mapParamsAt@.
+module ModuleSpec (tests) where
+
+import Data.IORef (newIORef)
+import Effectful (runEff)
+import Effectful.Error.Static (runErrorNoCallStack)
+import ProgramFixtures
+ ( Outline (..),
+ Topic (..),
+ markerBody,
+ mkResponse,
+ runScriptedLLM,
+ topicToOutline,
+ )
+import Shikumi.Error (ShikumiError)
+import Shikumi.Module (WithReasoning (..), chainOfThought, chainOfThoughtRaw, predict)
+import Shikumi.Program (Params (..), emptyParams, foldParams, mapParamsAt, runProgram)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (testCase, (@?=))
+
+tests :: TestTree
+tests =
+ testGroup
+ "ModuleSpec"
+ [ testCase "predict builds a single default-parameter node" $
+ foldParams (predict topicToOutline) @?= [emptyParams],
+ testCase "chainOfThoughtRaw yields a WithReasoning value through the stub" $ do
+ ref <-
+ newIORef
+ [ mkResponse
+ ( markerBody
+ [ ("reasoning", "first I consider the topic"),
+ ("value", "{\"points\": [\"a\", \"b\", \"c\"]}")
+ ]
+ )
+ ]
+ out <-
+ runEff . runErrorNoCallStack @ShikumiError . runScriptedLLM ref $
+ runProgram (chainOfThoughtRaw topicToOutline) (Topic "haskell")
+ out
+ @?= Right
+ (WithReasoning "first I consider the topic" (Outline {points = ["a", "b", "c"]})),
+ testCase "chainOfThought projects out the bare output" $ do
+ ref <-
+ newIORef
+ [ mkResponse
+ ( markerBody
+ [ ("reasoning", "thinking"),
+ ("value", "{\"points\": [\"x\", \"y\"]}")
+ ]
+ )
+ ]
+ out <-
+ runEff . runErrorNoCallStack @ShikumiError . runScriptedLLM ref $
+ runProgram (chainOfThought topicToOutline) (Topic "haskell")
+ out @?= Right (Outline {points = ["x", "y"]}),
+ testCase "the chain-of-thought node's params are reachable via mapParamsAt 0" $
+ foldParams (mapParamsAt 0 (\p -> p {instructionOverride = Just "cot"}) (chainOfThought topicToOutline))
+ @?= [emptyParams {instructionOverride = Just "cot"}]
+ ]
diff --git a/test/MultimodalAdapterSpec.hs b/test/MultimodalAdapterSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/MultimodalAdapterSpec.hs
@@ -0,0 +1,95 @@
+{-# LANGUAGE DataKinds #-}
+
+-- | EP-24 M2: the schema/adapter seam recognises an image input field and
+-- 'Shikumi.Adapter.render' lowers it to a baikai @UserImage@ block alongside the
+-- text fields, with the image field's name dropped from the prompt text. An
+-- image-free control input produces no image blocks (the regression-safe path).
+module MultimodalAdapterSpec (tests) where
+
+import Baikai
+ ( Context,
+ ImageContent (..),
+ Message (UserMessage),
+ TextContent (..),
+ UserContent (..),
+ )
+import Control.Lens ((^.))
+import Data.ByteString qualified as BS
+import Data.Generics.Labels ()
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Vector qualified as V
+import Fixtures (Article, sampleArticle)
+import GHC.Generics (Generic)
+import Shikumi.Adapter (Adapter (..), ToPrompt, fallbackAdapter, toPrompt)
+import Shikumi.Multimodal (Image, imageFromBytes)
+import Shikumi.Schema (FromModel, ToSchema, Validatable)
+import Shikumi.Schema.Types (Field, field)
+import Shikumi.Signature (Signature, mkSignature)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (testCase, (@?=))
+
+data Describe = Describe
+ { image :: Image,
+ question :: Field "What to ask about the image" Text
+ }
+ deriving stock (Generic, Show)
+ deriving anyclass (ToPrompt)
+
+newtype Answer = Answer
+ {answer :: Field "A short answer to the question" Text}
+ deriving stock (Generic, Show, Eq)
+ deriving anyclass (ToSchema, FromModel, ToPrompt)
+
+instance Validatable Answer
+
+-- | Every 'ImageContent' carried by a user message in the context, in order.
+userImageBlocks :: Context -> [ImageContent]
+userImageBlocks ctx =
+ [ ic
+ | UserMessage p <- V.toList (ctx ^. #messages),
+ UserImage ic <- V.toList (p ^. #content)
+ ]
+
+-- | Every user text block's text, in order.
+userTextBlocks :: Context -> [Text]
+userTextBlocks ctx =
+ [ t
+ | UserMessage p <- V.toList (ctx ^. #messages),
+ UserText (TextContent t) <- V.toList (p ^. #content)
+ ]
+
+rawPng :: BS.ByteString
+rawPng = BS.pack [0x89, 0x50, 0x4e, 0x47]
+
+describeInput :: Describe
+describeInput =
+ Describe
+ { image = imageFromBytes "image/png" rawPng,
+ question = field "What vehicle is shown?"
+ }
+
+describeSig :: Signature Describe Answer
+describeSig = mkSignature "Answer the question about the image"
+
+articleSig :: Signature Article Answer
+articleSig = mkSignature "Summarise"
+
+tests :: TestTree
+tests =
+ testGroup
+ "MultimodalAdapterSpec"
+ [ testCase "render lowers an image field to a UserImage block with the right bytes" $ do
+ let (ctx, _) = render fallbackAdapter describeSig describeInput
+ userImageBlocks ctx @?= [ImageContent {imageData = rawPng, mimeType = "image/png"}],
+ testCase "render keeps the text fields and drops the image field name" $ do
+ let (ctx, _) = render fallbackAdapter describeSig describeInput
+ texts = userTextBlocks ctx
+ any (T.isInfixOf "What vehicle is shown?") texts @?= True
+ -- the image field's name must not leak into the prompt text
+ any (T.isInfixOf "image:") texts @?= False,
+ testCase "image-free input produces no image blocks (text path unchanged)" $ do
+ let (ctx, _) = render fallbackAdapter articleSig sampleArticle
+ userImageBlocks ctx @?= []
+ userTextBlocks ctx @?= [toPrompt sampleArticle]
+ ]
diff --git a/test/MultimodalEndToEndSpec.hs b/test/MultimodalEndToEndSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/MultimodalEndToEndSpec.hs
@@ -0,0 +1,131 @@
+{-# LANGUAGE DataKinds #-}
+
+-- | EP-24 M3: the full render -> stub -> parse loop for an image-bearing
+-- signature. A capturing 'LLM' interpreter records the 'Context' it is handed so
+-- the test can prove the model genuinely received a @UserImage@ block, then
+-- returns a canned @[[ ## answer ## ]]@ body that decodes into a typed 'Answer'.
+-- A second test locks in the regression invariant: an image-free input renders a
+-- 'Context' with no image blocks and the unchanged text-only user turn.
+module MultimodalEndToEndSpec (tests) where
+
+import Baikai
+ ( AssistantContent (..),
+ Context,
+ ImageContent (..),
+ Message (UserMessage),
+ Model,
+ Response,
+ TextContent (..),
+ UserContent (..),
+ _Model,
+ _Response,
+ _TextContent,
+ )
+import Control.Lens ((&), (.~), (^.))
+import Data.ByteString qualified as BS
+import Data.Generics.Labels ()
+import Data.IORef (IORef, newIORef, readIORef, writeIORef)
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Vector qualified as V
+import Effectful (Eff, IOE, liftIO, runEff, type (:>))
+import Effectful.Dispatch.Dynamic (interpret)
+import Fixtures (Article, sampleArticle)
+import GHC.Generics (Generic)
+import Shikumi.Adapter (Adapter (..), ToPrompt, fallbackAdapter, toPrompt)
+import Shikumi.Error (ShikumiError)
+import Shikumi.LLM (LLM (..), complete)
+import Shikumi.Multimodal (Image, imageFromBytes)
+import Shikumi.Schema (FromModel, ToSchema, Validatable)
+import Shikumi.Schema.Types (Field, field)
+import Shikumi.Signature (Signature, mkSignature)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (testCase, (@?=))
+
+data Describe = Describe
+ { image :: Image,
+ question :: Field "What to ask about the image" Text
+ }
+ deriving stock (Generic, Show)
+ deriving anyclass (ToPrompt)
+
+newtype Answer = Answer
+ {answer :: Field "A short answer to the question" Text}
+ deriving stock (Generic, Show, Eq)
+ deriving anyclass (ToSchema, FromModel, ToPrompt)
+
+instance Validatable Answer
+
+-- | Interpret 'LLM' by capturing the request 'Context' into an 'IORef' and
+-- returning a fixed canned response. Runs under 'IOE' (the suite uses 'runEff').
+runCapturingLLM :: (IOE :> es) => IORef (Maybe Context) -> Response -> Eff (LLM : es) a -> Eff es a
+runCapturingLLM capture canned = interpret $ \_ -> \case
+ Complete _ ctx _ -> do
+ liftIO (writeIORef capture (Just ctx))
+ pure canned
+ Stream _ _ _ -> pure []
+
+-- | render -> call (captured) model -> parse.
+runSig :: (LLM :> es) => Adapter i o -> Signature i o -> i -> Eff es (Either ShikumiError o)
+runSig adapter signature i = do
+ let (ctx, opts) = render adapter signature i
+ resp <- complete (_Model :: Model) ctx opts
+ pure (parse adapter signature resp)
+
+mkResponse :: Text -> Response
+mkResponse t =
+ _Response & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t))
+
+cannedAnswer :: Response
+cannedAnswer =
+ mkResponse $
+ T.intercalate
+ "\n"
+ [ "[[ ## answer ## ]]",
+ "A red bicycle.",
+ "[[ ## completed ## ]]"
+ ]
+
+userImageBlocks :: Context -> [ImageContent]
+userImageBlocks ctx =
+ [ ic
+ | UserMessage p <- V.toList (ctx ^. #messages),
+ UserImage ic <- V.toList (p ^. #content)
+ ]
+
+userTextBlocks :: Context -> [Text]
+userTextBlocks ctx =
+ [ t
+ | UserMessage p <- V.toList (ctx ^. #messages),
+ UserText (TextContent t) <- V.toList (p ^. #content)
+ ]
+
+rawPng :: BS.ByteString
+rawPng = BS.pack [0x89, 0x50, 0x4e, 0x47]
+
+tests :: TestTree
+tests =
+ testGroup
+ "MultimodalEndToEndSpec"
+ [ testCase "Describe with an image: model receives a UserImage block; answer decodes" $ do
+ capture <- newIORef Nothing
+ let img = imageFromBytes "image/png" rawPng
+ input = Describe {image = img, question = field "What vehicle is shown?"}
+ sig = mkSignature "Answer the question about the image" :: Signature Describe Answer
+ out <-
+ runEff . runCapturingLLM capture cannedAnswer $
+ runSig fallbackAdapter sig input
+ out @?= Right (Answer {answer = field "A red bicycle."})
+ Just ctx <- readIORef capture
+ userImageBlocks ctx @?= [ImageContent {imageData = rawPng, mimeType = "image/png"}]
+ any (T.isInfixOf "What vehicle is shown?") (userTextBlocks ctx) @?= True,
+ testCase "image-free input renders the unchanged text-only Context" $ do
+ capture <- newIORef Nothing
+ let sig = mkSignature "Summarize" :: Signature Article Answer
+ _ <-
+ runEff . runCapturingLLM capture cannedAnswer $
+ runSig fallbackAdapter sig sampleArticle
+ Just ctx <- readIORef capture
+ userImageBlocks ctx @?= []
+ userTextBlocks ctx @?= [toPrompt sampleArticle]
+ ]
diff --git a/test/MultimodalSpec.hs b/test/MultimodalSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/MultimodalSpec.hs
@@ -0,0 +1,46 @@
+-- | EP-24 M1: the 'Image' value type and its lowering to baikai's 'ImageContent'.
+-- Hermetic: a known byte string is round-tripped through a file and through
+-- base64, and both paths must reach 'ImageContent' with the exact decoded bytes
+-- (never base64) and the correct MIME type.
+module MultimodalSpec (tests) where
+
+import Baikai (ImageContent (..))
+import Data.ByteString qualified as BS
+import Data.ByteString.Base64 qualified as Base64
+import Data.Text.Encoding (decodeUtf8)
+import Shikumi.Multimodal
+ ( imageBytes,
+ imageFromBase64,
+ imageFromFile,
+ imageMime,
+ imageToContent,
+ )
+import System.Directory (getTemporaryDirectory, removeFile)
+import System.IO (hClose, openTempFile)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (assertFailure, testCase, (@?=))
+
+tests :: TestTree
+tests =
+ testGroup
+ "MultimodalSpec"
+ [ testCase "imageFromBase64 round-trips decoded bytes into ImageContent" $ do
+ let raw = BS.pack [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a] -- PNG magic
+ b64 = decodeUtf8 (Base64.encode raw)
+ img <- either (assertFailure . show) pure (imageFromBase64 "image/png" b64)
+ imageBytes img @?= raw
+ let ic = imageToContent img
+ imageData ic @?= raw
+ mimeType ic @?= "image/png",
+ testCase "imageFromFile reads bytes and infers MIME from .png" $ do
+ tmp <- getTemporaryDirectory
+ (fp, h) <- openTempFile tmp "shikumi-img.png"
+ let raw = BS.pack [0x89, 0x50, 0x4e, 0x47]
+ BS.hPut h raw
+ hClose h
+ res <- imageFromFile fp
+ img <- either (assertFailure . show) pure res
+ imageBytes img @?= raw
+ imageMime img @?= "image/png"
+ removeFile fp
+ ]
diff --git a/test/ProgramAcceptanceSpec.hs b/test/ProgramAcceptanceSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/ProgramAcceptanceSpec.hs
@@ -0,0 +1,105 @@
+-- | EP-4 M5 — the end-to-end acceptance behaviors named in the plan's Purpose:
+--
+-- * __typed composition runs end-to-end__: two record-typed signatures composed
+-- with 'pipeline' run through the stub @LLM@ and return a typed @Draft@;
+-- * __the program is rewritable as data, and the rewrite reaches the wire__:
+-- editing node 0's instruction with 'mapParamsAt' shows the new instruction in
+-- both @foldParams@ (as data) and in the captured first-stage prompt (on the
+-- wire), while node 1 is untouched.
+module ProgramAcceptanceSpec (tests) where
+
+import Data.Aeson (Value, object, (.=))
+import Data.IORef (newIORef, readIORef)
+import Data.Text (Text)
+import Data.Text qualified as T
+import Effectful (runEff)
+import Effectful.Error.Static (runErrorNoCallStack)
+import ProgramFixtures
+ ( Draft (..),
+ Topic (..),
+ draftResponse,
+ outlineResponse,
+ outlineToDraft,
+ runRecordingLLM,
+ runScriptedLLM,
+ topicToOutline,
+ )
+import Shikumi.Error (ShikumiError)
+import Shikumi.Module (predict)
+import Shikumi.Program
+ ( Demo (..),
+ Params (..),
+ Program,
+ foldParams,
+ mapParamsAt,
+ pipeline,
+ runProgram,
+ )
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (assertBool, testCase, (@?=))
+
+-- The headline two-stage program: Topic -> Outline -> Draft.
+essay :: Program Topic Draft
+essay = pipeline (predict topicToOutline) (predict outlineToDraft)
+
+-- A JSON demo whose halves decode to a Topic and an Outline (so it can be spliced
+-- into the rendered prompt without error).
+sampleDemo :: Demo
+sampleDemo =
+ Demo
+ { input = topicJSON "demo topic",
+ output = outlineJSON ["alpha", "beta"]
+ }
+
+topicJSON :: Text -> Value
+topicJSON s = object ["subject" .= s]
+
+outlineJSON :: [Text] -> Value
+outlineJSON ps = object ["points" .= ps]
+
+tests :: TestTree
+tests =
+ testGroup
+ "ProgramAcceptanceSpec"
+ [ testCase "Behavior 1: a typed two-stage pipeline returns a Draft" $ do
+ ref <- newIORef [outlineResponse, draftResponse]
+ out <-
+ runEff . runErrorNoCallStack @ShikumiError . runScriptedLLM ref $
+ runProgram essay (Topic "haskell")
+ out @?= Right (Draft {prose = "A fine essay about the outline."}),
+ testCase "Behavior 2 (data): rewriting node 0 changes only node 0's params" $ do
+ let rewritten =
+ mapParamsAt
+ 0
+ (\p -> p {instructionOverride = Just "NEW INSTRUCTION", demos = [sampleDemo]})
+ essay
+ params = foldParams rewritten
+ length params @?= 2
+ instructionOverride (params !! 0) @?= Just "NEW INSTRUCTION"
+ demos (params !! 0) @?= [sampleDemo]
+ params !! 1 @?= emptyParamsLike,
+ testCase "Behavior 2 (wire): the new instruction appears in the captured first-stage prompt" $ do
+ capture <- newIORef []
+ ref <- newIORef [outlineResponse, draftResponse]
+ let rewritten =
+ mapParamsAt
+ 0
+ (\p -> p {instructionOverride = Just "NEW INSTRUCTION", demos = [sampleDemo]})
+ essay
+ out <-
+ runEff . runErrorNoCallStack @ShikumiError . runRecordingLLM capture ref $
+ runProgram rewritten (Topic "haskell")
+ out @?= Right (Draft {prose = "A fine essay about the outline."})
+ prompts <- readIORef capture
+ case prompts of
+ (firstStage : _) -> do
+ length prompts @?= 2
+ assertBool
+ "first-stage prompt should contain the new instruction"
+ ("NEW INSTRUCTION" `T.isInfixOf` firstStage)
+ [] -> assertBool "expected captured prompts" False
+ ]
+ where
+ -- node 1 should remain the default overlay
+ emptyParamsLike :: Params
+ emptyParamsLike = Params Nothing []
diff --git a/test/ProgramFixtures.hs b/test/ProgramFixtures.hs
new file mode 100644
--- /dev/null
+++ b/test/ProgramFixtures.hs
@@ -0,0 +1,177 @@
+-- | Shared fixtures for the EP-4 program specs: three record-typed signatures
+-- (@Topic -> Outline -> Draft@) plus a self-looping @Cell -> Cell@ used to build
+-- arbitrarily long pipelines for the ordering-law property; and two network-free
+-- 'LLM' interpreters — a /scripted/ one that pops canned responses in order, and a
+-- /recording/ one that also captures each rendered request's system prompt so a
+-- test can assert what the program actually sent on the wire.
+module ProgramFixtures
+ ( -- * Records
+ Topic (..),
+ Outline (..),
+ Draft (..),
+ Cell (..),
+
+ -- * Signatures
+ topicToOutline,
+ outlineToDraft,
+ cellSig,
+
+ -- * Canned responses
+ mkResponse,
+ outlineResponse,
+ draftResponse,
+ cellResponse,
+ markerBody,
+
+ -- * Fake LLM interpreters
+ runScriptedLLM,
+ runRecordingLLM,
+ )
+where
+
+import Baikai
+ ( AssistantContent (..),
+ Context,
+ Response,
+ _Response,
+ _TextContent,
+ )
+import Control.Lens ((&), (.~), (^.))
+import Data.Generics.Labels ()
+import Data.IORef (IORef, atomicModifyIORef', modifyIORef')
+import Data.Maybe (fromMaybe)
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Vector qualified as V
+import Effectful (Eff, IOE, liftIO, type (:>))
+import Effectful.Dispatch.Dynamic (interpret)
+import GHC.Generics (Generic)
+import Shikumi.Adapter (ToPrompt)
+import Shikumi.LLM (LLM (..))
+import Shikumi.Schema (FromModel, ToSchema, Validatable)
+import Shikumi.Signature (Signature, mkSignature)
+
+-- ---------------------------------------------------------------------------
+-- Records (plain fields — no Field wrappers needed for these tests)
+-- ---------------------------------------------------------------------------
+
+newtype Topic = Topic {subject :: Text}
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema Topic
+
+instance FromModel Topic
+
+instance ToPrompt Topic
+
+instance Validatable Topic
+
+newtype Outline = Outline {points :: [Text]}
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema Outline
+
+instance FromModel Outline
+
+instance ToPrompt Outline
+
+instance Validatable Outline
+
+newtype Draft = Draft {prose :: Text}
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema Draft
+
+instance FromModel Draft
+
+instance ToPrompt Draft
+
+instance Validatable Draft
+
+-- | A self-looping record, so @predict cellSig :: Program Cell Cell@ can be
+-- composed with itself to make an N-node pipeline.
+newtype Cell = Cell {cell :: Text}
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema Cell
+
+instance FromModel Cell
+
+instance ToPrompt Cell
+
+instance Validatable Cell
+
+-- ---------------------------------------------------------------------------
+-- Signatures
+-- ---------------------------------------------------------------------------
+
+topicToOutline :: Signature Topic Outline
+topicToOutline = mkSignature "Outline the topic into points"
+
+outlineToDraft :: Signature Outline Draft
+outlineToDraft = mkSignature "Draft prose from the outline"
+
+cellSig :: Signature Cell Cell
+cellSig = mkSignature "Echo the cell"
+
+-- ---------------------------------------------------------------------------
+-- Canned responses (rendered as the fallback adapter's [[ ## field ## ]] sections)
+-- ---------------------------------------------------------------------------
+
+-- | Build a fallback-style response body from @(field, value)@ sections, closing
+-- with the @completed@ marker.
+markerBody :: [(Text, Text)] -> Text
+markerBody fields = T.unlines (concatMap sect fields ++ ["[[ ## completed ## ]]"])
+ where
+ sect (k, v) = ["[[ ## " <> k <> " ## ]]", v]
+
+-- | An assistant 'Response' carrying @t@ as its single text block.
+mkResponse :: Text -> Response
+mkResponse t =
+ _Response & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t))
+
+outlineResponse :: Response
+outlineResponse = mkResponse (markerBody [("points", "[\"intro\", \"body\", \"end\"]")])
+
+draftResponse :: Response
+draftResponse = mkResponse (markerBody [("prose", "A fine essay about the outline.")])
+
+-- | Echo a fixed cell value (used by the @Cell -> Cell@ pipelines).
+cellResponse :: Response
+cellResponse = mkResponse (markerBody [("cell", "echoed")])
+
+-- ---------------------------------------------------------------------------
+-- Fake LLM interpreters
+-- ---------------------------------------------------------------------------
+
+-- | Pop the next scripted 'Response' for each 'Complete'; an exhausted queue
+-- yields an empty response (tests script enough). Ignores the request.
+runScriptedLLM :: (IOE :> es) => IORef [Response] -> Eff (LLM : es) a -> Eff es a
+runScriptedLLM ref = interpret $ \_ -> \case
+ Complete _ _ _ -> liftIO (pop ref)
+ Stream _ _ _ -> pure []
+
+-- | Like 'runScriptedLLM', but also append each request's rendered /system prompt/
+-- to @capture@, so a test can assert on what the program actually sent.
+runRecordingLLM ::
+ (IOE :> es) =>
+ IORef [Text] ->
+ IORef [Response] ->
+ Eff (LLM : es) a ->
+ Eff es a
+runRecordingLLM capture ref = interpret $ \_ -> \case
+ Complete _ ctx _ -> do
+ liftIO (modifyIORef' capture (++ [renderedPrompt ctx]))
+ liftIO (pop ref)
+ Stream _ _ _ -> pure []
+
+-- | What we capture from a request: its system prompt (where the instruction and
+-- output guide live).
+renderedPrompt :: Context -> Text
+renderedPrompt ctx = fromMaybe "" (ctx ^. #systemPrompt)
+
+pop :: IORef [Response] -> IO Response
+pop ref = atomicModifyIORef' ref step
+ where
+ step (x : xs) = (xs, x)
+ step [] = ([], mkResponse (markerBody []))
diff --git a/test/ProgramSpec.hs b/test/ProgramSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/ProgramSpec.hs
@@ -0,0 +1,84 @@
+-- | EP-4 M1 (run a node through the LLM effect) and M2 (the parameter interface:
+-- @foldParams@ / @mapParams@ / @mapParamsAt@ and the ordering law).
+module ProgramSpec (tests) where
+
+import Data.IORef (newIORef)
+import Data.List (foldl1')
+import Data.Text (Text)
+import Data.Text qualified as T
+import Effectful (runEff)
+import Effectful.Error.Static (runErrorNoCallStack)
+import ProgramFixtures
+ ( Cell (..),
+ Draft,
+ Outline (..),
+ Topic (..),
+ cellSig,
+ outlineResponse,
+ outlineToDraft,
+ runScriptedLLM,
+ topicToOutline,
+ )
+import Shikumi.Error (ShikumiError)
+import Shikumi.Program
+ ( Params (..),
+ Program (Compose, Predict),
+ emptyParams,
+ foldParams,
+ mapParams,
+ mapParamsAt,
+ pipeline,
+ runProgram,
+ )
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (testCase, (@?=))
+import Test.Tasty.QuickCheck (NonNegative (..), Positive (..), testProperty, (===))
+
+-- A two-stage program reused across the param tests.
+twoStage :: Program Topic Draft
+twoStage = pipeline (Predict topicToOutline emptyParams) (Predict outlineToDraft emptyParams)
+
+setInstr :: Text -> Params -> Params
+setInstr t p = p {instructionOverride = Just t}
+
+-- | A list update mirroring @mapParamsAt@'s contract, used to state the ordering
+-- law executably.
+adjust :: Int -> (a -> a) -> [a] -> [a]
+adjust n f xs = [if i == n then f x else x | (i, x) <- zip [0 ..] xs]
+
+-- | A pipeline of @k@ identical @Cell -> Cell@ predict nodes (k >= 1).
+cellPipeline :: Int -> Program Cell Cell
+cellPipeline k = foldl1' Compose (replicate k (Predict cellSig emptyParams))
+
+tests :: TestTree
+tests =
+ testGroup
+ "ProgramSpec"
+ [ testCase "M1: a single Predict node runs through the LLM effect to a typed output" $ do
+ ref <- newIORef [outlineResponse]
+ out <-
+ runEff . runErrorNoCallStack @ShikumiError . runScriptedLLM ref $
+ runProgram (Predict topicToOutline emptyParams) (Topic "haskell")
+ out @?= Right (Outline {points = ["intro", "body", "end"]}),
+ testCase "M2: foldParams returns nodes in stage order" $
+ foldParams twoStage @?= [emptyParams, emptyParams],
+ testCase "M2: mapParams touches every node" $
+ foldParams (mapParams (setInstr "X") twoStage)
+ @?= [setInstr "X" emptyParams, setInstr "X" emptyParams],
+ testCase "M2: mapParamsAt 0 edits only the first node" $
+ foldParams (mapParamsAt 0 (setInstr "first") twoStage)
+ @?= [setInstr "first" emptyParams, emptyParams],
+ testCase "M2: mapParamsAt 1 edits only the second node" $
+ foldParams (mapParamsAt 1 (setInstr "second") twoStage)
+ @?= [emptyParams, setInstr "second" emptyParams],
+ testCase "M2: mapParamsAt with an out-of-range index is the identity" $
+ foldParams (mapParamsAt 7 (setInstr "nope") twoStage)
+ @?= [emptyParams, emptyParams],
+ testProperty "M2: ordering law — foldParams . mapParamsAt n f == adjust n f . foldParams" $
+ \(Positive k') (NonNegative seed) ->
+ let k = 1 + (k' `mod` 8)
+ n = seed `mod` k
+ f = setInstr ("node " <> T.pack (show n))
+ prog = cellPipeline k
+ in foldParams (mapParamsAt n f prog) === adjust n f (foldParams prog)
+ ]
diff --git a/test/RefineSpec.hs b/test/RefineSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/RefineSpec.hs
@@ -0,0 +1,262 @@
+{-# LANGUAGE TypeApplications #-}
+
+-- | EP-18 acceptance: the reward-driven self-refinement modules
+-- ('Shikumi.Refine.bestOfN', 'Shikumi.Refine.refine',
+-- 'Shikumi.Refine.multiChainComparison').
+--
+-- Every test installs the production router (@routeLLM . runRouting _Model@) above a
+-- deterministic stub @LLM@ (see "RefineStub") so per-sample temperature reaches the
+-- wire exactly as in production — no network, no API key. The headline assertion
+-- (M4) shows a deliberately-weak program's score strictly improves once wrapped.
+module RefineSpec (tests) where
+
+import Baikai (Context, Options, _Model)
+import Data.IORef (newIORef, readIORef)
+import Data.Text (Text)
+import Data.Text qualified as T
+import Effectful (runEff)
+import Effectful.Concurrent (runConcurrent)
+import Effectful.Error.Static (runErrorNoCallStack)
+import RefineStub
+ ( Answer (..),
+ Sentence (..),
+ classifyProg,
+ decideMCC,
+ decideRefine,
+ decideTemp,
+ goodReward,
+ mccSynthSig,
+ reasonProg,
+ rewardByLabel,
+ runRecordingStub,
+ runStub,
+ runThrowingLLM,
+ )
+import Shikumi.Combinator (mapP, validate, (>>>))
+import Shikumi.Error (ShikumiError)
+import Shikumi.Module (predict)
+import Shikumi.Program
+ ( Program,
+ ProgramShape (ShapeEmbed),
+ foldParams,
+ programShape,
+ runProgram,
+ runProgramConc,
+ setProgramParams,
+ )
+import Shikumi.Refine (bestOfN, mkReward, multiChainComparison, refine)
+import Shikumi.Routing (routeLLM, runRouting)
+import Shikumi.Signature (Signature, mkSignature)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (assertBool, testCase, (@?=))
+
+tests :: TestTree
+tests =
+ testGroup
+ "Refine"
+ [ bestOfNReturnsHighest,
+ bestOfNShortCircuits,
+ bestOfNExhaustsFailCount,
+ refineClimbsViaAdvice,
+ multiChainSynthesizes,
+ moduleComposes,
+ moduleIsTransparent,
+ rewardDrivenRetryImproves
+ ]
+
+-- ---------------------------------------------------------------------------
+-- Run helpers — router over the stub, mirroring production install order
+-- ---------------------------------------------------------------------------
+
+-- | Run a program sequentially under @routeLLM . runRouting _Model@ over the
+-- decision-rule stub; returns the result and the completion count.
+runCounted ::
+ (Context -> Options -> Text) ->
+ Program i o ->
+ i ->
+ IO (Either ShikumiError o, Int)
+runCounted decide prog input = do
+ ref <- newIORef 0
+ res <-
+ runEff
+ . runErrorNoCallStack @ShikumiError
+ . runRouting _Model
+ . runStub ref decide
+ . routeLLM
+ $ runProgram prog input
+ n <- readIORef ref
+ pure (res, n)
+
+-- | As 'runCounted' but under the concurrent executor.
+runCountedConc ::
+ (Context -> Options -> Text) ->
+ Program i o ->
+ i ->
+ IO (Either ShikumiError o)
+runCountedConc decide prog input = do
+ ref <- newIORef 0
+ runEff
+ . runErrorNoCallStack @ShikumiError
+ . runConcurrent
+ . runRouting _Model
+ . runStub ref decide
+ . routeLLM
+ $ runProgramConc prog input
+
+-- | Run a program whose stub throws on every call; returns result and attempt
+-- count.
+runThrowing ::
+ Program i o ->
+ i ->
+ IO (Either ShikumiError o, Int)
+runThrowing prog input = do
+ ref <- newIORef 0
+ res <-
+ runEff
+ . runErrorNoCallStack @ShikumiError
+ . runRouting _Model
+ . runThrowingLLM ref
+ . routeLLM
+ $ runProgram prog input
+ n <- readIORef ref
+ pure (res, n)
+
+-- | Run recording each request's last user text (to assert what synthesis saw).
+runRecorded ::
+ (Context -> Options -> Text) ->
+ Program i o ->
+ i ->
+ IO (Either ShikumiError o, [Text])
+runRecorded decide prog input = do
+ cap <- newIORef []
+ res <-
+ runEff
+ . runErrorNoCallStack @ShikumiError
+ . runRouting _Model
+ . runRecordingStub cap decide
+ . routeLLM
+ $ runProgram prog input
+ texts <- readIORef cap
+ pure (res, texts)
+
+-- ---------------------------------------------------------------------------
+-- M1 — bestOfN
+-- ---------------------------------------------------------------------------
+
+-- | A reward table over the temperature-chosen labels.
+labelReward :: [(Text, Double)] -> Answer -> Double
+labelReward = rewardByLabel
+
+bestOfNReturnsHighest :: TestTree
+bestOfNReturnsHighest =
+ testCase "bestOfN returns the highest-reward sample (all attempts run)" $ do
+ let reward = mkReward (labelReward [("low", 0.2), ("mid", 0.9), ("high", 0.5)])
+ (res, n) <- runCounted decideTemp (bestOfN 3 1.0 reward classifyProg) (Sentence "x")
+ res @?= Right (Answer "mid")
+ n @?= 3
+
+bestOfNShortCircuits :: TestTree
+bestOfNShortCircuits =
+ testCase "bestOfN short-circuits on a passing threshold" $ do
+ let reward = mkReward (labelReward [("low", 0.2), ("mid", 1.0), ("high", 0.5)])
+ (res, n) <- runCounted decideTemp (bestOfN 3 1.0 reward classifyProg) (Sentence "x")
+ res @?= Right (Answer "mid")
+ n @?= 2
+
+bestOfNExhaustsFailCount :: TestTree
+bestOfNExhaustsFailCount =
+ testCase "bestOfN exhausts failCount then rethrows" $ do
+ let reward = mkReward (labelReward [("low", 0.2)])
+ (res, n) <- runThrowing (bestOfN 3 1.0 reward classifyProg) (Sentence "x")
+ assertBool "error propagates when every attempt throws" (isLeft res)
+ n @?= 3
+
+-- ---------------------------------------------------------------------------
+-- M2 — refine
+-- ---------------------------------------------------------------------------
+
+refineClimbsViaAdvice :: TestTree
+refineClimbsViaAdvice =
+ testCase "refine climbs from failing to passing via advice" $ do
+ let reward = mkReward goodReward
+ (single, _) <- runCounted decideRefine classifyProg (Sentence "x")
+ (wrapped, _) <- runCounted decideRefine (refine 2 1.0 reward classifyProg) (Sentence "x")
+ single @?= Right (Answer "bad")
+ wrapped @?= Right (Answer "good")
+ assertBool
+ "feedback strictly improves the reward"
+ (rewardOf goodReward wrapped > rewardOf goodReward single)
+
+-- ---------------------------------------------------------------------------
+-- M3 — multiChainComparison
+-- ---------------------------------------------------------------------------
+
+multiChainSynthesizes :: TestTree
+multiChainSynthesizes =
+ testCase "multiChainComparison synthesizes the modal consensus from all M attempts" $ do
+ (res, texts) <- runRecorded decideMCC (multiChainComparison 3 reasonProg mccSynthSig) (Sentence "x")
+ res @?= Right (Answer "Brussels")
+ case texts of
+ [] -> assertBool "expected at least the synthesis request" False
+ _ -> do
+ let synthText = last texts
+ assertBool "synthesis prompt shows attempt #1" (T.isInfixOf "Student Attempt #1" synthText)
+ assertBool "synthesis prompt shows attempt #2" (T.isInfixOf "Student Attempt #2" synthText)
+ assertBool "synthesis prompt shows attempt #3" (T.isInfixOf "Student Attempt #3" synthText)
+
+-- ---------------------------------------------------------------------------
+-- M4 — composition + the headline acceptance scenario
+-- ---------------------------------------------------------------------------
+
+moduleComposes :: TestTree
+moduleComposes =
+ testCase "module nests inside mapP and >>> and agrees across executors" $ do
+ let reward = mkReward (labelReward [("low", 0.2), ("mid", 0.9), ("high", 0.5)])
+ p = bestOfN 3 1.0 reward classifyProg
+ mapped = mapP 2 p
+ inputs = [Sentence "a", Sentence "b"]
+ (seqRes, _) <- runCounted decideTemp mapped inputs
+ concRes <- runCountedConc decideTemp mapped inputs
+ seqRes @?= Right [Answer "mid", Answer "mid"]
+ concRes @?= seqRes
+ -- a module feeding a downstream validator
+ let echoProg = predict (mkSignature "Echo the label." :: Signature Answer Answer)
+ guarded = p >>> validate (\(Answer a) -> a `elem` ["low", "mid", "high"]) "unknown label" echoProg
+ (gRes, _) <- runCounted decideTemp guarded (Sentence "x")
+ assertBool "module >>> validator runs without error" (isRight gRes)
+
+moduleIsTransparent :: TestTree
+moduleIsTransparent =
+ testCase "module is structurally transparent (ShapeEmbed, no params)" $ do
+ let reward = mkReward (labelReward [("mid", 1.0)])
+ p = bestOfN 3 1.0 reward classifyProg
+ programShape p @?= ShapeEmbed
+ foldParams p @?= []
+ assertBool "empty param vector loads onto the wrapper" (isRight (setProgramParams [] p))
+
+rewardDrivenRetryImproves :: TestTree
+rewardDrivenRetryImproves =
+ testCase "reward-driven retry demonstrably improves a deliberately-weak program's score" $ do
+ let table = [("low", 0.2), ("mid", 0.9), ("high", 0.5)]
+ reward = mkReward (labelReward table)
+ (single, _) <- runCounted decideTemp classifyProg (Sentence "x")
+ (wrapped, _) <- runCounted decideTemp (bestOfN 3 1.0 reward classifyProg) (Sentence "x")
+ let singleScore = rewardOf (labelReward table) single
+ wrappedScore = rewardOf (labelReward table) wrapped
+ assertBool
+ ("wrapped reward " <> show wrappedScore <> " should exceed single-shot " <> show singleScore)
+ (wrappedScore > singleScore)
+
+-- ---------------------------------------------------------------------------
+-- Small helpers
+-- ---------------------------------------------------------------------------
+
+rewardOf :: (Answer -> Double) -> Either ShikumiError Answer -> Double
+rewardOf f (Right a) = f a
+rewardOf _ (Left _) = -1
+
+isRight :: Either a b -> Bool
+isRight = either (const False) (const True)
+
+isLeft :: Either a b -> Bool
+isLeft = either (const True) (const False)
diff --git a/test/RefineStub.hs b/test/RefineStub.hs
new file mode 100644
--- /dev/null
+++ b/test/RefineStub.hs
@@ -0,0 +1,268 @@
+{-# LANGUAGE LambdaCase #-}
+
+-- | Hermetic, network-free fixtures and deterministic stub @LLM@ interpreters for
+-- the EP-18 self-refinement spec ('Shikumi.Refine').
+--
+-- The whole point of a self-refinement test is that __a later attempt scores
+-- better than the first__. So the stub is not a constant: it inspects the rendered
+-- request and the per-sample temperature ('Options.temperature', made live by the
+-- router installed above it) and answers by a rule:
+--
+-- * a /classify/ call answers a label chosen by the request's temperature
+-- (@< 0.5@ → @low@, @< 1.0@ → @mid@, else @high@), so 'Shikumi.Refine.bestOfN'
+-- — which fans its attempts out across temperatures — finds a better-scoring
+-- answer than a single low-temperature shot;
+-- * a /classify/ call whose system prompt carries an injected advice hint answers
+-- @good@ (otherwise @bad@), so 'Shikumi.Refine.refine' climbs once advice lands;
+-- * an /advice/ call (its output guide names an @advice@ field) returns a fixed
+-- advice string;
+-- * a /reasoning/ call (chain-of-thought; its guide names a @reasoning@ field)
+-- answers a temperature-chosen candidate, and the /synthesis/ call returns the
+-- modal candidate it is shown.
+module RefineStub
+ ( -- * Task records
+ Sentence (..),
+ Answer (..),
+
+ -- * Signatures and programs
+ classifySig,
+ classifyProg,
+ reasonProg,
+ synthInstruction,
+ mccSynthSig,
+
+ -- * Rewards
+ rewardByLabel,
+ goodReward,
+
+ -- * Decision rules
+ decideTemp,
+ decideRefine,
+ decideMCC,
+
+ -- * Stub interpreters
+ runStub,
+ runRecordingStub,
+ runThrowingLLM,
+
+ -- * Request readers (re-used by assertions)
+ lastUserText,
+ )
+where
+
+import Baikai (Context, Options, TextContent (..))
+import Baikai.Content (UserContent (..))
+import Baikai.Message (Message (..), UserPayload (..))
+import Control.Lens ((^.))
+import Data.Generics.Labels ()
+import Data.IORef (IORef, atomicModifyIORef', modifyIORef')
+import Data.Maybe (fromMaybe)
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Vector qualified as V
+import Effectful (Eff, IOE, liftIO, (:>))
+import Effectful.Dispatch.Dynamic (interpret)
+import Effectful.Error.Static (Error, throwError)
+import GHC.Generics (Generic)
+import ProgramFixtures (markerBody, mkResponse)
+import Shikumi.Adapter (ToPrompt)
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.LLM (LLM (..))
+import Shikumi.Module (WithReasoning, chainOfThoughtRaw, predict)
+import Shikumi.Program (Program, modal)
+import Shikumi.Refine (MultiChainInput, multiChainSig)
+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)
+
+instance ToSchema Sentence
+
+instance FromModel Sentence
+
+instance ToPrompt Sentence
+
+instance Validatable Sentence
+
+-- | A one-field answer label.
+newtype Answer = Answer {answer :: Text}
+ deriving stock (Generic, Show, Eq, Ord)
+
+instance ToSchema Answer
+
+instance FromModel Answer
+
+instance ToPrompt Answer
+
+instance Validatable Answer
+
+-- ---------------------------------------------------------------------------
+-- Signatures and programs
+-- ---------------------------------------------------------------------------
+
+-- | A deliberately bland classify signature (no @advice@/@reasoning@ words, so the
+-- stub's call-type detection stays unambiguous).
+classifySig :: Signature Sentence Answer
+classifySig = mkSignature "Classify the input into a label."
+
+-- | The single-node classify program.
+classifyProg :: Program Sentence Answer
+classifyProg = predict classifySig
+
+-- | The chain-of-thought reasoning program (produces @(reasoning, answer)@).
+reasonProg :: Program Sentence (WithReasoning Answer)
+reasonProg = chainOfThoughtRaw classifySig
+
+-- | The synthesis instruction — deliberately free of the words @reasoning@ and
+-- @advice@ so the stub recognises a synthesis call by elimination.
+synthInstruction :: Text
+synthInstruction = "Compare the student attempts and give the single consensus label."
+
+-- | The synthesis signature for 'Shikumi.Refine.multiChainComparison'.
+mccSynthSig :: Signature (MultiChainInput Sentence Answer) Answer
+mccSynthSig = multiChainSig synthInstruction
+
+-- ---------------------------------------------------------------------------
+-- Rewards
+-- ---------------------------------------------------------------------------
+
+-- | Score an answer by its label, defaulting unknown labels to @0@.
+rewardByLabel :: [(Text, Double)] -> Answer -> Double
+rewardByLabel table (Answer a) = fromMaybe 0 (lookup a table)
+
+-- | The reward used by the advice-driven test: @good@ is the only passing label.
+goodReward :: Answer -> Double
+goodReward (Answer a) = if a == "good" then 1.0 else 0.0
+
+-- ---------------------------------------------------------------------------
+-- Decision rules (a pure @Context -> Options -> Text@ response body)
+-- ---------------------------------------------------------------------------
+
+-- | The per-sample temperature carried on the request (0 when unset).
+reqTemp :: Options -> Double
+reqTemp opts = fromMaybe 0 (opts ^. #temperature)
+
+-- | The label a classify call answers, chosen by request temperature.
+labelForTemp :: Double -> Text
+labelForTemp t
+ | t < 0.5 = "low"
+ | t < 1.0 = "mid"
+ | otherwise = "high"
+
+sysPrompt :: Context -> Text
+sysPrompt ctx = fromMaybe "" (ctx ^. #systemPrompt)
+
+hasField :: Text -> Context -> Bool
+hasField name ctx = T.isInfixOf name (sysPrompt ctx)
+
+-- | M1/M4 classify rule: answer a temperature-chosen label.
+decideTemp :: Context -> Options -> Text
+decideTemp _ opts = markerBody [("answer", labelForTemp (reqTemp opts))]
+
+-- | M2 refine rule: an advice call returns fixed advice; a classify call answers
+-- @good@ when an advice hint is present, @bad@ otherwise.
+decideRefine :: Context -> Options -> Text
+decideRefine ctx _
+ | hasField "advice" ctx = markerBody [("advice", "State the keyword clearly.")]
+ | hasField "Hint from a previous attempt" ctx = markerBody [("answer", "good")]
+ | otherwise = markerBody [("answer", "bad")]
+
+-- | M3 multi-chain rule: a reasoning call answers a temperature-chosen candidate; a
+-- synthesis call returns the modal candidate it was shown.
+decideMCC :: Context -> Options -> Text
+decideMCC ctx opts
+ | hasField "reasoning" ctx =
+ markerBody [("reasoning", "considering the input"), ("value", answerJSON (candidateForTemp (reqTemp opts)))]
+ | otherwise = markerBody [("answer", modalCandidate (lastUserText ctx))]
+
+-- | The reasoning candidate for a temperature: low → @Paris@, otherwise
+-- @Brussels@ (so 3 spread samples give @[Paris, Brussels, Brussels]@, modal
+-- @Brussels@).
+candidateForTemp :: Double -> Text
+candidateForTemp t = if t < 0.5 then "Paris" else "Brussels"
+
+-- | Render an 'Answer' as the nested JSON the fallback adapter expects under a
+-- @value@ section.
+answerJSON :: Text -> Text
+answerJSON a = "{\"answer\": \"" <> a <> "\"}"
+
+-- | The modal answer among the @"my answer is ..."@ fragments in a synthesis
+-- prompt.
+modalCandidate :: Text -> Text
+modalCandidate t = case extractAnswers t of
+ [] -> ""
+ xs -> modal xs
+
+-- | Pull each attempt's answer out of the rendered 'MultiChainInput' prompt.
+extractAnswers :: Text -> [Text]
+extractAnswers t = map clean (drop 1 (T.splitOn "my answer is " t))
+ where
+ clean p =
+ let line = T.takeWhile (/= '\n') p
+ in T.strip (fromMaybe line (T.stripPrefix "answer:" (T.strip line)))
+
+-- ---------------------------------------------------------------------------
+-- Stub interpreters
+-- ---------------------------------------------------------------------------
+
+-- | Interpret @LLM@ with a pure decision rule, counting completions (for the
+-- short-circuit / budget assertions). Concurrency-safe via 'atomicModifyIORef''.
+runStub ::
+ (IOE :> es) =>
+ IORef Int ->
+ (Context -> Options -> Text) ->
+ Eff (LLM : es) a ->
+ Eff es a
+runStub ref decide = interpret $ \_ -> \case
+ Complete _ ctx opts -> do
+ liftIO (atomicModifyIORef' ref (\n -> (n + 1, ())))
+ pure (mkResponse (decide ctx opts))
+ Stream {} -> pure []
+
+-- | Like 'runStub' but records each request's last user text (for asserting what a
+-- synthesis call was shown).
+runRecordingStub ::
+ (IOE :> es) =>
+ IORef [Text] ->
+ (Context -> Options -> Text) ->
+ Eff (LLM : es) a ->
+ Eff es a
+runRecordingStub cap decide = interpret $ \_ -> \case
+ Complete _ ctx opts -> do
+ liftIO (modifyIORef' cap (++ [lastUserText ctx]))
+ pure (mkResponse (decide ctx opts))
+ Stream {} -> pure []
+
+-- | An @LLM@ interpreter that throws on every completion (the always-fails case),
+-- counting attempts.
+runThrowingLLM ::
+ (IOE :> es, Error ShikumiError :> es) =>
+ IORef Int ->
+ Eff (LLM : es) a ->
+ Eff es a
+runThrowingLLM ref = interpret $ \_ -> \case
+ Complete _ _ _ -> do
+ liftIO (atomicModifyIORef' ref (\n -> (n + 1, ())))
+ throwError (ProviderFailure "RefineStub: scripted failure")
+ Stream {} -> pure []
+
+-- ---------------------------------------------------------------------------
+-- Request readers
+-- ---------------------------------------------------------------------------
+
+-- | The text of the last user message — the actual input being processed.
+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)]
diff --git a/test/ResilienceSpec.hs b/test/ResilienceSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/ResilienceSpec.hs
@@ -0,0 +1,102 @@
+-- | The resilience interpreter's observable behaviors: transient failures are
+-- retried with backoff until success or exhaustion; non-transient failures are
+-- not retried; a budget refuses the call that would cross the ceiling; and a
+-- rate limit caps in-flight calls.
+module ResilienceSpec (tests) where
+
+import Baikai (flattenAssistantBlocks)
+import Control.Concurrent (forkIO)
+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
+import Control.Concurrent.STM (newTVarIO, readTVarIO)
+import Data.IORef (newIORef, readIORef)
+import Data.Ratio ((%))
+import Effectful (runEff)
+import Effectful.Concurrent (runConcurrent)
+import Effectful.Error.Static (runErrorNoCallStack)
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.LLM
+ ( LLMConfig (..),
+ RetryPolicy (..),
+ complete,
+ defaultLLMConfig,
+ newRateLimiter,
+ runLLMResilient,
+ )
+import Shikumi.LLM.Budget (newBudget)
+import StubProvider
+ ( concurrencyStubRegistry,
+ costStubRegistry,
+ failingStubRegistry,
+ flattenAssistantText,
+ invalidStubRegistry,
+ stubContext,
+ stubModel,
+ stubOptions,
+ )
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=))
+
+tests :: TestTree
+tests =
+ testGroup
+ "ResilienceSpec"
+ [ testCase "retry recovers after 2 failures" $ do
+ ref <- newIORef 0
+ reg <- failingStubRegistry ref 2 "ok"
+ let cfg = (defaultLLMConfig reg) {retryPolicy = RetryPolicy 3 1 5}
+ res <- runText cfg
+ res @?= Right "ok"
+ n <- readIORef ref
+ n @?= 3,
+ testCase "retry exhausts after maxAttempts" $ do
+ ref <- newIORef 0
+ reg <- failingStubRegistry ref 2 "ok"
+ let cfg = (defaultLLMConfig reg) {retryPolicy = RetryPolicy 2 1 5}
+ res <- runText cfg
+ case res of
+ Left (ProviderFailure _) -> pure ()
+ other -> assertFailure ("expected Left (ProviderFailure ...), got " <> show other)
+ n <- readIORef ref
+ n @?= 2,
+ testCase "non-transient error not retried" $ do
+ ref <- newIORef 0
+ reg <- invalidStubRegistry ref
+ let cfg = (defaultLLMConfig reg) {retryPolicy = RetryPolicy 5 1 5}
+ res <- runText cfg
+ case res of
+ Left (SchemaMismatch _) -> pure ()
+ other -> assertFailure ("expected Left (SchemaMismatch ...), got " <> show other)
+ n <- readIORef ref
+ n @?= 1,
+ testCase "budget refuses the second call" $ do
+ reg <- costStubRegistry (1 % 100) "ok"
+ b <- newBudget (Just (1 % 100))
+ let cfg = (defaultLLMConfig reg) {budget = Just b}
+ res1 <- runText cfg
+ res1 @?= Right "ok"
+ res2 <- runText cfg
+ case res2 of
+ Left (BudgetExceeded _) -> pure ()
+ other -> assertFailure ("expected Left (BudgetExceeded ...), got " <> show other),
+ testCase "rate limit caps concurrency at 1" $ do
+ cur <- newTVarIO 0
+ mx <- newTVarIO 0
+ reg <- concurrencyStubRegistry cur mx "ok"
+ rl <- newRateLimiter 1
+ let cfg = (defaultLLMConfig reg) {rateLimit = Just rl}
+ v1 <- newEmptyMVar
+ v2 <- newEmptyMVar
+ _ <- forkIO (runText cfg >>= putMVar v1)
+ _ <- forkIO (runText cfg >>= putMVar v2)
+ r1 <- takeMVar v1
+ r2 <- takeMVar v2
+ r1 @?= Right "ok"
+ r2 @?= Right "ok"
+ observedMax <- readTVarIO mx
+ assertBool ("observed max concurrency " <> show observedMax <> " should be <= 1") (observedMax <= 1)
+ ]
+ where
+ runText cfg =
+ runEff . runConcurrent . runErrorNoCallStack @ShikumiError . runLLMResilient cfg $ do
+ r <- complete stubModel stubContext stubOptions
+ pure (flattenAssistantText (flattenAssistantBlocks r))
diff --git a/test/RoutingSpec.hs b/test/RoutingSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/RoutingSpec.hs
@@ -0,0 +1,176 @@
+{-# LANGUAGE TypeApplications #-}
+
+-- | EP-14 acceptance: ambient model routing and live native structured output.
+--
+-- Every test installs a /capturing stub/ @LLM@ that records the @(Model, Options)@
+-- of each outgoing request, with the production router @routeLLM@ above it and a
+-- 'runRouting' supplying the ambient model. No network, no API key. The same
+-- assertions fail against the pre-EP-14 code (model id @""@, no @responseFormat@,
+-- identical temperatures).
+module RoutingSpec (tests) where
+
+import Baikai (Model, Options, ResponseFormat (..), _Model)
+import Baikai.Models.Generated (openai_gpt_4o_mini)
+import Control.Lens ((^.))
+import Data.Generics.Labels ()
+import Data.IORef (IORef, modifyIORef', newIORef, readIORef)
+import Data.List (nub, sort)
+import Data.Map.Strict qualified as Map
+import Data.Maybe (mapMaybe)
+import Effectful (Eff, IOE, liftIO, runEff, type (:>))
+import Effectful.Concurrent (runConcurrent)
+import Effectful.Dispatch.Dynamic (interpret)
+import Effectful.Error.Static (runErrorNoCallStack)
+import ProgramFixtures (Outline, Topic (..), outlineResponse, topicToOutline)
+import Shikumi.Adapter (metaResponseSchemaKey, metaTemperatureKey)
+import Shikumi.Combinator (majorityVote)
+import Shikumi.Error (ShikumiError)
+import Shikumi.LLM (LLM (..), Response)
+import Shikumi.Module (predict)
+import Shikumi.Program (Program, TempSchedule (..), runProgram, runProgramConc)
+import Shikumi.Routing (routeLLM, runRouting)
+import Shikumi.Schema (deriveSchema)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (assertBool, testCase, (@?=))
+
+tests :: TestTree
+tests =
+ testGroup
+ "Routing"
+ [ routesModelId,
+ nativeAttachesSchema,
+ fallbackLeavesSchemaUnset,
+ spreadSetsDistinctTemps,
+ fixedSetsExactTemps,
+ concurrentAgreesOnTemps
+ ]
+
+-- ---------------------------------------------------------------------------
+-- Capturing stub + run helpers
+-- ---------------------------------------------------------------------------
+
+-- | A base @LLM@ interpreter that records every request's @(Model, Options)@ and
+-- returns a fixed response.
+runCapturingLLM ::
+ (IOE :> es) =>
+ IORef [(Model, Options)] ->
+ Response ->
+ Eff (LLM : es) a ->
+ Eff es a
+runCapturingLLM ref resp = interpret $ \_ -> \case
+ Complete m _ o -> do
+ liftIO (modifyIORef' ref (++ [(m, o)]))
+ pure resp
+ Stream _ _ _ -> pure []
+
+-- | Run a program sequentially under @routeLLM . runRouting model@ over a capturing
+-- stub and return the captured requests.
+captureRouted :: Model -> Program Topic Outline -> Topic -> IO [(Model, Options)]
+captureRouted model prog input = do
+ ref <- newIORef []
+ res <-
+ runEff
+ . runErrorNoCallStack @ShikumiError
+ . runRouting model
+ . runCapturingLLM ref outlineResponse
+ . routeLLM
+ $ runProgram prog input
+ assertBool "routed program decodes without error" (isRight res)
+ readIORef ref
+
+-- | As 'captureRouted' but under the concurrent executor.
+captureRoutedConc :: Model -> Program Topic Outline -> Topic -> IO [(Model, Options)]
+captureRoutedConc model prog input = do
+ ref <- newIORef []
+ res <-
+ runEff
+ . runErrorNoCallStack @ShikumiError
+ . runConcurrent
+ . runRouting model
+ . runCapturingLLM ref outlineResponse
+ . routeLLM
+ $ runProgramConc prog input
+ assertBool "routed program decodes without error" (isRight res)
+ readIORef ref
+
+isRight :: Either a b -> Bool
+isRight = either (const False) (const True)
+
+-- ---------------------------------------------------------------------------
+-- M1 — the ambient model id reaches the wire
+-- ---------------------------------------------------------------------------
+
+routesModelId :: TestTree
+routesModelId =
+ testCase "routes the ambient model id onto the wire" $ do
+ captured <- captureRouted openai_gpt_4o_mini (predict topicToOutline) (Topic "cats")
+ map (^. #modelId) (map fst captured) @?= [openai_gpt_4o_mini ^. #modelId]
+
+-- ---------------------------------------------------------------------------
+-- M2 — native responseFormat on the wire; fallback leaves it unset
+-- ---------------------------------------------------------------------------
+
+nativeAttachesSchema :: TestTree
+nativeAttachesSchema =
+ testCase "native model attaches responseFormat with the derived schema" $ do
+ captured <- captureRouted openai_gpt_4o_mini (predict topicToOutline) (Topic "cats")
+ case captured of
+ [(_, o)] -> do
+ o ^. #responseFormat
+ @?= Just (JsonSchema {name = "output", schema = deriveSchema @Outline, strict = True})
+ assertBool
+ "private schema key stripped before transport"
+ (Map.notMember metaResponseSchemaKey (o ^. #metadata))
+ other -> assertBool ("expected exactly one request, got " <> show (length other)) False
+
+fallbackLeavesSchemaUnset :: TestTree
+fallbackLeavesSchemaUnset =
+ testCase "fallback model leaves responseFormat unset" $ do
+ captured <- captureRouted _Model (predict topicToOutline) (Topic "cats")
+ case captured of
+ [(_, o)] -> do
+ o ^. #responseFormat @?= Nothing
+ assertBool
+ "private schema key stripped before transport"
+ (Map.notMember metaResponseSchemaKey (o ^. #metadata))
+ other -> assertBool ("expected exactly one request, got " <> show (length other)) False
+
+-- ---------------------------------------------------------------------------
+-- M3 — per-sample temperature from MajorityVote's TempSchedule
+-- ---------------------------------------------------------------------------
+
+capturedTemps :: [(Model, Options)] -> [Double]
+capturedTemps = mapMaybe (\(_, o) -> o ^. #temperature)
+
+spreadSetsDistinctTemps :: TestTree
+spreadSetsDistinctTemps =
+ testCase "majorityVote spread sets distinct per-sample temperatures" $ do
+ captured <-
+ captureRouted openai_gpt_4o_mini (majorityVote 3 (TempSpread 0.5 0.4) (predict topicToOutline)) (Topic "cats")
+ let temps = capturedTemps captured
+ length temps @?= 3
+ assertBool "temperatures are not all equal" (length (nub temps) > 1)
+ map round3 (sort temps) @?= [0.1, 0.5, 0.9]
+ assertBool
+ "private temperature key stripped before transport"
+ (all (\(_, o) -> Map.notMember metaTemperatureKey (o ^. #metadata)) captured)
+
+fixedSetsExactTemps :: TestTree
+fixedSetsExactTemps =
+ testCase "majorityVote fixed schedule sets exactly the listed temperatures" $ do
+ captured <-
+ captureRouted openai_gpt_4o_mini (majorityVote 2 (TempFixed [0.1, 0.9]) (predict topicToOutline)) (Topic "cats")
+ sort (capturedTemps captured) @?= [0.1, 0.9]
+
+concurrentAgreesOnTemps :: TestTree
+concurrentAgreesOnTemps =
+ testCase "runProgramConc produces the same temperature multiset" $ do
+ captured <-
+ captureRoutedConc openai_gpt_4o_mini (majorityVote 3 (TempSpread 0.5 0.4) (predict topicToOutline)) (Topic "cats")
+ let temps = capturedTemps captured
+ length temps @?= 3
+ map round3 (sort temps) @?= [0.1, 0.5, 0.9]
+
+-- | Round to three decimal places to compare 'TempSpread' floating-point output.
+round3 :: Double -> Double
+round3 d = fromIntegral (round (d * 1000) :: Integer) / 1000
diff --git a/test/SchemaSpec.hs b/test/SchemaSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/SchemaSpec.hs
@@ -0,0 +1,90 @@
+{-# LANGUAGE DataKinds #-}
+
+-- | M2 (schema derivation) + M3 (total decode and the failure taxonomy) for
+-- "Shikumi.Schema". Schema equality is compared as an aeson 'Value' (key order
+-- does not matter); every failing input asserts the exact located 'ShikumiError'.
+module SchemaSpec (tests) where
+
+import Data.Aeson (Value, object, (.=))
+import Data.Text (Text)
+import Fixtures (Author (..), Sentiment (..), Summary (..))
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.Schema (deriveSchema, parseOutput)
+import Shikumi.Schema.Types (field)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (assertFailure, testCase, (@?=))
+
+-- | The schema we expect @deriveSchema \@Summary@ to produce.
+expectedSummarySchema :: Value
+expectedSummarySchema =
+ object
+ [ "type" .= s "object",
+ "properties"
+ .= object
+ [ "headline" .= object ["type" .= s "string", "description" .= s "A one-line summary"],
+ "bullets"
+ .= object
+ [ "type" .= s "array",
+ "items" .= object ["type" .= s "string"],
+ "description" .= s "Three to five key points"
+ ],
+ "author"
+ .= object
+ [ "type" .= s "object",
+ "properties"
+ .= object
+ ["name" .= object ["type" .= s "string", "description" .= s "Author full name"]],
+ "required" .= ["name" :: Text],
+ "additionalProperties" .= False
+ ],
+ "sentiment" .= object ["enum" .= (["Positive", "Neutral", "Negative"] :: [Text])],
+ "note" .= object ["anyOf" .= [object ["type" .= s "string"], object ["type" .= s "null"]]]
+ ],
+ "required" .= (["headline", "bullets", "author", "sentiment"] :: [Text]),
+ "additionalProperties" .= False
+ ]
+ where
+ s = id :: Text -> Text
+
+-- | A valid provider reply.
+goodBody :: Text
+goodBody =
+ "{\"headline\":\"H\",\"bullets\":[\"a\",\"b\",\"c\"],\
+ \\"author\":{\"name\":\"Ada\"},\"sentiment\":\"Positive\",\"note\":null}"
+
+decode :: Text -> Either ShikumiError Summary
+decode = parseOutput
+
+tests :: TestTree
+tests =
+ testGroup
+ "SchemaSpec"
+ [ testCase "deriveSchema @Summary matches the expected JSON Schema" $
+ deriveSchema @Summary @?= expectedSummarySchema,
+ testCase "good provider JSON decodes to the expected Summary" $
+ decode goodBody
+ @?= Right
+ Summary
+ { headline = field "H",
+ bullets = field ["a", "b", "c"],
+ author = Author {name = field "Ada"},
+ sentiment = Positive,
+ note = Nothing
+ },
+ testCase "missing required field -> MissingField (located)" $
+ decode "{\"bullets\":[\"a\",\"b\",\"c\"],\"author\":{\"name\":\"Ada\"},\"sentiment\":\"Positive\"}"
+ @?= Left (MissingField "headline"),
+ testCase "wrong JSON type -> SchemaMismatch (located)" $
+ decode "{\"headline\":\"H\",\"bullets\":\"oops\",\"author\":{\"name\":\"Ada\"},\"sentiment\":\"Positive\"}"
+ @?= Left (SchemaMismatch "bullets: expected array, got string"),
+ testCase "bad enum value -> SchemaMismatch listing the choices" $
+ decode "{\"headline\":\"H\",\"bullets\":[\"a\",\"b\",\"c\"],\"author\":{\"name\":\"Ada\"},\"sentiment\":\"Angry\"}"
+ @?= Left (SchemaMismatch "sentiment: expected one of [Positive, Neutral, Negative]"),
+ testCase "validation rule violation -> ValidationFailure" $
+ decode "{\"headline\":\"H\",\"bullets\":[\"a\"],\"author\":{\"name\":\"Ada\"},\"sentiment\":\"Positive\"}"
+ @?= Left (ValidationFailure "bullets: must have 3 to 5 items"),
+ testCase "non-JSON body -> InvalidJSON" $
+ case decode "not json" of
+ Left (InvalidJSON _) -> pure ()
+ other -> assertFailure ("expected Left (InvalidJSON ...), got " <> show other)
+ ]
diff --git a/test/SerializeSpec.hs b/test/SerializeSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/SerializeSpec.hs
@@ -0,0 +1,61 @@
+-- | EP-4 M3: parameter-state serialization (never closures). A program's
+-- @programParams@ vector JSON-round-trips and re-applies onto a default-parameter
+-- program of the same shape; @programShape@ is stable across parameter changes;
+-- a wrong-length vector is rejected.
+module SerializeSpec (tests) where
+
+import Data.Aeson (Value (String), decode, encode)
+import ProgramFixtures (Draft, Topic, outlineToDraft, topicToOutline)
+import Shikumi.Program
+ ( Demo (..),
+ Params (..),
+ Program (Predict),
+ ProgramShape (..),
+ ProgramShapeError (..),
+ emptyParams,
+ foldParams,
+ mapParamsAt,
+ pipeline,
+ programParams,
+ programShape,
+ setProgramParams,
+ )
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (assertFailure, testCase, (@?=))
+
+-- A fresh, default-parameter two-stage program.
+defaultProg :: Program Topic Draft
+defaultProg = pipeline (Predict topicToOutline emptyParams) (Predict outlineToDraft emptyParams)
+
+-- The same program with distinct parameters on each node.
+tunedProg :: Program Topic Draft
+tunedProg =
+ mapParamsAt 0 (const node0Params)
+ . mapParamsAt 1 (const node1Params)
+ $ defaultProg
+
+node0Params :: Params
+node0Params = Params (Just "outline carefully") [Demo (String "a topic") (String "an outline")]
+
+node1Params :: Params
+node1Params = Params (Just "draft fluently") []
+
+tests :: TestTree
+tests =
+ testGroup
+ "SerializeSpec"
+ [ testCase "programParams JSON round-trips to an equal vector" $
+ decode (encode (programParams tunedProg)) @?= Just [node0Params, node1Params],
+ testCase "setProgramParams reapplies a saved vector onto a default program" $
+ fmap foldParams (setProgramParams [node0Params, node1Params] defaultProg)
+ @?= Right (foldParams tunedProg),
+ testCase "programShape is stable across parameter changes" $
+ programShape tunedProg @?= programShape defaultProg,
+ testCase "programShape records the constructor tree and per-node labels" $
+ programShape defaultProg
+ @?= ShapeCompose (ShapePredict "points") (ShapePredict "prose"),
+ testCase "setProgramParams with a wrong-length vector is rejected" $
+ case setProgramParams [emptyParams] defaultProg of
+ Left e -> e @?= ParamCountMismatch (2, 1)
+ Right _ -> assertFailure "expected Left on a wrong-length vector"
+ ]
diff --git a/test/Shikumi/Effect/TimeSpec.hs b/test/Shikumi/Effect/TimeSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/Shikumi/Effect/TimeSpec.hs
@@ -0,0 +1,25 @@
+-- | Exercises the 'Shikumi.Effect.Time' effect against its real interpreter
+-- 'runTime': wall-clock reads succeed and the monotonic counter never runs
+-- backwards. This pins the effect's behavior so a future fake-clock interpreter
+-- can be checked against the same contract.
+module Shikumi.Effect.TimeSpec (tests) where
+
+import Effectful (runEff)
+import Shikumi.Effect.Time (getCurrentTime, getMonotonicTimeNSec, runTime)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (assertBool, testCase)
+
+tests :: TestTree
+tests =
+ testGroup
+ "Shikumi.Effect.TimeSpec"
+ [ testCase "getCurrentTime returns a renderable timestamp" $ do
+ t <- runEff (runTime getCurrentTime)
+ assertBool "timestamp renders to a non-empty string" (not (null (show t))),
+ testCase "getMonotonicTimeNSec is non-decreasing across two reads" $ do
+ (a, b) <- runEff $ runTime $ do
+ a <- getMonotonicTimeNSec
+ b <- getMonotonicTimeNSec
+ pure (a, b)
+ assertBool "second read is >= first" (b >= a)
+ ]
diff --git a/test/Shikumi/LLM/Mock.hs b/test/Shikumi/LLM/Mock.hs
new file mode 100644
--- /dev/null
+++ b/test/Shikumi/LLM/Mock.hs
@@ -0,0 +1,82 @@
+-- | A scripted, deterministic interpreter of the @LLM@ effect for the combinator
+-- tests (EP-5). Each 'Shikumi.LLM.Complete' pops the next 'MockReply': a 'MockOk'
+-- yields a canned 'Response' (build them with @ProgramFixtures.mkResponse@ /
+-- @markerBody@), a 'MockFail' throws a 'ShikumiError' through the error channel —
+-- which is exactly what 'Shikumi.Program.Retry' / 'Shikumi.Program.Validate'
+-- react to. 'runMockLLMCounting' additionally bumps a caller-supplied counter on
+-- every completion, so a retry test can assert the precise attempt count.
+--
+-- The counter is caller-owned (an 'IORef' in 'IO') rather than returned as a
+-- tuple deliberately: a 'MockFail' short-circuits the @Error ShikumiError@
+-- channel, so a tuple result would be discarded on the all-fail path — but the
+-- 'IORef', mutated before the throw, survives and is read after the run resolves.
+--
+-- This is the richer sibling of @ProgramFixtures.runScriptedLLM@ (which only
+-- scripts successes); it lives in the test tree, not the library.
+module Shikumi.LLM.Mock
+ ( MockReply (..),
+ runMockLLM,
+ runMockLLMCounting,
+ )
+where
+
+import Data.IORef (IORef, atomicModifyIORef', newIORef)
+import Effectful (Eff, IOE, liftIO, type (:>))
+import Effectful.Dispatch.Dynamic (interpret)
+import Effectful.Error.Static (Error, throwError)
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.LLM (LLM (..), Response)
+
+-- | One scripted reply to a single completion: a canned success or a typed
+-- failure.
+data MockReply
+ = MockOk Response
+ | MockFail ShikumiError
+
+-- | Interpret @LLM@ by popping scripted replies in order. A 'MockFail' is thrown
+-- through @Error ShikumiError@; the streaming op is unused (returns @[]@). An
+-- exhausted script throws a clear 'ProviderFailure' rather than blocking.
+runMockLLM ::
+ (IOE :> es, Error ShikumiError :> es) =>
+ [MockReply] ->
+ Eff (LLM : es) a ->
+ Eff es a
+runMockLLM replies act = do
+ ref <- liftIO (newIORef replies)
+ interpret
+ ( \_ -> \case
+ Complete {} -> popReply ref
+ Stream {} -> pure []
+ )
+ act
+
+-- | Like 'runMockLLM', but bumps the caller-supplied counter on every completion
+-- (counting every attempt, including those that fail). The bump happens before
+-- the reply is consumed, so a thrown 'MockFail' still counts.
+runMockLLMCounting ::
+ (IOE :> es, Error ShikumiError :> es) =>
+ IORef Int ->
+ [MockReply] ->
+ Eff (LLM : es) a ->
+ Eff es a
+runMockLLMCounting cnt replies act = do
+ ref <- liftIO (newIORef replies)
+ interpret
+ ( \_ -> \case
+ Complete {} -> do
+ liftIO (atomicModifyIORef' cnt (\n -> (n + 1, ())))
+ popReply ref
+ Stream {} -> pure []
+ )
+ act
+
+-- | Pop and act on the next reply, throwing on failure or an exhausted script.
+popReply :: (IOE :> es, Error ShikumiError :> es) => IORef [MockReply] -> Eff es Response
+popReply ref = do
+ r <- liftIO (atomicModifyIORef' ref step)
+ case r of
+ MockOk resp -> pure resp
+ MockFail e -> throwError e
+ where
+ step (x : xs) = (xs, x)
+ step [] = ([], MockFail (ProviderFailure "mock: script exhausted"))
diff --git a/test/SignatureSpec.hs b/test/SignatureSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/SignatureSpec.hs
@@ -0,0 +1,51 @@
+-- | M4: 'Shikumi.Signature'. Field metadata is derived from the record types
+-- (matching the schema), and the instruction/demos are first-class replaceable
+-- parameters (read back what you write).
+module SignatureSpec (tests) where
+
+import Fixtures (Article, Summary, sampleArticle, sampleSummary)
+import Shikumi.Schema.Types (FieldMeta (..))
+import Shikumi.Signature
+ ( Demo (..),
+ Signature,
+ getDemos,
+ getInstruction,
+ inputFields,
+ mkSignature,
+ outputFields,
+ setDemos,
+ setInstruction,
+ )
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (testCase, (@?=))
+
+sig0 :: Signature Article Summary
+sig0 = mkSignature "Summarize the article"
+
+tests :: TestTree
+tests =
+ testGroup
+ "SignatureSpec"
+ [ testCase "instruction is the one supplied" $
+ getInstruction sig0 @?= "Summarize the article",
+ testCase "demos start empty" $
+ getDemos sig0 @?= [],
+ testCase "input field metadata is derived from the input record" $
+ inputFields sig0
+ @?= [ FieldMeta "title" (Just "The article's headline"),
+ FieldMeta "body" (Just "The full article text")
+ ],
+ testCase "output field metadata matches the output record (descriptions and all)" $
+ outputFields sig0
+ @?= [ FieldMeta "headline" (Just "A one-line summary"),
+ FieldMeta "bullets" (Just "Three to five key points"),
+ FieldMeta "author" Nothing,
+ FieldMeta "sentiment" Nothing,
+ FieldMeta "note" Nothing
+ ],
+ testCase "instruction is replaceable" $
+ getInstruction (setInstruction "New instruction" sig0) @?= "New instruction",
+ testCase "demos are replaceable" $
+ getDemos (setDemos [Demo sampleArticle sampleSummary] sig0)
+ @?= [Demo sampleArticle sampleSummary]
+ ]
diff --git a/test/StreamSpec.hs b/test/StreamSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/StreamSpec.hs
@@ -0,0 +1,220 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+-- | EP-25: program-level streaming. Verified entirely hermetically against a
+-- scripted streaming 'LLM' interpreter (no network, no routing). The decisive
+-- proof of fidelity: in M2/M3 the streamed return value is asserted equal to what
+-- the blocking 'runProgram' produces for the same scripted events — so streaming is
+-- an additive parallel path over the same decode, and 'runProgram'\'s contract is
+-- untouched.
+module StreamSpec (tests) where
+
+import Baikai
+ ( AssistantContent (..),
+ AssistantMessageEvent (..),
+ BlockEndPayload (..),
+ DeltaPayload (..),
+ IndexPayload (..),
+ Message (..),
+ Response,
+ StartPayload (..),
+ StopReason (..),
+ TerminalPayload (..),
+ _Context,
+ _Model,
+ _Options,
+ _Response,
+ _TextContent,
+ )
+import Control.Lens ((&), (.~), (^.))
+import Data.Generics.Labels ()
+import Data.IORef (IORef, atomicModifyIORef', modifyIORef', newIORef, readIORef)
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Vector qualified as V
+import Effectful (Eff, IOE, liftIO, runEff, type (:>))
+import Effectful.Dispatch.Dynamic (interpret)
+import Effectful.Error.Static (runErrorNoCallStack)
+import GHC.Generics (Generic)
+import ProgramFixtures (Draft (..), Topic (..), markerBody, outlineToDraft, topicToOutline)
+import Shikumi.Adapter (ToPrompt, responseText)
+import Shikumi.Combinator ((>>>))
+import Shikumi.Error (ShikumiError)
+import Shikumi.LLM (LLM (..))
+import Shikumi.Module (predict)
+import Shikumi.Program (Program, runProgram)
+import Shikumi.Schema (FromModel, ToSchema, Validatable)
+import Shikumi.Signature (Signature, mkSignature)
+import Shikumi.Stream
+ ( FieldChunk (..),
+ Status (..),
+ StatusPhase (..),
+ StreamEvent (..),
+ streamComplete,
+ streamProgram,
+ )
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (testCase, (@?=))
+
+-- A single-text-field signature, the honest field-chunk case.
+newtype Question = Question {question :: Text}
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema Question
+
+instance FromModel Question
+
+instance ToPrompt Question
+
+instance Validatable Question
+
+newtype Answer = Answer {answer :: Text}
+ deriving stock (Generic, Show, Eq)
+
+instance ToSchema Answer
+
+instance FromModel Answer
+
+instance ToPrompt Answer
+
+instance Validatable Answer
+
+qToAnswer :: Signature Question Answer
+qToAnswer = mkSignature "Answer the question"
+
+-- ---------------------------------------------------------------------------
+-- Scripted streaming LLM: one event-list per call. Stream returns it; Complete
+-- derives a Response from its terminal event, so the SAME script drives both the
+-- streamed and the blocking run (and the equality assertion is meaningful).
+-- ---------------------------------------------------------------------------
+
+runStreamingLLM :: (IOE :> es) => IORef [[AssistantMessageEvent]] -> Eff (LLM : es) a -> Eff es a
+runStreamingLLM ref = interpret $ \_ -> \case
+ Complete _ _ _ -> liftIO (terminalResponse <$> popEvents ref)
+ Stream _ _ _ -> liftIO (popEvents ref)
+
+popEvents :: IORef [[AssistantMessageEvent]] -> IO [AssistantMessageEvent]
+popEvents ref = atomicModifyIORef' ref step
+ where
+ step (x : xs) = (xs, x)
+ step [] = ([], [])
+
+-- | A response carrying @t@ as its single assistant text block.
+mkResponse :: Text -> Response
+mkResponse t =
+ _Response & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t))
+
+-- | The 'Response' assembled from a stream's terminal event.
+terminalResponse :: [AssistantMessageEvent] -> Response
+terminalResponse evs =
+ case [p | EventDone (TerminalPayload _ (AssistantMessage p)) <- evs] of
+ (p : _) -> _Response & #message .~ p
+ [] -> mkResponse ""
+
+-- | A valid event sequence: @deltas@ stream as text chunks, and the terminal event
+-- carries @terminalText@ as the fully-assembled message (the structured body the
+-- fallback adapter parses). The two are decoupled here as a fake: the deltas model
+-- the value streaming incrementally, the terminal carries the whole structured reply.
+streamEventsFor :: [Text] -> Text -> [AssistantMessageEvent]
+streamEventsFor deltas terminalText =
+ [ EventStart (StartPayload (AssistantMessage (_Response ^. #message))),
+ TextStart (IndexPayload 0)
+ ]
+ ++ [TextDelta (DeltaPayload 0 d) | d <- deltas]
+ ++ [ TextEnd (BlockEndPayload 0 (T.concat deltas)),
+ EventDone (TerminalPayload Stop (AssistantMessage (payloadWith terminalText)))
+ ]
+ where
+ payloadWith t = (_Response ^. #message) & #content .~ V.singleton (AssistantText (_TextContent & #text .~ t))
+
+-- | A recording callback: append each event to the IORef in order.
+recorder :: (IOE :> es) => IORef [StreamEvent] -> StreamEvent -> Eff es ()
+recorder rec ev = liftIO (modifyIORef' rec (++ [ev]))
+
+-- ---------------------------------------------------------------------------
+-- Tests
+-- ---------------------------------------------------------------------------
+
+tests :: TestTree
+tests =
+ testGroup
+ "StreamSpec"
+ [ testCase "M1: streamComplete folds deltas into field chunks and assembles the response" $ do
+ ref <- newIORef [streamEventsFor ["Hel", "lo"] "Hello"]
+ rec <- newIORef []
+ resp <-
+ runEff . runStreamingLLM ref $
+ streamComplete "answer" _Model _Context _Options (recorder rec)
+ evs <- readIORef rec
+ evs
+ @?= [ StreamFieldChunk (FieldChunk "answer" "Hel" False),
+ StreamFieldChunk (FieldChunk "answer" "lo" True)
+ ]
+ responseText resp @?= "Hello",
+ testCase "M2: streamProgram streams a single Predict's field, bracketed by LM status" $ do
+ ref <- newIORef [streamEventsFor ["Hel", "lo"] (markerBody [("answer", "Hello")])]
+ rec <- newIORef []
+ out <-
+ runEff . runErrorNoCallStack @ShikumiError . runStreamingLLM ref $
+ streamProgram (predict qToAnswer) (Question "q?") (recorder rec)
+ evs <- readIORef rec
+ evs
+ @?= [ StreamStatus (Status LmStart "LM call started"),
+ StreamFieldChunk (FieldChunk "answer" "Hel" False),
+ StreamFieldChunk (FieldChunk "answer" "lo" True),
+ StreamStatus (Status LmEnd "LM call finished")
+ ]
+ out @?= Right (Answer "Hello"),
+ testCase "M2: streamProgram returns the same value as runProgram" $ do
+ let script = [streamEventsFor ["Hel", "lo"] (markerBody [("answer", "Hello")])]
+ ref1 <- newIORef script
+ streamed <-
+ runEff . runErrorNoCallStack @ShikumiError . runStreamingLLM ref1 $
+ streamProgram (predict qToAnswer) (Question "q?") (\_ -> pure ())
+ ref2 <- newIORef script
+ blocking <-
+ runEff . runErrorNoCallStack @ShikumiError . runStreamingLLM ref2 $
+ runProgram (predict qToAnswer) (Question "q?")
+ streamed @?= blocking,
+ testCase "M3: a chain of predicts brackets each LM with status and streams both fields" $ do
+ let script =
+ [ streamEventsFor ["pts"] (markerBody [("points", "[\"intro\", \"body\", \"end\"]")]),
+ streamEventsFor ["txt"] (markerBody [("prose", "A fine essay about the outline.")])
+ ]
+ chain :: Program Topic Draft
+ chain = predict topicToOutline >>> predict outlineToDraft
+ ref <- newIORef script
+ rec <- newIORef []
+ out <-
+ runEff . runErrorNoCallStack @ShikumiError . runStreamingLLM ref $
+ streamProgram chain (Topic "haskell") (recorder rec)
+ evs <- readIORef rec
+ evs
+ @?= [ StreamStatus (Status NodeStart "node started"),
+ StreamStatus (Status LmStart "LM call started"),
+ StreamFieldChunk (FieldChunk "points" "pts" True),
+ StreamStatus (Status LmEnd "LM call finished"),
+ StreamStatus (Status LmStart "LM call started"),
+ StreamFieldChunk (FieldChunk "prose" "txt" True),
+ StreamStatus (Status LmEnd "LM call finished"),
+ StreamStatus (Status NodeEnd "node finished")
+ ]
+ out @?= Right (Draft {prose = "A fine essay about the outline."})
+
+ -- and the streamed value equals the blocking run for the same script
+ ref2 <- newIORef script
+ blocking <-
+ runEff . runErrorNoCallStack @ShikumiError . runStreamingLLM ref2 $
+ runProgram chain (Topic "haskell")
+ Right (Draft {prose = "A fine essay about the outline."}) @?= blocking,
+ testCase "M3: ToolStart/ToolEnd status events flow through the callback in order" $ do
+ rec <- newIORef []
+ runEff $ do
+ let cb = recorder rec
+ cb (StreamStatus (Status ToolStart "tool started"))
+ cb (StreamStatus (Status ToolEnd "tool finished"))
+ evs <- readIORef rec
+ evs
+ @?= [ StreamStatus (Status ToolStart "tool started"),
+ StreamStatus (Status ToolEnd "tool finished")
+ ]
+ ]
diff --git a/test/StubProvider.hs b/test/StubProvider.hs
new file mode 100644
--- /dev/null
+++ b/test/StubProvider.hs
@@ -0,0 +1,157 @@
+-- | Hand-rolled, network-free baikai providers for the hermetic shikumi tests.
+--
+-- Every helper builds an isolated 'ProviderRegistry' (never the global one) whose
+-- single 'ApiProvider' serves a known 'Custom' 'Api' tag, so the default
+-- @cabal test all@ needs no network and no API key. The variants differ only in
+-- what @complete@ does: return fixed text, return text with a fixed dollar cost,
+-- fail a fixed number of times then succeed, always fail with a non-transient
+-- error, or instrument observed concurrency.
+module StubProvider
+ ( stubModel,
+ stubContext,
+ stubOptions,
+ flattenAssistantText,
+ stubRegistry,
+ costStubRegistry,
+ failingStubRegistry,
+ invalidStubRegistry,
+ concurrencyStubRegistry,
+ )
+where
+
+import Baikai
+import Baikai.Prelude
+import Control.Concurrent (threadDelay)
+import Control.Concurrent.STM
+ ( TVar,
+ atomically,
+ modifyTVar',
+ readTVar,
+ )
+import Control.Exception (bracket_, throwIO)
+import Data.IORef (IORef, atomicModifyIORef')
+import Data.Text qualified as T
+import Data.Vector qualified as V
+import Streamly.Data.Stream qualified as Stream
+
+-- | Concatenate the text of every 'AssistantText' block, ignoring thinking and
+-- tool-call blocks. baikai has no library equivalent (its own smoke tests define
+-- this same helper locally).
+flattenAssistantText :: V.Vector AssistantContent -> Text
+flattenAssistantText = T.concat . V.toList . V.mapMaybe textOf
+ where
+ textOf (AssistantText (TextContent t)) = Just t
+ textOf _ = Nothing
+
+-- | The 'Api' tag every stub serves. A 'Custom' tag needs no catalog entry.
+stubApi :: Api
+stubApi = Custom "stub"
+
+-- | A hand-built 'Model' whose 'api' routes to the stub provider.
+stubModel :: Model
+stubModel =
+ _Model
+ & #api
+ .~ stubApi
+ & #modelId
+ .~ "stub-model"
+ & #name
+ .~ "stub"
+ & #provider
+ .~ "stub"
+
+-- | A minimal request context (one user turn).
+stubContext :: Context
+stubContext = _Context & #messages .~ V.singleton (user "ping")
+
+-- | Default request options.
+stubOptions :: Options
+stubOptions = _Options
+
+-- | An assistant payload carrying the given text as its single text block.
+stubPayloadWith :: Text -> AssistantPayload
+stubPayloadWith t =
+ (_Response ^. #message)
+ & #content
+ .~ V.singleton (AssistantText (_TextContent & #text .~ t))
+
+-- | A fixed assistant response carrying the given text as its single text block.
+stubResponse :: Text -> Response
+stubResponse t =
+ _Response
+ & #message
+ .~ stubPayloadWith t
+ & #model
+ .~ stubModel
+ & #api
+ .~ stubApi
+ & #provider
+ .~ "stub"
+
+-- | A deterministic, valid event sequence for the given text.
+stubEvents :: Text -> [AssistantMessageEvent]
+stubEvents t =
+ [ EventStart StartPayload {partial = AssistantMessage (_Response ^. #message)},
+ TextStart IndexPayload {contentIndex = 0},
+ TextDelta DeltaPayload {contentIndex = 0, delta = t},
+ TextEnd BlockEndPayload {contentIndex = 0, content = t},
+ EventDone TerminalPayload {reason = Stop, message = AssistantMessage (stubPayloadWith t)}
+ ]
+
+-- | Build an isolated registry from a @complete@ implementation, reusing the
+-- deterministic stub stream for the streaming path.
+mkRegistry :: Text -> (Model -> Context -> Options -> IO Response) -> IO ProviderRegistry
+mkRegistry streamText completeFn = do
+ reg <- newProviderRegistry
+ registerApiProviderWith
+ reg
+ ApiProvider
+ { apiTag = stubApi,
+ complete = completeFn,
+ stream = \_ _ _ -> Stream.fromList (stubEvents streamText)
+ }
+ pure reg
+
+-- | A registry whose @complete@ returns the given text.
+stubRegistry :: Text -> IO ProviderRegistry
+stubRegistry t = mkRegistry t (\_ _ _ -> pure (stubResponse t))
+
+-- | A registry whose @complete@ returns the given text with a fixed dollar cost
+-- recorded in @Usage.cost.usd@ (so the budget interpreter charges it).
+costStubRegistry :: Rational -> Text -> IO ProviderRegistry
+costStubRegistry c t =
+ mkRegistry t (\_ _ _ -> pure (stubResponse t & #message . #usage . #cost . #usd .~ c))
+
+-- | A registry whose @complete@ throws 'ProviderError' (a transient failure) the
+-- first @failTimes@ calls, then returns the given text. The 'IORef' records the
+-- total number of attempts (used by the retry tests).
+failingStubRegistry :: IORef Int -> Int -> Text -> IO ProviderRegistry
+failingStubRegistry ref failTimes t =
+ mkRegistry t $ \_ _ _ -> do
+ n <- atomicModifyIORef' ref (\k -> (k + 1, k + 1))
+ if n <= failTimes
+ then throwIO (ProviderError ("stub transient failure #" <> T.pack (show n)))
+ else pure (stubResponse t)
+
+-- | A registry whose @complete@ always throws 'RequestInvalid' (a non-transient
+-- failure that maps to 'Shikumi.Error.SchemaMismatch'). The 'IORef' records the
+-- number of attempts (a retry must not increase it).
+invalidStubRegistry :: IORef Int -> IO ProviderRegistry
+invalidStubRegistry ref =
+ mkRegistry "" $ \_ _ _ -> do
+ _ <- atomicModifyIORef' ref (\k -> (k + 1, k + 1))
+ throwIO (RequestInvalid "stub bad request")
+
+-- | A registry whose @complete@ records observed concurrency: it bumps @cur@ on
+-- entry (updating the running maximum in @mx@), sleeps briefly, and decrements on
+-- exit. Used to assert that a rate limit caps in-flight calls.
+concurrencyStubRegistry :: TVar Int -> TVar Int -> Text -> IO ProviderRegistry
+concurrencyStubRegistry cur mx t =
+ mkRegistry t $ \_ _ _ ->
+ bracket_ enter leave (threadDelay 30000 >> pure (stubResponse t))
+ where
+ enter = atomically $ do
+ modifyTVar' cur (+ 1)
+ c <- readTVar cur
+ modifyTVar' mx (max c)
+ leave = atomically (modifyTVar' cur (subtract 1))
diff --git a/test/TwoStepSpec.hs b/test/TwoStepSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/TwoStepSpec.hs
@@ -0,0 +1,97 @@
+-- | EP-26 M3: the 'Shikumi.Module.twoStep' combinator. It issues two model calls
+-- — a free-form prose answer, then a structured extraction — and returns a typed
+-- @o@. Driven against a stub LLM scripted with prose then a marker reply, it must
+-- return the expected 'Summary' and consume both scripted responses (the
+-- observable that distinguishes it from a single-call adapter).
+module TwoStepSpec (tests) where
+
+import Baikai
+ ( AssistantContent (..),
+ Response,
+ _Response,
+ _TextContent,
+ )
+import Control.Lens ((&), (.~))
+import Data.Generics.Labels ()
+import Data.IORef (newIORef, readIORef)
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Vector qualified as V
+import Effectful (runEff)
+import Effectful.Error.Static (runErrorNoCallStack)
+import Fixtures (Article, Author (..), Sentiment (..), Summary (..), sampleArticle, sampleSummary)
+import ProgramFixtures (runScriptedLLM)
+import Shikumi.Error (ShikumiError)
+import Shikumi.Module (twoStep)
+import Shikumi.Program (Program, runProgram)
+import Shikumi.Schema.Types (field)
+import Shikumi.Signature (Demo (..), Signature, mkSignature, setDemos)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (testCase, (@?=))
+
+sig :: Signature Article Summary
+sig = setDemos [Demo sampleArticle sampleSummary] (mkSignature "Summarize the article")
+
+prog :: Program Article Summary
+prog = twoStep sig
+
+-- | Build a response whose single assistant text block carries the given body.
+mkResponse :: Text -> Response
+mkResponse t =
+ _Response & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t))
+
+-- | The first scripted call's reply: free-form prose (no structure requested).
+freeFormResponse :: Response
+freeFormResponse =
+ mkResponse
+ "The headline is 'Shikumi types LM programs'. Three points: records in, \
+ \records out, errors are typed. Author Ada. Sentiment positive."
+
+-- | The second scripted call's reply: a structured [[ ## field ## ]] marker body,
+-- the extraction target the fallback adapter parses.
+structuredResponse :: Response
+structuredResponse =
+ mkResponse $
+ T.intercalate
+ "\n"
+ [ "[[ ## headline ## ]]",
+ "Shikumi types LM programs",
+ "[[ ## bullets ## ]]",
+ "[\"records in\", \"records out\", \"errors are typed\"]",
+ "[[ ## author ## ]]",
+ "{\"name\": \"Ada\"}",
+ "[[ ## sentiment ## ]]",
+ "Positive",
+ "[[ ## note ## ]]",
+ "null",
+ "[[ ## completed ## ]]"
+ ]
+
+expectedSummary :: Summary
+expectedSummary =
+ Summary
+ { headline = field "Shikumi types LM programs",
+ bullets = field ["records in", "records out", "errors are typed"],
+ author = Author {name = field "Ada"},
+ sentiment = Positive,
+ note = Nothing
+ }
+
+tests :: TestTree
+tests =
+ testGroup
+ "TwoStepSpec"
+ [ testCase "twoStep: free-form then extract yields the expected Summary" $ do
+ ref <- newIORef [freeFormResponse, structuredResponse]
+ out <-
+ runEff . runErrorNoCallStack @ShikumiError . runScriptedLLM ref $
+ runProgram prog sampleArticle
+ out @?= Right expectedSummary,
+ testCase "twoStep: both scripted calls are consumed (two model calls)" $ do
+ ref <- newIORef [freeFormResponse, structuredResponse]
+ _ <-
+ runEff . runErrorNoCallStack @ShikumiError . runScriptedLLM ref $
+ runProgram prog sampleArticle
+ remaining <- readIORef ref
+ length remaining @?= 0
+ ]
diff --git a/test/XmlAdapterSpec.hs b/test/XmlAdapterSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/XmlAdapterSpec.hs
@@ -0,0 +1,100 @@
+-- | EP-26 M1: 'Shikumi.Adapter.xmlAdapter'. A third wire format on the typed seam.
+-- @render@ asks for @…@ tags; @parse@ reads them back into the
+-- typed output via the same 'sectionsToObject' + decode path the fallback adapter
+-- uses, so a missing tag yields the same located 'MissingField'.
+module XmlAdapterSpec (tests) where
+
+import Baikai
+ ( AssistantContent (..),
+ Response,
+ _Response,
+ _TextContent,
+ )
+import Control.Lens ((&), (.~), (^.))
+import Data.Generics.Labels ()
+import Data.Maybe (fromMaybe)
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Vector qualified as V
+import Fixtures (Article, Author (..), Sentiment (..), Summary (..), sampleArticle, sampleSummary)
+import Shikumi.Adapter (Adapter (..), xmlAdapter)
+import Shikumi.Error (ShikumiError (..))
+import Shikumi.Schema.Types (field)
+import Shikumi.Signature (Demo (..), Signature, mkSignature, setDemos)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (testCase, (@?=))
+
+sig :: Signature Article Summary
+sig = setDemos [Demo sampleArticle sampleSummary] (mkSignature "Summarize the article")
+
+-- | Build a response whose single assistant text block carries the given body.
+mkResponse :: Text -> Response
+mkResponse t =
+ _Response & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t))
+
+-- | A hand-written XML reply wrapping each 'Summary' field in tags.
+xmlBody :: Text
+xmlBody =
+ T.intercalate
+ "\n"
+ [ "",
+ "Shikumi types LM programs",
+ "",
+ "",
+ "[\"records in\", \"records out\", \"errors are typed\"]",
+ "",
+ "",
+ "{\"name\": \"Ada\"}",
+ "",
+ "",
+ "Positive",
+ "",
+ "",
+ "null",
+ ""
+ ]
+
+-- | The same body with the @@ tag omitted (required field missing).
+xmlBodyNoBullets :: Text
+xmlBodyNoBullets =
+ T.intercalate
+ "\n"
+ [ "",
+ "Shikumi types LM programs",
+ "",
+ "",
+ "{\"name\": \"Ada\"}",
+ "",
+ "",
+ "Positive",
+ "",
+ "",
+ "null",
+ ""
+ ]
+
+expectedSummary :: Summary
+expectedSummary =
+ Summary
+ { headline = field "Shikumi types LM programs",
+ bullets = field ["records in", "records out", "errors are typed"],
+ author = Author {name = field "Ada"},
+ sentiment = Positive,
+ note = Nothing
+ }
+
+sysOf :: Adapter Article Summary -> Text
+sysOf adapter = fromMaybe "" (fst (render adapter sig sampleArticle) ^. #systemPrompt)
+
+tests :: TestTree
+tests =
+ testGroup
+ "XmlAdapterSpec"
+ [ testCase "xml render: system prompt has the instruction and an XML tag" $ do
+ T.isInfixOf "Summarize the article" (sysOf xmlAdapter) @?= True
+ T.isInfixOf "" (sysOf xmlAdapter) @?= True,
+ testCase "xml parse: tagged body decodes to the expected Summary" $
+ parse xmlAdapter sig (mkResponse xmlBody) @?= Right expectedSummary,
+ testCase "xml parse: a missing tag -> MissingField (located)" $
+ parse xmlAdapter sig (mkResponse xmlBodyNoBullets) @?= Left (MissingField "bullets")
+ ]