shikumi (empty) → 0.1.0.0
raw patch · 48 files changed
+7160/−0 lines, 48 filesdep +QuickCheckdep +aesondep +baikai
Dependencies added: QuickCheck, aeson, baikai, baikai-claude, baikai-effectful, baikai-openai, base, base64-bytestring, bytestring, containers, directory, effectful, filepath, generic-lens, lens, scientific, shikumi, stm, streamly-core, tasty, tasty-hunit, tasty-quickcheck, text, time, vector
Files
- CHANGELOG.md +10/−0
- shikumi.cabal +131/−0
- src/Shikumi/Adapter.hs +471/−0
- src/Shikumi/Combinator.hs +169/−0
- src/Shikumi/Effect/Time.hs +52/−0
- src/Shikumi/Error.hs +62/−0
- src/Shikumi/LLM.hs +293/−0
- src/Shikumi/LLM/Budget.hs +55/−0
- src/Shikumi/Module.hs +216/−0
- src/Shikumi/Multimodal.hs +182/−0
- src/Shikumi/Prelude.hs +10/−0
- src/Shikumi/Program.hs +653/−0
- src/Shikumi/Refine.hs +385/−0
- src/Shikumi/Reward.hs +37/−0
- src/Shikumi/Routing.hs +113/−0
- src/Shikumi/Schema.hs +458/−0
- src/Shikumi/Schema/Types.hs +192/−0
- src/Shikumi/Signature.hs +79/−0
- src/Shikumi/Stream.hs +288/−0
- test/AdapterSpec.hs +115/−0
- test/CombinatorSpec.hs +330/−0
- test/ConstraintSpec.hs +90/−0
- test/EndToEndSpec.hs +106/−0
- test/ErrorSpec.hs +35/−0
- test/Fixtures.hs +88/−0
- test/LLMSpec.hs +42/−0
- test/LiveSpec.hs +48/−0
- test/Main.hs +58/−0
- test/ModuleSpec.hs +63/−0
- test/MultimodalAdapterSpec.hs +95/−0
- test/MultimodalEndToEndSpec.hs +131/−0
- test/MultimodalSpec.hs +46/−0
- test/ProgramAcceptanceSpec.hs +105/−0
- test/ProgramFixtures.hs +177/−0
- test/ProgramSpec.hs +84/−0
- test/RefineSpec.hs +262/−0
- test/RefineStub.hs +268/−0
- test/ResilienceSpec.hs +102/−0
- test/RoutingSpec.hs +176/−0
- test/SchemaSpec.hs +90/−0
- test/SerializeSpec.hs +61/−0
- test/Shikumi/Effect/TimeSpec.hs +25/−0
- test/Shikumi/LLM/Mock.hs +82/−0
- test/SignatureSpec.hs +51/−0
- test/StreamSpec.hs +220/−0
- test/StubProvider.hs +157/−0
- test/TwoStepSpec.hs +97/−0
- test/XmlAdapterSpec.hs +100/−0
+ CHANGELOG.md view
@@ -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.
+ shikumi.cabal view
@@ -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
+ src/Shikumi/Adapter.hs view
@@ -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 @\<field\>…\</field\>@ 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 @\<field\>…\</field\>@ 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, @\<name\>@.+openTag :: Text -> Text+openTag name = "<" <> name <> ">"++-- | An XML close tag, @\</name\>@.+closeTag :: Text -> Text+closeTag name = "</" <> name <> ">"++-- | Render a demo output as @\<field\>…\</field\>@ 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 @\<name\>…\</name\>@ 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+-- @\<name\>@ and its next @\</name\>@ — a non-greedy match, the same as DSPy's+-- @\<(?P\<name\>\\w+)\>(?P\<content\>.*?)\</\\1\>@ 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 @\<name\>@ and its next @\</name\>@, 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
+ src/Shikumi/Combinator.hs view
@@ -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
+ src/Shikumi/Effect/Time.hs view
@@ -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
+ src/Shikumi/Error.hs view
@@ -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
+ src/Shikumi/LLM.hs view
@@ -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
+ src/Shikumi/LLM/Budget.hs view
@@ -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)
+ src/Shikumi/Module.hs view
@@ -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 : <o's schema>}@.+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+ }
+ src/Shikumi/Multimodal.hs view
@@ -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 _ = []
+ src/Shikumi/Prelude.hs view
@@ -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)
+ src/Shikumi/Program.hs view
@@ -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'')
+ src/Shikumi/Refine.hs view
@@ -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 <reasoning>; my answer is <answer>"@ (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})
+ src/Shikumi/Reward.hs view
@@ -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)
+ src/Shikumi/Routing.hs view
@@ -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
+ src/Shikumi/Schema.hs view
@@ -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"
+ src/Shikumi/Schema/Types.hs view
@@ -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":<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
+ src/Shikumi/Signature.hs view
@@ -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}
+ src/Shikumi/Stream.hs view
@@ -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
+ test/AdapterSpec.hs view
@@ -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+ ]
+ test/CombinatorSpec.hs view
@@ -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+ ]
+ test/ConstraintSpec.hs view
@@ -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.<field>@ 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)))+ ]+ )
+ test/EndToEndSpec.hs view
@@ -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")+ ]
+ test/ErrorSpec.hs view
@@ -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+ ]
+ test/Fixtures.hs view
@@ -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+ }
+ test/LLMSpec.hs view
@@ -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)+ ]
+ test/LiveSpec.hs view
@@ -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"+ ]
+ test/Main.hs view
@@ -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+ ]
+ test/ModuleSpec.hs view
@@ -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"}]+ ]
+ test/MultimodalAdapterSpec.hs view
@@ -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]+ ]
+ test/MultimodalEndToEndSpec.hs view
@@ -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]+ ]
+ test/MultimodalSpec.hs view
@@ -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+ ]
+ test/ProgramAcceptanceSpec.hs view
@@ -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 []
+ test/ProgramFixtures.hs view
@@ -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 []))
+ test/ProgramSpec.hs view
@@ -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)+ ]
+ test/RefineSpec.hs view
@@ -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)
+ test/RefineStub.hs view
@@ -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)]
+ test/ResilienceSpec.hs view
@@ -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))
+ test/RoutingSpec.hs view
@@ -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
+ test/SchemaSpec.hs view
@@ -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)+ ]
+ test/SerializeSpec.hs view
@@ -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"+ ]
+ test/Shikumi/Effect/TimeSpec.hs view
@@ -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)+ ]
+ test/Shikumi/LLM/Mock.hs view
@@ -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"))
+ test/SignatureSpec.hs view
@@ -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]+ ]
+ test/StreamSpec.hs view
@@ -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")+ ]+ ]
+ test/StubProvider.hs view
@@ -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))
+ test/TwoStepSpec.hs view
@@ -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+ ]
+ test/XmlAdapterSpec.hs view
@@ -0,0 +1,100 @@+-- | EP-26 M1: 'Shikumi.Adapter.xmlAdapter'. A third wire format on the typed seam.+-- @render@ asks for @<field>…</field>@ 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"+ [ "<headline>",+ "Shikumi types LM programs",+ "</headline>",+ "<bullets>",+ "[\"records in\", \"records out\", \"errors are typed\"]",+ "</bullets>",+ "<author>",+ "{\"name\": \"Ada\"}",+ "</author>",+ "<sentiment>",+ "Positive",+ "</sentiment>",+ "<note>",+ "null",+ "</note>"+ ]++-- | The same body with the @<bullets>@ tag omitted (required field missing).+xmlBodyNoBullets :: Text+xmlBodyNoBullets =+ T.intercalate+ "\n"+ [ "<headline>",+ "Shikumi types LM programs",+ "</headline>",+ "<author>",+ "{\"name\": \"Ada\"}",+ "</author>",+ "<sentiment>",+ "Positive",+ "</sentiment>",+ "<note>",+ "null",+ "</note>"+ ]++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 "<headline>" (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")+ ]