diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,10 @@ +# Changelog + +## Unreleased + +## 0.1.0.0 - 2026-06-13 + +### Added + +- Initial Hackage release of hierarchical tracing and deterministic replay for shikumi. +- Trace spans, trace trees, program-node tracing, replay, trace storage, feedback logs, rendering, and a trace demo executable. diff --git a/app/Main.hs b/app/Main.hs new file mode 100644 --- /dev/null +++ b/app/Main.hs @@ -0,0 +1,14 @@ +-- | The @shikumi-trace-demo@ executable (EP-7, M5). Delegates to the testable +-- core 'Shikumi.Trace.Demo.demoMain': +-- +-- @ +-- cabal run shikumi-trace-demo -- live: prints the tree, writes trace.json +-- SHIKUMI_OFFLINE=1 cabal run shikumi-trace-demo -- --replay trace.json +-- @ +module Main (main) where + +import Shikumi.Trace.Demo (demoMain) +import System.Environment (getArgs) + +main :: IO () +main = getArgs >>= demoMain diff --git a/shikumi-trace.cabal b/shikumi-trace.cabal new file mode 100644 --- /dev/null +++ b/shikumi-trace.cabal @@ -0,0 +1,104 @@ +cabal-version: 3.4 +name: shikumi-trace +version: 0.1.0.0 +synopsis: + Hierarchical tracing, observability, and deterministic replay for shikumi (EP-7) + +category: AI +description: + Hierarchical (nested) tracing and deterministic offline replay for shikumi. + Running a program inside the @Trace@ effect produces a /tree/ of spans (program + / module / combinator / LM-call), captured by interposing on EP-1's @LLM@ + effect — so each provider call becomes a leaf span carrying model, latency, + tokens, cost, and the EP-6 content-addressed cache key. The tree pretty-prints + to the terminal and serializes to a stable JSON trace file; @runLLMReplay@ then + re-runs the same program with the network disabled, serving every call from the + recording (and raising a precise @ReplayDivergence@ for any unrecorded call). + +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.Trace + Shikumi.Trace.Demo + Shikumi.Trace.Feedback + Shikumi.Trace.Internal.Spike + Shikumi.Trace.Node + Shikumi.Trace.Program + Shikumi.Trace.Replay + Shikumi.Trace.ResponseJSON + Shikumi.Trace.Store + + build-depends: + , aeson + , baikai + , base >=4.20 && <5 + , bytestring + , containers + , directory + , effectful + , generic-lens + , lens ^>=5.3 + , scientific + , shikumi ^>=0.1.0.0 + , shikumi-cache ^>=0.1.0.0 + , text ^>=2.1 + , time + , vector + +executable shikumi-trace-demo + import: common-options + hs-source-dirs: app + main-is: Main.hs + ghc-options: -threaded -with-rtsopts=-N + build-depends: + , base + , shikumi-trace ^>=0.1.0.0 + +test-suite shikumi-trace-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: TraceFixtures + build-depends: + , aeson + , baikai + , base + , bytestring + , containers + , effectful + , generic-lens + , lens ^>=5.3 + , QuickCheck + , shikumi ^>=0.1.0.0 + , shikumi-cache ^>=0.1.0.0 + , shikumi-trace ^>=0.1.0.0 + , tasty + , tasty-hunit + , tasty-quickcheck + , temporary + , text ^>=2.1 + , time + , vector diff --git a/src/Shikumi/Trace.hs b/src/Shikumi/Trace.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Trace.hs @@ -0,0 +1,405 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Hierarchical tracing for shikumi (EP-7, M1). +-- +-- The 'Trace' effect produces a /tree/ of timed, attributed spans for a program +-- run. A span is a node of one of four 'SpanKind's — a whole program, a module +-- such as @predict@, a combinator such as @Pipeline@, or a single LM call. Each +-- span records its parent, so nesting @withSpan@s builds a real tree even though +-- neither baikai nor EP-1's @LLM@ effect has any parent\/child concept. +-- +-- The hierarchy is shikumi's: 'runTrace' keeps a stack of span ids, and 'withSpan' +-- pushes\/pops. LM-call leaves are captured automatically by 'tracedLLM', which +-- __interposes__ on EP-1's @LLM@ effect (the same seam EP-6's @cachedLLM@ uses): +-- on each 'Shikumi.LLM.complete' it opens an 'LlmCallSpan' under the active span +-- and fills its attributes — model, provider, latency, tokens, cost, tool calls, +-- the recorded response JSON, and the EP-6 content-addressed 'Shikumi.Cache.Key.cacheKey' +-- (integration point #7) — from the returned 'Baikai.Response'. +-- +-- 'renderTree' pretty-prints the tree; 'Shikumi.Trace.Store' serializes it and +-- 'Shikumi.Trace.Replay' replays it offline. +module Shikumi.Trace + ( -- * Span and tree types + SpanKind (..), + SpanId (..), + ToolCallRecord (..), + SpanAttrs (..), + emptyAttrs, + Span (..), + TraceTree (..), + childrenOf, + + -- * Node identity (re-exported from "Shikumi.Trace.Node") + NodePath (..), + + -- * The effect + Trace, + withSpan, + currentSpanId, + bumpRetry, + recordToolCall, + annotateSpan, + + -- * Interpreters and capture + runTrace, + tracedLLM, + llmLabel, + llmAttrs, + + -- * Rendering + renderTree, + ) +where + +import Baikai + ( AssistantContent (..), + Context, + Model, + Options, + Response, + flattenAssistantBlocks, + ) +import Control.Lens ((^.)) +import Data.Aeson (FromJSON, FromJSONKey, ToJSON, ToJSONKey, Value, toJSON) +import Data.Generics.Labels () +import Data.List (sortOn) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Maybe (fromMaybe) +import Data.Scientific (Scientific) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Time.Clock (UTCTime, diffUTCTime) +import Data.Vector qualified as V +import Effectful (Dispatch (Dynamic), DispatchOf, Eff, Effect, (:>)) +import Effectful.Dispatch.Dynamic (interpose, interpret, localSeqUnlift, send) +import Effectful.Exception (bracket) +import Effectful.Prim (Prim) +import Effectful.Prim.IORef + ( IORef, + atomicModifyIORef', + modifyIORef', + newIORef, + readIORef, + ) +import GHC.Generics (Generic) +import Numeric.Natural (Natural) +import Shikumi.Cache.Key (CacheKey (..), requestToCanonicalValue) +import Shikumi.Cache.Key qualified as Key +import Shikumi.Effect.Time (Time, getCurrentTime) +import Shikumi.LLM (LLM (..), complete, stream) +import Shikumi.Trace.Node (NodePath (..)) +import Shikumi.Trace.ResponseJSON () + +-- --------------------------------------------------------------------------- +-- Span and tree types +-- --------------------------------------------------------------------------- + +-- | What sort of work a span records. +data SpanKind = ProgramSpan | ModuleSpan | CombinatorSpan | LlmCallSpan + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +-- | A span identifier (a fresh @"span-N"@ string allocated by 'runTrace'). +newtype SpanId = SpanId Text + deriving stock (Eq, Ord, Show, Generic) + deriving newtype (ToJSON, FromJSON, ToJSONKey, FromJSONKey) + +-- | A tool call observed on an LM response: the tool name and its JSON arguments. +data ToolCallRecord = ToolCallRecord + { name :: !Text, + arguments :: !Value + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +-- | The spec-mandated attribute bag a span can carry. Most fields are populated +-- only on 'LlmCallSpan' nodes; the structural nodes (program\/module\/combinator) +-- leave them empty. +data SpanAttrs = SpanAttrs + { -- | model id, for LM-call spans + model :: !(Maybe Text), + provider :: !(Maybe Text), + -- | the request as canonical JSON (system, messages, tools, options) + prompt :: !(Maybe Value), + -- | the response payload as JSON (for replay + inspection) + response :: !(Maybe Value), + latencyMs :: !(Maybe Integer), + inputTokens :: !(Maybe Natural), + outputTokens :: !(Maybe Natural), + costUsd :: !(Maybe Scientific), + -- | how many times this span retried its body + retries :: !Int, + -- | tool calls observed on the response + toolCalls :: ![ToolCallRecord], + -- | the EP-6 content-addressed key, present on LM-call spans + cacheKey :: !(Maybe Text), + -- | the structural path of the @Program@ node that issued this span's LM call + -- (EP-16). Present only on model-call spans produced by @runProgramTraced@; + -- 'Nothing' for spans opened by bare 'withSpan' or a non-node-correlated run. + nodePath :: !(Maybe NodePath) + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +-- | A span with everything empty: the base every node starts from. +emptyAttrs :: SpanAttrs +emptyAttrs = + SpanAttrs + { model = Nothing, + provider = Nothing, + prompt = Nothing, + response = Nothing, + latencyMs = Nothing, + inputTokens = Nothing, + outputTokens = Nothing, + costUsd = Nothing, + retries = 0, + toolCalls = [], + cacheKey = Nothing, + nodePath = Nothing + } + +-- | One node of the trace tree. +data Span = Span + { spanId :: !SpanId, + parent :: !(Maybe SpanId), + kind :: !SpanKind, + label :: !Text, + startedAt :: !UTCTime, + endedAt :: !(Maybe UTCTime), + attrs :: !SpanAttrs + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +-- | A finished trace tree: a flat 'Map' of spans plus the root id. The nesting is +-- reconstructed from each span's 'parent' pointer ('childrenOf'). +data TraceTree = TraceTree + { root :: !SpanId, + spans :: !(Map SpanId Span) + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +-- | The children of a span, in creation order (sorted by start time, ties broken +-- by span id). +childrenOf :: TraceTree -> SpanId -> [SpanId] +childrenOf t sid = + map spanId $ + sortOn (\s -> (startedAt s, spanId s)) $ + [s | s <- Map.elems (spans t), parent s == Just sid] + +-- --------------------------------------------------------------------------- +-- The effect +-- --------------------------------------------------------------------------- + +-- | The hierarchical tracing effect. +data Trace :: Effect where + WithSpan :: SpanKind -> Text -> m a -> Trace m a + CurrentSpanId :: Trace m (Maybe SpanId) + BumpRetry :: Trace m () + RecordToolCall :: ToolCallRecord -> Trace m () + AnnotateSpan :: (SpanAttrs -> SpanAttrs) -> Trace m () + +type instance DispatchOf Trace = 'Dynamic + +-- | Open a span of the given kind and label around an inner computation: it is +-- pushed as a child of the currently-active span, timed, and recorded on exit. +withSpan :: (Trace :> es) => SpanKind -> Text -> Eff es a -> Eff es a +withSpan k lbl act = send (WithSpan k lbl act) + +-- | The id of the currently-active span, if any. +currentSpanId :: (Trace :> es) => Eff es (Maybe SpanId) +currentSpanId = send CurrentSpanId + +-- | Increment the active span's retry counter. +bumpRetry :: (Trace :> es) => Eff es () +bumpRetry = send BumpRetry + +-- | Record a tool call on the active span. +recordToolCall :: (Trace :> es) => ToolCallRecord -> Eff es () +recordToolCall = send . RecordToolCall + +-- | Apply a function to the active span's attributes (used by 'tracedLLM' to fill +-- an LM-call span after the response arrives). +annotateSpan :: (Trace :> es) => (SpanAttrs -> SpanAttrs) -> Eff es () +annotateSpan = send . AnnotateSpan + +-- --------------------------------------------------------------------------- +-- The interpreter +-- --------------------------------------------------------------------------- + +-- | Mutable building state for one trace, all in 'IORef's. +data TraceState = TraceState + { tsCounter :: !(IORef Int), + tsStack :: !(IORef [SpanId]), + tsSpans :: !(IORef (Map SpanId Span)), + tsRoot :: !(IORef (Maybe SpanId)) + } + +-- | Run a traced computation, returning its result and the finished tree. +-- +-- The building state lives in 'IORef's reached through the 'Prim' effect (so no +-- open-ended @IOE@ is needed here — only in-process mutation), and span +-- timestamps come from shikumi's own 'Time' effect. Both are discharged at the +-- program edge by 'runPrim' and 'runTime'. +runTrace :: (Prim :> es, Time :> es) => Eff (Trace : es) a -> Eff es (a, TraceTree) +runTrace act = do + st <- newTraceState + a <- + interpret + ( \env -> \case + CurrentSpanId -> safeHead <$> readIORef (tsStack st) + BumpRetry -> modifyActive st (\a' -> a' {retries = retries a' + 1}) + RecordToolCall tc -> modifyActive st (\a' -> a' {toolCalls = toolCalls a' ++ [tc]}) + AnnotateSpan f -> modifyActive st f + WithSpan k lbl inner -> + bracket + (openSpan st k lbl) + (closeSpan st) + (\_ -> localSeqUnlift env (\unlift -> unlift inner)) + ) + act + tree <- freezeTree st + pure (a, tree) + +newTraceState :: (Prim :> es) => Eff es TraceState +newTraceState = + TraceState <$> newIORef 0 <*> newIORef [] <*> newIORef Map.empty <*> newIORef Nothing + +-- | Allocate a fresh span as a child of the current top-of-stack, insert it +-- (open, 'endedAt' = 'Nothing'), push it, and record it as the root if it is the +-- first parentless span. +openSpan :: (Prim :> es, Time :> es) => TraceState -> SpanKind -> Text -> Eff es SpanId +openSpan st k lbl = do + n <- atomicModifyIORef' (tsCounter st) (\i -> (i + 1, i + 1)) + let sid = SpanId ("span-" <> T.pack (show n)) + par <- safeHead <$> readIORef (tsStack st) + now <- getCurrentTime + let s = Span sid par k lbl now Nothing emptyAttrs + modifyIORef' (tsSpans st) (Map.insert sid s) + modifyIORef' (tsStack st) (sid :) + case par of + Nothing -> modifyIORef' (tsRoot st) (Just . fromMaybe sid) + Just _ -> pure () + pure sid + +-- | Close a span: stamp its 'endedAt' and pop it off the stack. +closeSpan :: (Prim :> es, Time :> es) => TraceState -> SpanId -> Eff es () +closeSpan st sid = do + now <- getCurrentTime + modifyIORef' (tsSpans st) (Map.adjust (\s -> s {endedAt = Just now}) sid) + modifyIORef' (tsStack st) (drop 1) + +-- | Apply a function to the active span's attributes (a no-op with no active span). +modifyActive :: (Prim :> es) => TraceState -> (SpanAttrs -> SpanAttrs) -> Eff es () +modifyActive st f = do + stk <- readIORef (tsStack st) + case stk of + (sid : _) -> modifyIORef' (tsSpans st) (Map.adjust (\s -> s {attrs = f (attrs s)}) sid) + [] -> pure () + +-- | Freeze the building state into an immutable 'TraceTree'. +freezeTree :: (Prim :> es) => TraceState -> Eff es TraceTree +freezeTree st = do + sp <- readIORef (tsSpans st) + r <- readIORef (tsRoot st) + pure (TraceTree (fromMaybe (SpanId "") r) sp) + +safeHead :: [a] -> Maybe a +safeHead [] = Nothing +safeHead (x : _) = Just x + +-- --------------------------------------------------------------------------- +-- LM-call capture (interpose on the LLM effect) +-- --------------------------------------------------------------------------- + +-- | Capture every LM call as a leaf 'LlmCallSpan' under the active span. Interpose +-- on EP-1's @LLM@ effect: open a span, delegate to the underlying handler, then +-- fill the span's attributes from the returned 'Response' (and the request). The +-- streaming op is wrapped in a span but its attributes are left empty (streams +-- carry the same data incrementally; the demo/replay path uses 'complete'). +tracedLLM :: (Trace :> es, LLM :> es) => Eff es a -> Eff es a +tracedLLM = interpose $ \_ -> \case + Complete m c o -> withSpan LlmCallSpan (llmLabel m) $ do + resp <- complete m c o + annotateSpan (const (llmAttrs m c o resp)) + pure resp + Stream m c o -> withSpan LlmCallSpan (llmLabel m) (stream m c o) + +-- | The label for an LM-call span: @provider/model-id@. +llmLabel :: Model -> Text +llmLabel m = (m ^. #provider) <> "/" <> (m ^. #modelId) + +-- | Build the attribute bag for a completed LM call from the request and response. +llmAttrs :: Model -> Context -> Options -> Response -> SpanAttrs +llmAttrs m c o resp = + emptyAttrs + { model = Just (m ^. #modelId), + provider = Just (m ^. #provider), + prompt = Just (requestToCanonicalValue m c o), + response = Just (toJSON resp), + latencyMs = Just (resp ^. #latencyMs), + inputTokens = Just (resp ^. #message . #usage . #inputTokens), + outputTokens = Just (resp ^. #message . #usage . #outputTokens), + costUsd = Just (realToFrac (resp ^. #message . #usage . #cost . #usd :: Rational)), + toolCalls = toolCallsOf resp, + cacheKey = Just (unCacheKey (Key.cacheKey m c o)) + } + +-- | Extract the tool calls from a response's assistant content blocks. +toolCallsOf :: Response -> [ToolCallRecord] +toolCallsOf resp = + [ ToolCallRecord (tc ^. #name) (tc ^. #arguments) + | AssistantToolCall tc <- V.toList (flattenAssistantBlocks resp) + ] + +-- --------------------------------------------------------------------------- +-- Rendering +-- --------------------------------------------------------------------------- + +-- | Pretty-print a trace tree as an indented outline, one line per span: kind, +-- label, wall-clock duration, and (for LM-call spans) tokens and cost. +renderTree :: TraceTree -> Text +renderTree t + | Map.null (spans t) = "(empty trace)\n" + | otherwise = T.concat (go (0 :: Int) (root t)) + where + go depth sid = case Map.lookup sid (spans t) of + Nothing -> [] + Just s -> line depth s : concatMap (go (depth + 1)) (childrenOf t sid) + line depth s = + T.replicate (depth * 2) " " + <> kindTag (kind s) + <> " " + <> label s + <> durationOf s + <> llmStats s + <> "\n" + durationOf s = case endedAt s of + Just e -> + let ms = round (realToFrac (diffUTCTime e (startedAt s)) * 1000 :: Double) :: Integer + in " " <> tshow ms <> "ms" + Nothing -> "" + llmStats s + | kind s == LlmCallSpan = + let a = attrs s + in " in=" + <> maybe "?" tshow (inputTokens a) + <> " out=" + <> maybe "?" tshow (outputTokens a) + <> " $" + <> maybe "0" tshow (costUsd a) + | otherwise = "" + +kindTag :: SpanKind -> Text +kindTag = \case + ProgramSpan -> "program" + ModuleSpan -> "module" + CombinatorSpan -> "combinator" + LlmCallSpan -> "llm-call" + +tshow :: (Show a) => a -> Text +tshow = T.pack . show diff --git a/src/Shikumi/Trace/Demo.hs b/src/Shikumi/Trace/Demo.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Trace/Demo.hs @@ -0,0 +1,169 @@ +-- | The end-to-end demo (EP-7, M5): a two-stage pipeline that drafts a one-line +-- summary of a fixed article and then critiques the draft, tying together the +-- whole story — live trace → pretty-print → persist → deterministic offline +-- replay with identical output and zero provider calls. +-- +-- The pipeline ('demoPipeline') is the shared core used by both the +-- @shikumi-trace-demo@ executable ('demoMain') and the @-p e2e@ acceptance test. +-- It is provider-agnostic: it issues 'Shikumi.LLM.complete' calls through the +-- @LLM@ effect and nests its stages in 'withSpan's, so the live run captures a +-- tree and the replay run is byte-identical. +-- +-- For reproducibility the demo uses a deterministic in-process stub +-- ('demoResponder') rather than a real provider — no API keys, no network — so a +-- novice can run it anywhere and replay is well-defined. (The plan sketched an +-- optional live-provider path; the canned stub is the hermetic, always-runnable +-- choice. Replaying never contacts a provider regardless: 'runLLMReplay' has no +-- registry.) +module Shikumi.Trace.Demo + ( demoArticle, + demoModel, + demoPipeline, + demoResponder, + runStubLLM, + demoMain, + ) +where + +import Baikai + ( Api (Custom), + AssistantContent (..), + Context, + Message (UserMessage), + Model, + Options, + Response, + TextContent (..), + UserContent (UserText), + user, + _Context, + _Model, + _Options, + _Response, + _TextContent, + ) +import Control.Lens ((&), (.~), (^.)) +import Data.Generics.Labels () +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.IO qualified as TIO +import Data.Vector qualified as V +import Effectful (Eff, runEff, (:>)) +import Effectful.Dispatch.Dynamic (interpret) +import Effectful.Prim (runPrim) +import Shikumi.Effect.Time (runTime) +import Shikumi.LLM (LLM (..), complete) +import Shikumi.Trace + ( SpanKind (ModuleSpan, ProgramSpan), + Trace, + renderTree, + runTrace, + tracedLLM, + withSpan, + ) +import Shikumi.Trace.Replay (runLLMReplay) +import Shikumi.Trace.Store (readTraceFile, replayIndex, writeTraceFile) +import System.Environment (lookupEnv) + +-- | The fixed input article the demo summarizes. +demoArticle :: Text +demoArticle = "Shikumi turns LM calls into typed, traceable, replayable programs." + +-- | The demo model's routing identity (LM-call spans read @stub/stub-model@). +demoModel :: Model +demoModel = + _Model + & #provider .~ "stub" + & #modelId .~ "stub-model" + & #api .~ Custom "stub" + +-- | The two-stage pipeline: draft a one-line summary, then critique the draft. +-- Stage two's request depends on stage one's output, so replay must reproduce +-- stage one exactly to reach stage two. Returns the critique. +demoPipeline :: (Trace :> es, LLM :> es) => Text -> Eff es Text +demoPipeline article = + withSpan ProgramSpan "summarize-and-critique" $ do + draft <- + withSpan ModuleSpan "predict:Draft" $ + textOf <$> complete demoModel (ctxFor ("draft: " <> article)) opts + withSpan ModuleSpan "predict:Critique" $ + textOf <$> complete demoModel (ctxFor ("critique: " <> draft)) opts + +-- | The deterministic stub: a draft request yields a fixed summary; a critique +-- request returns a fixed verdict. Deterministic per request, so replay is +-- well-defined. +demoResponder :: Context -> Response +demoResponder c = + if "draft:" `T.isPrefixOf` lastUserText c + then mkResponse "Shikumi makes LM programs typed and replayable." + else mkResponse "Accurate and concise; ship it." + +-- | A base @LLM@ interpreter driving the demo from 'demoResponder'. +runStubLLM :: Eff (LLM : es) a -> Eff es a +runStubLLM = interpret $ \_ -> \case + Complete _ c _ -> pure (demoResponder c) + Stream {} -> pure [] + +-- --------------------------------------------------------------------------- +-- Executable entry point +-- --------------------------------------------------------------------------- + +-- | The demo's command-line core. With no arguments it runs the pipeline live, +-- prints the trace tree and the final critique, and writes @trace.json@. With +-- @--replay PATH@ it loads the trace, replays the pipeline offline, and prints +-- the same critique plus @provider calls: 0@. +demoMain :: [String] -> IO () +demoMain args = case args of + [] -> liveMode + ["--replay", path] -> replayMode path + _ -> TIO.putStrLn "usage: shikumi-trace-demo [--replay PATH]" + +liveMode :: IO () +liveMode = do + (critique, tree) <- + runEff . runPrim . runTime . runTrace . runStubLLM . tracedLLM $ demoPipeline demoArticle + TIO.putStr (renderTree tree) + writeTraceFile "trace.json" tree + TIO.putStrLn ("FINAL: " <> critique) + +replayMode :: FilePath -> IO () +replayMode path = do + _ <- lookupEnv "SHIKUMI_OFFLINE" -- offline is structural: replay contacts no provider + loaded <- readTraceFile path + case loaded of + Left err -> TIO.putStrLn ("trace load error: " <> err) + Right tree -> do + let idx = replayIndex tree + (critique, _) <- runEff . runPrim . runTime . runTrace . runLLMReplay idx $ demoPipeline demoArticle + TIO.putStrLn ("FINAL: " <> critique) + TIO.putStrLn "provider calls: 0" + +-- --------------------------------------------------------------------------- +-- helpers +-- --------------------------------------------------------------------------- + +opts :: Options +opts = _Options + +ctxFor :: Text -> Context +ctxFor t = _Context & #messages .~ V.singleton (user t) + +mkResponse :: Text -> Response +mkResponse t = + _Response + & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t)) + & #message . #usage . #inputTokens .~ 24 + & #message . #usage . #outputTokens .~ 9 + & #latencyMs .~ 7 + +textOf :: Response -> Text +textOf resp = + mconcat [t | AssistantText (TextContent t) <- V.toList (resp ^. #message . #content)] + +lastUserText :: Context -> Text +lastUserText c = + mconcat + [ t + | UserMessage up <- V.toList (c ^. #messages), + UserText (TextContent t) <- V.toList (up ^. #content) + ] diff --git a/src/Shikumi/Trace/Feedback.hs b/src/Shikumi/Trace/Feedback.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Trace/Feedback.hs @@ -0,0 +1,78 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE TypeFamilies #-} + +-- | The per-node feedback channel (EP-16, M3). +-- +-- A 'FeedbackLog' is a sibling of the trace (deliberately /not/ part of +-- 'Shikumi.Trace.TraceTree', whose serialized format is pinned and consumed by +-- replay): it maps a node's 'NodePath' to the textual critiques attached to it. A +-- metric or LM-judge writes a critique with 'attachFeedback' during evaluation; an +-- optimizer reads all critiques for a node with 'feedbackFor' during a proposal +-- step. GEPA (@docs/plans/22-gepa-reflective-optimizer.md@) is the headline +-- consumer. +module Shikumi.Trace.Feedback + ( FeedbackLog (..), + emptyFeedback, + feedbackFor, + Feedback (..), + attachFeedback, + runFeedback, + ) +where + +import Data.Aeson (FromJSON (..), ToJSON (..)) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Text (Text) +import Effectful (Dispatch (Dynamic), DispatchOf, Eff, Effect, (:>)) +import Effectful.Dispatch.Dynamic (interpret, send) +import Effectful.Prim (Prim) +import Effectful.Prim.IORef (modifyIORef', newIORef, readIORef) +import Shikumi.Trace.Node (NodePath) + +-- | All feedback gathered during a run: critiques keyed by the node they target. A +-- node may accumulate several critiques (e.g. one per example), kept in attach +-- order. +newtype FeedbackLog = FeedbackLog (Map NodePath [Text]) + deriving stock (Eq, Show) + +-- | Serialized as a list of @(NodePath, [Text])@ pairs, so 'NodePath' needs only +-- 'ToJSON'\/'FromJSON' (not a text 'ToJSONKey'). +instance ToJSON FeedbackLog where + toJSON (FeedbackLog m) = toJSON (Map.toList m) + +instance FromJSON FeedbackLog where + parseJSON v = FeedbackLog . Map.fromList <$> parseJSON v + +-- | The empty log. +emptyFeedback :: FeedbackLog +emptyFeedback = FeedbackLog Map.empty + +-- | All critiques attached to a node, in attach order (empty if none). +feedbackFor :: NodePath -> FeedbackLog -> [Text] +feedbackFor p (FeedbackLog m) = Map.findWithDefault [] p m + +-- | The effect a metric or judge uses to attach feedback during a run. +data Feedback :: Effect where + AttachFeedback :: NodePath -> Text -> Feedback m () + +type instance DispatchOf Feedback = 'Dynamic + +-- | Attach a short textual critique to a node. +attachFeedback :: (Feedback :> es) => NodePath -> Text -> Eff es () +attachFeedback p t = send (AttachFeedback p t) + +-- | Run the feedback effect, returning the result paired with the collected log. +-- Accumulates into an @IORef@ via 'Prim' (the same mechanism @runTrace@ uses), +-- appending each critique to its node's list. +runFeedback :: (Prim :> es) => Eff (Feedback : es) a -> Eff es (a, FeedbackLog) +runFeedback act = do + ref <- newIORef Map.empty + a <- + interpret + ( \_ -> \case + AttachFeedback p t -> modifyIORef' ref (Map.insertWith (\new old -> old ++ new) p [t]) + ) + act + m <- readIORef ref + pure (a, FeedbackLog m) diff --git a/src/Shikumi/Trace/Internal/Spike.hs b/src/Shikumi/Trace/Internal/Spike.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Trace/Internal/Spike.hs @@ -0,0 +1,142 @@ +{-# LANGUAGE GADTs #-} + +-- | M0 de-risking spike for EP-7 +-- (@docs/plans/7-hierarchical-tracing-observability-and-replay.md@). +-- +-- The one genuinely novel piece of the tracing plan is /how the span hierarchy is +-- formed/, since neither baikai nor EP-1's @LLM@ effect carries a parent\/child +-- relationship. This spike proves the mechanism the production effect (M1) is +-- built on: +-- +-- * a manually-maintained __stack of span ids__ (an @IORef [SpanId]@), and +-- * an __interpose__ over EP-1's @LLM@ effect that, on each 'Shikumi.LLM.complete', +-- reads the id currently on top of the stack and tags the call with it before +-- delegating to the real handler. +-- +-- This is the same @interpose@ seam EP-6's @cachedLLM@ already uses, so the risk +-- is low; the spike confirms that the id on top of the stack /at the moment the +-- call is sent/ is the enclosing span, which is exactly what M1 needs. +-- +-- Decision (EP-7, 2026-06-08): capture LM calls by interposing on the @LLM@ +-- effect rather than installing a baikai 'Baikai.Trace.Sink.TraceSink'. EP-1's +-- interpreters expose no sink parameter, and @LLM.complete@ returns the full +-- baikai 'Baikai.Response' (latency, usage, cost, tool blocks) — strictly more +-- than baikai's flat @TraceEvent@. See the plan's Decision Log. +module Shikumi.Trace.Internal.Spike + ( SpanId (..), + runSpike, + ) +where + +import Baikai + ( AssistantContent (..), + Context, + Response, + TextContent (..), + _Context, + _Model, + _Options, + _Response, + _TextContent, + ) +import Control.Lens ((&), (.~), (^.)) +import Data.Generics.Labels () +import Data.IORef (IORef, modifyIORef', newIORef, readIORef) +import Data.Text (Text) +import Data.Vector qualified as V +import Effectful (Eff, IOE, liftIO, runEff, (:>)) +import Effectful.Dispatch.Dynamic (interpose, interpret) +import Shikumi.LLM (LLM (..), complete) + +-- | A span identifier. The production effect (M1) reuses this newtype. +newtype SpanId = SpanId Text + deriving stock (Eq, Ord, Show) + +-- | Run two fake LM calls nested under two distinct span ids and return, for +-- each captured call, the @(enclosing span id, response text)@ pair. If the +-- mechanism works, the first call is tagged with @"draft"@ and the second with +-- @"critique"@ — the ids that were on top of the stack when each call ran. +runSpike :: IO [(SpanId, Text)] +runSpike = runEff $ do + stack <- liftIO (newIORef []) + captured <- liftIO (newIORef []) + runScripted [stubResponse "first draft", stubResponse "a sharp critique"] + . traceSpike stack captured + $ do + _ <- withSpike stack (SpanId "draft") oneCall + _ <- withSpike stack (SpanId "critique") oneCall + pure () + liftIO (readIORef captured) + where + oneCall = complete _Model neutralCtx _Options + +neutralCtx :: Context +neutralCtx = _Context + +-- --------------------------------------------------------------------------- +-- The mechanism under test +-- --------------------------------------------------------------------------- + +-- | Push @sid@ on the stack for the dynamic extent of @act@, then pop it. The +-- enclosing span is whatever is on top while @act@ runs. +withSpike :: (IOE :> es) => IORef [SpanId] -> SpanId -> Eff es a -> Eff es a +withSpike stack sid act = do + liftIO (modifyIORef' stack (sid :)) + r <- act + liftIO (modifyIORef' stack (drop 1)) + pure r + +-- | Interpose on @LLM@: on each 'Complete', read the top-of-stack span id, +-- delegate to the underlying handler, and record @(top, responseText)@. +traceSpike :: + (IOE :> es, LLM :> es) => + IORef [SpanId] -> + IORef [(SpanId, Text)] -> + Eff es a -> + Eff es a +traceSpike stack captured = interpose $ \_ -> \case + Complete m c o -> do + top <- topOf <$> liftIO (readIORef stack) + resp <- complete m c o + liftIO (modifyIORef' captured (++ [(top, responseText resp)])) + pure resp + Stream m c o -> stream' m c o + where + topOf [] = SpanId "" + topOf (x : _) = x + stream' m c o = complete m c o >> pure [] + +-- --------------------------------------------------------------------------- +-- A network-free scripted @LLM@ interpreter (pops canned responses in order) +-- --------------------------------------------------------------------------- + +runScripted :: (IOE :> es) => [Response] -> Eff (LLM : es) a -> Eff es a +runScripted rs0 act = do + ref <- liftIO (newIORef rs0) + interpret + ( \_ -> \case + Complete {} -> liftIO (pop ref) + Stream {} -> pure [] + ) + act + where + pop :: IORef [Response] -> IO Response + pop ref = do + rs <- readIORef ref + case rs of + (x : xs) -> modifyIORef' ref (const xs) >> pure x + [] -> pure (stubResponse "") + +-- --------------------------------------------------------------------------- +-- Tiny helpers +-- --------------------------------------------------------------------------- + +-- | A 'Response' carrying @t@ as its single assistant text block. +stubResponse :: Text -> Response +stubResponse t = + _Response & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t)) + +-- | Concatenate the assistant text blocks of a response. +responseText :: Response -> Text +responseText resp = + mconcat [t | AssistantText (TextContent t) <- V.toList (resp ^. #message . #content)] diff --git a/src/Shikumi/Trace/Node.hs b/src/Shikumi/Trace/Node.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Trace/Node.hs @@ -0,0 +1,126 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Node identity for the trace (EP-16, M1). +-- +-- A 'NodePath' names a 'Shikumi.Program.Program' node's structural position as the +-- list of branch steps from the program root. 'programNodePaths' enumerates the +-- path of every @Predict@ node in /exactly/ the left-to-right depth-first order +-- 'Shikumi.Program.foldParams' yields their @Params@ — so @programNodePaths p !! n@ +-- is the path to the same node @mapParamsAt n@ edits and @foldParams p !! n@ +-- parameterizes. That agreement is the parent MasterPlan's integration point #3, +-- and it is structural (the enumeration reuses the identical depth-first descent), +-- not a coincidence two functions must independently maintain. +module Shikumi.Trace.Node + ( NodeStep (..), + NodePath (..), + programNodePaths, + renderNodePath, + nodeFields, + ) +where + +import Data.Aeson (FromJSON, ToJSON) +import Data.Text (Text) +import Data.Text qualified as T +import GHC.Generics (Generic) +import Shikumi.Program + ( NodeFields, + Program + ( Compose, + Embed, + Ensemble, + FMap, + MajorityVote, + Map, + Parallel, + Predict, + Retry, + RetryWhen, + Validate + ), + nodeFieldsIndexed, + ) + +-- | One structural step from a parent node to a child, naming which branch was +-- taken. The labels mirror the 'Shikumi.Program.Program' constructors so a path is +-- human-readable and shape-stable: two programs of the same shape yield identical +-- paths. +data NodeStep + = -- | into the left side of @Compose@ + StepComposeL + | -- | into the right side of @Compose@ + StepComposeR + | -- | through an @FMap@ + StepFMap + | -- | through a @Map@ + StepMap + | -- | into the left side of @Parallel@ + StepParallelL + | -- | into the right side of @Parallel@ + StepParallelR + | -- | through a @Retry@ + StepRetry + | -- | through a @RetryWhen@ + StepRetryWhen + | -- | through a @Validate@ + StepValidate + | -- | through a @MajorityVote@ + StepMajorityVote + | -- | into the i-th member of an @Ensemble@ (0-based) + StepEnsemble !Int + deriving stock (Eq, Ord, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +-- | The structural position of a node within a 'Shikumi.Program.Program', as the +-- list of steps from the program root to that node, outermost first. A bare +-- root @Predict@ has the empty path. +newtype NodePath = NodePath [NodeStep] + deriving stock (Eq, Ord, Show, Generic) + deriving newtype (ToJSON, FromJSON) + +-- | Enumerate the 'NodePath' of every @Predict@ node, in the same left-to-right +-- depth-first order as 'Shikumi.Program.foldParams'. Reuses that identical descent +-- so length and node order agree by construction. +programNodePaths :: Program i o -> [NodePath] +programNodePaths = go [] + where + -- @prefix@ is the reversed list of steps taken so far; we reverse on emit. + go :: forall x y. [NodeStep] -> Program x y -> [NodePath] + go prefix (Predict _ _) = [NodePath (reverse prefix)] + go prefix (Compose a b) = go (StepComposeL : prefix) a ++ go (StepComposeR : prefix) b + go prefix (FMap _ p) = go (StepFMap : prefix) p + go prefix (Map _ p) = go (StepMap : prefix) p + go prefix (Parallel a b) = go (StepParallelL : prefix) a ++ go (StepParallelR : prefix) b + go prefix (Retry _ p) = go (StepRetry : prefix) p + go prefix (RetryWhen _ _ p) = go (StepRetryWhen : prefix) p + go prefix (Validate _ p) = go (StepValidate : prefix) p + go prefix (MajorityVote _ _ p) = go (StepMajorityVote : prefix) p + go prefix (Ensemble ps _) = concat (zipWith (\i p -> go (StepEnsemble i : prefix) p) [0 ..] ps) + go _ (Embed _) = [] + +-- | Render a 'NodePath' to a short, stable string (e.g. @compose.0/predict@-style +-- slash-joined steps), suitable as a trace/OTel attribute value. The empty path +-- renders as @\"root\"@. +renderNodePath :: NodePath -> Text +renderNodePath (NodePath []) = "root" +renderNodePath (NodePath steps) = T.intercalate "/" (map stepLabel steps) + where + stepLabel = \case + StepComposeL -> "compose.l" + StepComposeR -> "compose.r" + StepFMap -> "fmap" + StepMap -> "map" + StepParallelL -> "parallel.l" + StepParallelR -> "parallel.r" + StepRetry -> "retry" + StepRetryWhen -> "retryWhen" + StepValidate -> "validate" + StepMajorityVote -> "majorityVote" + StepEnsemble i -> "ensemble." <> T.pack (show i) + +-- | Associate each @Predict@ node's 'NodePath' with its input/output field names. +-- A consumer with a 'NodePath' from a trace span zips against this to map +-- path → field metadata; one with a @foldParams@ index uses +-- 'Shikumi.Program.nodeFieldsIndexed' directly. +nodeFields :: Program i o -> [(NodePath, NodeFields)] +nodeFields p = zip (programNodePaths p) (nodeFieldsIndexed p) diff --git a/src/Shikumi/Trace/Program.hs b/src/Shikumi/Trace/Program.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Trace/Program.hs @@ -0,0 +1,186 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Node-correlated program execution for the trace (EP-16, M2). +-- +-- 'runProgramTraced' is an /additive/ entry point: it runs a +-- 'Shikumi.Program.Program' exactly like 'Shikumi.Program.runProgram' — reusing +-- its render\/parse\/retry\/vote semantics by delegating each @Predict@ leaf to +-- @runProgram@ — but it also opens a trace span per node and threads the active +-- node's 'NodePath' so that each model-call span is tagged with the structural +-- position of the node that issued it. 'Shikumi.Program.runProgram' / +-- 'Shikumi.Program.runProgramConc' are untouched (MasterPlan integration point #4). +-- +-- The correlation is threaded by a tiny reader-style 'CurrentNode' effect held in +-- an @IORef@ cell: 'runProgramTraced' wraps each @Predict@ in 'localNode' as it +-- descends (with the same step prefix 'Shikumi.Trace.Node.programNodePaths' uses), +-- and the node-aware capture 'tracedNodeLLM' reads it ('askNode') to stamp the +-- span. The capture and the tag happen inside one 'withSpan', so there is no +-- cross-interpreter ordering subtlety (see the Decision Log of EP-16): use +-- 'tracedNodeLLM' /instead of/ 'Shikumi.Trace.tracedLLM' for a traced-program run. +-- +-- Stack shape (note 'runCurrentNode' is outer of 'tracedNodeLLM' so both share the +-- same node cell): +-- +-- @ +-- runEff . runPrim . runTime . runTrace . runCurrentNode +-- . runKeyedLLM responder -- base LLM interpreter (stub or real) +-- . tracedNodeLLM -- opens the LlmCallSpan, fills attrs, tags nodePath +-- $ runProgramTraced program input +-- @ +module Shikumi.Trace.Program + ( -- * The current-node effect + CurrentNode (..), + askNode, + localNode, + runCurrentNode, + + -- * Node-aware capture and execution + tracedNodeLLM, + runProgramTraced, + ) +where + +import Effectful (Dispatch (Dynamic), DispatchOf, Eff, Effect, (:>)) +import Effectful.Dispatch.Dynamic (interpose, interpret, localSeqUnlift, send) +import Effectful.Error.Static (Error) +import Effectful.Exception (bracket_) +import Effectful.Prim (Prim) +import Effectful.Prim.IORef (newIORef, readIORef, writeIORef) +import Shikumi.Error (ShikumiError) +import Shikumi.LLM (LLM (..), complete, stream) +import Shikumi.Program + ( Program + ( Compose, + Embed, + Ensemble, + FMap, + MajorityVote, + Map, + Parallel, + Predict, + Retry, + RetryWhen, + Validate + ), + acceptOrReject, + modal, + retryWith, + runProgram, + sampleTemps, + withSampleTemp, + ) +import Shikumi.Trace + ( SpanAttrs (..), + SpanKind (CombinatorSpan, LlmCallSpan, ModuleSpan), + Trace, + annotateSpan, + llmAttrs, + llmLabel, + withSpan, + ) +import Shikumi.Trace.Node (NodePath (..), NodeStep (..)) + +-- --------------------------------------------------------------------------- +-- The current-node effect +-- --------------------------------------------------------------------------- + +-- | Carries the 'NodePath' of the node currently executing. 'localNode' sets it +-- for the duration of an inner action; 'askNode' reads it. A dedicated effect +-- (rather than @Effectful.Reader@) keeps the value dynamically scoped as execution +-- descends the program tree without the caller threading it, and keeps +-- 'Shikumi.Program.runProgram'\'s row untouched. +data CurrentNode :: Effect where + AskNode :: CurrentNode m (Maybe NodePath) + LocalNode :: NodePath -> m a -> CurrentNode m a + +type instance DispatchOf CurrentNode = 'Dynamic + +-- | The node currently executing, if any. +askNode :: (CurrentNode :> es) => Eff es (Maybe NodePath) +askNode = send AskNode + +-- | Run an action with the given node marked as current. +localNode :: (CurrentNode :> es) => NodePath -> Eff es a -> Eff es a +localNode p act = send (LocalNode p act) + +-- | Discharge 'CurrentNode' with a mutable cell (via 'Prim'). Because the cell is +-- shared, an 'askNode' issued from a lower interpose ('tracedNodeLLM') sees the +-- value a higher 'localNode' set, which is exactly how a model call is correlated +-- to its issuing node. +runCurrentNode :: (Prim :> es) => Eff (CurrentNode : es) a -> Eff es a +runCurrentNode act = do + ref <- newIORef Nothing + interpret + ( \env -> \case + AskNode -> readIORef ref + LocalNode p inner -> do + old <- readIORef ref + bracket_ + (writeIORef ref (Just p)) + (writeIORef ref old) + (localSeqUnlift env (\unlift -> unlift inner)) + ) + act + +-- --------------------------------------------------------------------------- +-- Node-aware capture +-- --------------------------------------------------------------------------- + +-- | Like 'Shikumi.Trace.tracedLLM' but also stamps the active 'NodePath' onto each +-- model-call span. The capture (model\/prompt\/response\/cost) and the node tag are +-- written inside the /same/ 'withSpan', so the tag always lands on the LM-call span +-- — no dependence on interpose ordering relative to a separate capture layer. +tracedNodeLLM :: (Trace :> es, CurrentNode :> es, LLM :> es) => Eff es a -> Eff es a +tracedNodeLLM = interpose $ \_ -> \case + Complete m c o -> withSpan LlmCallSpan (llmLabel m) $ do + resp <- complete m c o + mp <- askNode + annotateSpan (\_ -> (llmAttrs m c o resp) {nodePath = mp}) + pure resp + Stream m c o -> withSpan LlmCallSpan (llmLabel m) (stream m c o) + +-- --------------------------------------------------------------------------- +-- Node-correlated execution +-- --------------------------------------------------------------------------- + +-- | Run a program like 'Shikumi.Program.runProgram', opening a span per node and +-- tagging each model-call span with the issuing node's 'NodePath'. The step prefix +-- accumulated as it descends is the same one 'Shikumi.Trace.Node.programNodePaths' +-- builds, so a @Predict@ leaf's path here equals the path that enumeration assigns +-- it. Each @Predict@ leaf delegates to @runProgram@ (reusing its exact semantics); +-- the combinators are mirrored with spans and the same control flow. +runProgramTraced :: + forall i o es. + (LLM :> es, Trace :> es, CurrentNode :> es, Error ShikumiError :> es) => + Program i o -> + i -> + Eff es o +runProgramTraced = go [] + where + go :: forall x y. [NodeStep] -> Program x y -> x -> Eff es y + go prefix node@(Predict _ _) i = + withSpan ModuleSpan "Predict" (localNode (NodePath (reverse prefix)) (runProgram node i)) + go prefix (Compose f g) i = + withSpan CombinatorSpan "Compose" (go (StepComposeL : prefix) f i >>= go (StepComposeR : prefix) g) + go prefix (FMap k p) i = + withSpan CombinatorSpan "FMap" (k <$> go (StepFMap : prefix) p i) + go prefix (Map _ p) xs = + withSpan CombinatorSpan "Map" (traverse (go (StepMap : prefix) p) xs) + go prefix (Parallel a b) i = + withSpan CombinatorSpan "Parallel" ((,) <$> go (StepParallelL : prefix) a i <*> go (StepParallelR : prefix) b i) + go prefix (Retry n p) i = + withSpan CombinatorSpan "Retry" (retryWith (go (StepRetry : prefix)) (const True) n p i) + go prefix (RetryWhen ok n p) i = + withSpan CombinatorSpan "RetryWhen" (retryWith (go (StepRetryWhen : prefix)) ok n p i) + go prefix (Validate v p) i = + withSpan CombinatorSpan "Validate" (go (StepValidate : prefix) p i >>= acceptOrReject v) + go prefix (MajorityVote k sched p) i = + withSpan CombinatorSpan "MajorityVote" $ + modal <$> traverse (\mt -> withSampleTemp mt (go (StepMajorityVote : prefix) p i)) (sampleTemps (max 1 k) sched) + go prefix (Ensemble ps reduce) i = + withSpan CombinatorSpan "Ensemble" $ + reduce <$> sequence [go (StepEnsemble idx : prefix) p i | (idx, p) <- zip [0 ..] ps] + go _ (Embed f) i = + withSpan CombinatorSpan "Embed" (f i) diff --git a/src/Shikumi/Trace/Replay.hs b/src/Shikumi/Trace/Replay.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Trace/Replay.hs @@ -0,0 +1,95 @@ +{-# LANGUAGE GADTs #-} + +-- | Deterministic offline replay of the @LLM@ effect (EP-7, M3). +-- +-- 'runLLMReplay' is a drop-in /alternative interpreter/ of EP-1's @LLM@ effect: +-- where the real interpreter dispatches to a provider, this one answers every +-- 'Shikumi.LLM.complete' from a recorded trace. It computes the EP-6 +-- content-addressed 'Shikumi.Cache.Key.cacheKey' of the request, looks it up in a +-- replay index (built by 'Shikumi.Trace.Store.replayIndex'), and returns the +-- recorded 'Baikai.Response' — decoded via the 'Shikumi.Trace.ResponseJSON' +-- instances. Because it replaces the @LLM@ effect at the /same boundary/, the rest +-- of the program (modules, combinators, decoding) runs exactly as in live mode, so +-- the typed outputs are identical; only the leaf LM calls are redirected. +-- +-- Divergence is __fail-closed and loud__: a request whose key is not in the trace +-- raises a typed 'ReplayDivergence' carrying the key, the model id, and a redacted +-- prompt summary. It never falls through to the network and never fabricates a +-- response. (There is no registry in this interpreter at all, so "zero provider +-- calls" is structural, not merely policy.) +module Shikumi.Trace.Replay + ( ReplayDivergence (..), + runLLMReplay, + ) +where + +import Baikai (Context, Model, Response) +import Control.Lens ((^.)) +import Data.Aeson (Result (Error, Success), Value, fromJSON) +import Data.Generics.Labels () +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Text (Text) +import Data.Text qualified as T +import Effectful (Eff) +import Effectful.Dispatch.Dynamic (interpret) +import Effectful.Exception (Exception, throwIO) +import GHC.Generics (Generic) +import Shikumi.Cache.Key (CacheKey (..), cacheKey) +import Shikumi.LLM (LLM (..)) +import Shikumi.Trace.ResponseJSON () + +-- | Raised when a replayed program issues an LM call that the recorded trace +-- never saw. +data ReplayDivergence = ReplayDivergence + { -- | the content-addressed key that was absent from the trace + divergedKey :: !CacheKey, + -- | the model the diverging call targeted + divergedModel :: !Text, + -- | a redacted one-line summary of the diverging request + promptSummary :: !Text + } + deriving stock (Eq, Show, Generic) + deriving anyclass (Exception) + +-- | Interpret the @LLM@ effect by lookup in a replay index instead of calling a +-- provider. A hit returns the recorded response; a miss raises 'ReplayDivergence'. +-- Streaming completions are not replayable and raise a divergence naming the key. +-- +-- No @IOE@ is required: the lookup and decode are pure, and 'ReplayDivergence' is +-- raised with @effectful@'s pure-in-@Eff@ 'throwIO'. The contract the plan sketched +-- as @(IOE :> es)@ is therefore satisfied with a strictly weaker constraint. +runLLMReplay :: Map CacheKey Value -> Eff (LLM : es) a -> Eff es a +runLLMReplay idx = interpret $ \_ -> \case + Complete m c o -> do + let key = cacheKey m c o + case Map.lookup key idx of + Nothing -> throwIO (divergence key m c) + Just v -> case fromJSON v of + Success r -> pure (r :: Response) + Error e -> + throwIO + (divergence key m c) {promptSummary = "recorded response failed to decode: " <> T.pack e} + Stream m c o -> + throwIO + (divergence (cacheKey m c o) m c) {promptSummary = "replay does not support streaming completions"} + +-- | Build a 'ReplayDivergence' for a request. +divergence :: CacheKey -> Model -> Context -> ReplayDivergence +divergence key m c = + ReplayDivergence + { divergedKey = key, + divergedModel = m ^. #modelId, + promptSummary = summarizeRequest m c + } + +-- | A redacted, one-line request summary for divergence errors: the routing +-- identity and the message count (no message bodies — those may be sensitive). +summarizeRequest :: Model -> Context -> Text +summarizeRequest m c = + (m ^. #provider) + <> "/" + <> (m ^. #modelId) + <> " (" + <> T.pack (show (length (c ^. #messages))) + <> " messages)" diff --git a/src/Shikumi/Trace/ResponseJSON.hs b/src/Shikumi/Trace/ResponseJSON.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Trace/ResponseJSON.hs @@ -0,0 +1,17 @@ +-- | A faithful JSON round-trip for baikai's 'Baikai.Response.Response' graph. +-- +-- This module exists for backward compatibility: it re-exports the baikai +-- 'Response'-graph orphan instances ('ToJSON'/'FromJSON' for @Response@, +-- @AssistantPayload@, @Usage@, @Cost@, @CostBreakdown@) from their single home, +-- "Shikumi.Cache.ResponseJSON" in @shikumi-cache@ (EP-6 owns the cache key that +-- @shikumi-trace@ already depends on, so the dependency direction is natural). +-- +-- The instances originally lived here (EP-7). They moved down to @shikumi-cache@ +-- when EP-6's persistent backends — which need the same round-trip to store a +-- response as JSON — landed, so that a module importing both the cache and the +-- trace package never sees a duplicate orphan instance. Importing this module +-- (@import Shikumi.Trace.ResponseJSON ()@) still brings the instances into scope +-- exactly as before. +module Shikumi.Trace.ResponseJSON () where + +import Shikumi.Cache.ResponseJSON () diff --git a/src/Shikumi/Trace/Store.hs b/src/Shikumi/Trace/Store.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Trace/Store.hs @@ -0,0 +1,91 @@ +-- | Persisting a 'TraceTree' to a stable on-disk JSON format, and deriving the +-- replay index from it (EP-7, M2). +-- +-- A trace file is a single JSON document — a 'TraceFile' carrying a +-- @formatVersion@ integer and the whole 'TraceTree'. Persisting the tree as a +-- tree (rather than a flat event log) keeps replay a simple key lookup and keeps +-- the file inspectable with @jq@ (e.g. @jq '.tree.spans | length'@). +-- +-- 'writeTraceFile' is atomic (write to a sibling @.tmp@ then rename), so a crash +-- mid-write never leaves a half-written file. 'readTraceFile' returns 'Left' on a +-- parse error or a format-version mismatch — replay refuses to start on a file it +-- does not understand rather than producing wrong outputs. +-- +-- 'replayIndex' projects the tree into the @Map CacheKey Value@ the replay +-- interpreter ('Shikumi.Trace.Replay') consults: each LM-call span's EP-6 +-- 'Shikumi.Cache.Key.CacheKey' mapped to its recorded response JSON. +module Shikumi.Trace.Store + ( TraceFile (..), + currentFormatVersion, + writeTraceFile, + readTraceFile, + replayIndex, + ) +where + +import Control.Lens ((^.)) +import Data.Aeson (FromJSON, ToJSON, Value, eitherDecode, encode) +import Data.ByteString.Lazy qualified as BL +import Data.Generics.Labels () +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Text (Text) +import Data.Text qualified as T +import GHC.Generics (Generic) +import Shikumi.Cache.Key (CacheKey (..)) +import Shikumi.Trace (SpanKind (LlmCallSpan), TraceTree (..)) +import System.Directory (renameFile) + +-- | The on-disk wrapper: a format version plus the tree. Bump +-- 'currentFormatVersion' on any breaking schema change so old files are rejected +-- loudly rather than mis-decoded. +data TraceFile = TraceFile + { formatVersion :: !Int, + tree :: !TraceTree + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +-- | The trace-file schema version this build reads and writes. Bumped 1→2 in +-- EP-16 for the additive optional @SpanAttrs.nodePath@ field. +currentFormatVersion :: Int +currentFormatVersion = 2 + +-- | Write a tree to @path@ atomically: encode to a sibling @.tmp@ file, then +-- 'renameFile' it into place (an atomic rename on the same filesystem). +writeTraceFile :: FilePath -> TraceTree -> IO () +writeTraceFile path t = do + let tmp = path <> ".tmp" + BL.writeFile tmp (encode (TraceFile currentFormatVersion t)) + renameFile tmp path + +-- | Read a tree from @path@. 'Left' on a JSON parse error or a @formatVersion@ +-- this build does not understand. +readTraceFile :: FilePath -> IO (Either Text TraceTree) +readTraceFile path = do + bs <- BL.readFile path + pure $ case eitherDecode bs of + Left e -> Left ("trace parse error: " <> T.pack e) + Right tf + | formatVersion tf /= currentFormatVersion -> + Left + ( "unsupported trace formatVersion " + <> T.pack (show (formatVersion tf)) + <> " (expected " + <> T.pack (show currentFormatVersion) + <> ")" + ) + | otherwise -> Right (tree tf) + +-- | The replay index: every LM-call span's content-addressed key mapped to its +-- recorded response JSON. A span missing either its @cacheKey@ or its @response@ +-- contributes nothing. +replayIndex :: TraceTree -> Map CacheKey Value +replayIndex t = + Map.fromList + [ (CacheKey ck, v) + | s <- Map.elems (spans t), + (s ^. #kind) == LlmCallSpan, + Just ck <- [s ^. #attrs . #cacheKey], + Just v <- [s ^. #attrs . #response] + ] diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,533 @@ +-- | EP-7 acceptance test suite. Milestones are selected with tasty's @-p@: +-- @spike@ (M0), @tree@ (M1), @store@ (M2), @replay@ (M3), @e2e@ (M5). +module Main (main) where + +import Baikai (Context, Model, Options, Response, user, _Context, _Model, _Options) +import Control.Exception (try) +import Control.Lens ((&), (.~)) +import Data.Aeson (Value (Object), decode, encode, object, (.=)) +import Data.ByteString.Lazy qualified as BL +import Data.Generics.Labels () +import Data.IORef (newIORef, readIORef, writeIORef) +import Data.List (sortOn) +import Data.Map.Strict qualified as Map +import Data.Text (Text) +import Data.Text qualified as T +import Data.Time.Clock (UTCTime, addUTCTime) +import Data.Vector qualified as V +import Effectful (Eff, runEff, type (:>)) +import Effectful.Error.Static (runErrorNoCallStack) +import Effectful.Prim (runPrim) +import GHC.Generics (Generic) +import Shikumi.Adapter (ToPrompt) +import Shikumi.Cache.Key (CacheKey (..)) +import Shikumi.Cache.Key qualified as Key +import Shikumi.Combinator (majorityVote, majorityVoteBy, parallel2, (>>>)) +import Shikumi.Effect.Time (runTime) +import Shikumi.Error (ShikumiError) +import Shikumi.LLM (LLM, complete) +import Shikumi.Module (predict) +import Shikumi.Program + ( NodeFields (..), + Params (..), + Program, + TempSchedule (TempFixed), + foldParams, + mapParamsAt, + ) +import Shikumi.Schema (FromModel, ToSchema, Validatable) +import Shikumi.Signature (Signature, mkSignature) +import Shikumi.Trace + ( Span (..), + SpanAttrs (..), + SpanId (..), + SpanKind (..), + TraceTree (..), + childrenOf, + emptyAttrs, + renderTree, + runTrace, + tracedLLM, + withSpan, + ) +import Shikumi.Trace.Demo (demoArticle, demoPipeline, demoResponder) +import Shikumi.Trace.Feedback + ( FeedbackLog, + attachFeedback, + feedbackFor, + runFeedback, + ) +import Shikumi.Trace.Internal.Spike qualified as Spike +import Shikumi.Trace.Node + ( NodePath (..), + NodeStep (..), + nodeFields, + programNodePaths, + ) +import Shikumi.Trace.Program + ( runCurrentNode, + runProgramTraced, + tracedNodeLLM, + ) +import Shikumi.Trace.Replay (ReplayDivergence (..), runLLMReplay) +import Shikumi.Trace.Store + ( TraceFile (..), + currentFormatVersion, + readTraceFile, + replayIndex, + writeTraceFile, + ) +import System.IO.Temp (withSystemTempDirectory) +import Test.Tasty (TestTree, defaultMain, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) +import Test.Tasty.QuickCheck (Gen, choose, elements, forAll, testProperty, vectorOf, (===)) +import TraceFixtures + ( ctxFor, + lastUserText, + mkResponse, + optsFor, + responseText, + runFixedLLM, + runKeyedCountingLLM, + stubModel, + ) + +main :: IO () +main = + defaultMain $ + testGroup + "shikumi-trace" + [spikeTests, treeTests, storeTests, replayTests, e2eTests, nodeTests, correlateTests, feedbackTests] + +-- --------------------------------------------------------------------------- +-- M0 +-- --------------------------------------------------------------------------- + +-- | M0: the interpose-plus-span-stack mechanism tags each LM call with the +-- enclosing span that was on top of the stack when the call ran. +spikeTests :: TestTree +spikeTests = + testGroup + "spike" + [ testCase "each captured LM call is tagged with its enclosing span id" $ do + captured <- Spike.runSpike + map fst captured @?= [Spike.SpanId "draft", Spike.SpanId "critique"] + map snd captured @?= ["first draft", "a sharp critique"] + ] + +-- --------------------------------------------------------------------------- +-- M1 +-- --------------------------------------------------------------------------- + +-- | M1: a two-@withSpan@ nest containing two stubbed LM calls produces a tree +-- with one root, two module children, each holding exactly one LM-call leaf; +-- @renderTree@ shows the model lines indented under their modules. +treeTests :: TestTree +treeTests = + testGroup + "tree" + [ testCase "nested withSpans build a program -> module -> llm-call tree" $ do + tree <- buildTree + let roots = [s | s <- Map.elems (spans tree), parent s == Nothing] + length roots @?= 1 + kind (spanAt tree (root tree)) @?= ProgramSpan + let kids = childrenOf tree (root tree) + length kids @?= 2 + map (kind . spanAt tree) kids @?= [ModuleSpan, ModuleSpan] + let grandkids = concatMap (childrenOf tree) kids + length grandkids @?= 2 + map (kind . spanAt tree) grandkids @?= [LlmCallSpan, LlmCallSpan], + testCase "renderTree shows the model lines indented under their modules" $ do + tree <- buildTree + let out = renderTree tree + assertBool "two llm-call model lines present" (T.count "stub/stub-model" out == 2) + assertBool "llm-call lines are indented under modules" (T.isInfixOf " llm-call stub/stub-model" out) + assertBool "module lines present" (T.count "module" out == 2) + ] + +-- --------------------------------------------------------------------------- +-- M2 +-- --------------------------------------------------------------------------- + +-- | M2: the trace tree serializes to a stable, versioned JSON document and reads +-- back; the replay index projects LM-call cacheKeys to responses; and the cache +-- key reproduces EP-6's pinned digest (integration point #7). +storeTests :: TestTree +storeTests = + testGroup + "store" + [ testProperty "trace tree survives a JSON round-trip" $ + forAll genTree $ \t -> + decode (encode (TraceFile currentFormatVersion t)) === Just (TraceFile currentFormatVersion t), + testCase "writeTraceFile/readTraceFile round-trips on disk" $ + withSystemTempDirectory "shikumi-trace" $ \dir -> do + tree <- buildTree + let p = dir <> "/trace.json" + writeTraceFile p tree + res <- readTraceFile p + res @?= Right tree, + testCase "reading a file with a foreign formatVersion is Left" $ + withSystemTempDirectory "shikumi-trace" $ \dir -> do + tree <- buildTree + let p = dir <> "/bad.json" + BL.writeFile p (encode (TraceFile 999 tree)) + res <- readTraceFile p + case res of + Left msg -> assertBool "the error names the offending version" (T.isInfixOf "999" msg) + Right _ -> assertFailure "expected Left on a foreign formatVersion", + testCase "replayIndex maps each llm-call cacheKey to its response" $ do + tree <- buildTree + let idx = replayIndex tree + -- two LM-call leaves with distinct requests => two distinct keys. + Map.size idx @?= 2 + assertBool "every indexed value is a JSON object (a recorded response)" (all isObject (Map.elems idx)), + testCase "cacheKey reproduces EP-6's pinned digest (integration point #7)" $ do + let CacheKey hex = Key.cacheKey fixModel fixCtx fixOpts + hex @?= pinnedKey + ] + +-- --------------------------------------------------------------------------- +-- M3 +-- --------------------------------------------------------------------------- + +-- | M3: a two-stage pipeline captured live, persisted, and replayed offline +-- produces byte-identical outputs while making zero provider calls; a mutated +-- request raises a 'ReplayDivergence' naming the missing key. +replayTests :: TestTree +replayTests = + testGroup + "replay" + [ testCase "replay reproduces the live outputs with zero provider calls" $ + withSystemTempDirectory "shikumi-trace" $ \dir -> do + calls <- newIORef (0 :: Int) + -- (1) live: run the pipeline, count provider calls, capture the tree. + (live, tree) <- + runEff . runPrim . runTime . runTrace . runKeyedCountingLLM calls responder . tracedLLM $ + twoStage "the article" + liveCalls <- readIORef calls + liveCalls @?= 2 + -- (2)+(3) persist and reload, build the replay index. + let p = dir <> "/trace.json" + writeTraceFile p tree + Right tree' <- readTraceFile p + let idx = replayIndex tree' + -- (4) replay the SAME pipeline; the counting provider must not be hit. + writeIORef calls 0 + replayed <- runEff . runLLMReplay idx $ twoStage "the article" + replayCalls <- readIORef calls + -- (5)+(6)+(7) identical outputs, zero provider calls. + replayed @?= live + replayCalls @?= 0, + testCase "a request absent from the trace raises ReplayDivergence" $ + withSystemTempDirectory "shikumi-trace" $ \dir -> do + calls <- newIORef (0 :: Int) + (_, tree) <- + runEff . runPrim . runTime . runTrace . runKeyedCountingLLM calls responder . tracedLLM $ + twoStage "the article" + let p = dir <> "/trace.json" + writeTraceFile p tree + Right tree' <- readTraceFile p + let idx = replayIndex tree' + -- A different article => a first-stage key that was never recorded. + res <- try (runEff . runLLMReplay idx $ twoStage "a different article") + case res of + Left (ReplayDivergence {divergedKey = CacheKey k}) -> + assertBool "the divergence names a non-empty key" (not (T.null k)) + Right _ -> assertFailure "expected a ReplayDivergence on an unrecorded request" + ] + +-- --------------------------------------------------------------------------- +-- M5 (e2e) +-- --------------------------------------------------------------------------- + +-- | M5: the demo's two-stage pipeline, captured live then replayed offline, prints +-- a byte-identical final critique while making zero provider calls — the headline +-- end-to-end story, run in-process. +e2eTests :: TestTree +e2eTests = + testGroup + "e2e" + [ testCase "live capture then offline replay: identical FINAL, zero provider calls" $ + withSystemTempDirectory "shikumi-trace" $ \dir -> do + calls <- newIORef (0 :: Int) + -- live run of the demo pipeline; capture the tree, count provider calls. + (liveFinal, tree) <- + runEff . runPrim . runTime . runTrace . runKeyedCountingLLM calls demoResponder . tracedLLM $ + demoPipeline demoArticle + liveCalls <- readIORef calls + liveCalls @?= 2 + -- persist, reload, replay offline. + let p = dir <> "/trace.json" + writeTraceFile p tree + Right tree' <- readTraceFile p + let idx = replayIndex tree' + writeIORef calls 0 + (replayFinal, _) <- runEff . runPrim . runTime . runTrace . runLLMReplay idx $ demoPipeline demoArticle + replayCalls <- readIORef calls + replayFinal @?= liveFinal + replayCalls @?= 0 + ] + +-- | A dependent two-stage pipeline: draft a summary of the article, then critique +-- the draft. The second stage's request depends on the first stage's output, so +-- replay must reproduce stage one exactly to reach the same stage-two request. +twoStage :: (LLM :> es) => Text -> Eff es (Text, Text) +twoStage article = do + draft <- responseText <$> complete stubModel (ctxFor ("draft: " <> article)) optsFor + crit <- responseText <$> complete stubModel (ctxFor ("critique: " <> draft)) optsFor + pure (draft, crit) + +-- | The deterministic stub: a draft request yields a fixed summary; a critique +-- request echoes the draft it was given. Deterministic per request, so replay is +-- well-defined. +responder :: Context -> Response +responder c = + let t = lastUserText c + in if "draft:" `T.isPrefixOf` t + then mkResponse "a tight one-line summary" + else mkResponse ("looks good: " <> T.drop (T.length "critique: ") t) + +-- | The EP-6 golden fixture and its pinned digest, copied here so the two plans +-- are proven byte-for-byte identical. +fixModel :: Model +fixModel = _Model & #modelId .~ "claude-sonnet-4-6" & #provider .~ "anthropic" + +fixCtx :: Context +fixCtx = + _Context + & #systemPrompt .~ Just "You are helpful." + & #messages .~ V.singleton (user "ping") + +fixOpts :: Options +fixOpts = _Options & #temperature .~ Just 0.0 & #maxTokens .~ Just 1024 + +pinnedKey :: Text +pinnedKey = "30b2015562ec8b5cd4fdb64c7cc671c84f56f80d24891deec6676c521f008113" + +-- --------------------------------------------------------------------------- +-- Shared helpers +-- --------------------------------------------------------------------------- + +-- | Run a two-stage traced computation: a program span containing two module +-- spans, each issuing one stubbed LM call (captured automatically by 'tracedLLM'). +buildTree :: IO TraceTree +buildTree = do + let resp = mkResponse "ok" + ((), tree) <- + runEff . runPrim . runTime . runTrace . runFixedLLM resp . tracedLLM $ + withSpan ProgramSpan "summarize-and-critique" $ do + _ <- withSpan ModuleSpan "predict:Draft" (complete stubModel (ctxFor "draft this") optsFor) + _ <- withSpan ModuleSpan "predict:Critique" (complete stubModel (ctxFor "critique that") optsFor) + pure () + pure tree + +spanAt :: TraceTree -> SpanId -> Span +spanAt tree sid = spans tree Map.! sid + +isObject :: Value -> Bool +isObject v = case v of + Object {} -> True + _ -> False + +-- --------------------------------------------------------------------------- +-- A small generator of random trees for the round-trip property +-- --------------------------------------------------------------------------- + +baseTime :: UTCTime +baseTime = read "2026-06-08 00:00:00 UTC" + +-- | A bare shape: a node kind and its children shapes. +data Shape = Shape SpanKind [Shape] + +genShape :: Int -> Gen Shape +genShape d = do + k <- elements [ProgramSpan, ModuleSpan, CombinatorSpan, LlmCallSpan] + kids <- + if d <= 0 + then pure [] + else do + n <- choose (0, 3) + vectorOf n (genShape (d - 1)) + pure (Shape k kids) + +-- | Generate a random (single-rooted) tree of up to depth 3, with sequential span +-- ids, integer-second timestamps (exact JSON round-trip), and LM-call leaves +-- carrying a cacheKey + response so the serialized form exercises every field. +genTree :: Gen TraceTree +genTree = do + shape <- genShape 3 + let (ss, _) = flattenShape Nothing 0 shape + m = Map.fromList [(spanId s, s) | s <- ss] + rootId = spanId (firstSpan ss) + pure (TraceTree rootId m) + where + firstSpan (s : _) = s + firstSpan [] = error "genTree: empty" + +-- | Flatten a shape into spans, assigning ids/parents/times from a running +-- counter. Returns the produced spans (root first) and the next free counter. +flattenShape :: Maybe SpanId -> Int -> Shape -> ([Span], Int) +flattenShape par n (Shape k kids) = + let sid = SpanId ("span-" <> T.pack (show n)) + s = + Span + { spanId = sid, + parent = par, + kind = k, + label = "lbl-" <> T.pack (show n), + startedAt = addUTCTime (fromIntegral (2 * n)) baseTime, + endedAt = Just (addUTCTime (fromIntegral (2 * n + 1)) baseTime), + attrs = attrsFor k n + } + (kidSpans, n') = foldl step ([], n + 1) kids + step (acc, cnt) kidShape = + let (ks, cnt') = flattenShape (Just sid) cnt kidShape + in (acc ++ ks, cnt') + in (s : kidSpans, n') + +-- | LM-call leaves carry a key + response (exercising replayIndex and the Value +-- fields); structural nodes are empty. +attrsFor :: SpanKind -> Int -> SpanAttrs +attrsFor LlmCallSpan n = + emptyAttrs + { model = Just "m", + cacheKey = Just ("key-" <> T.pack (show n)), + response = Just (object ["text" .= ("resp-" <> T.pack (show n))]), + inputTokens = Just (fromIntegral n), + costUsd = Just (read "0.001") + } +attrsFor _ _ = emptyAttrs + +-- --------------------------------------------------------------------------- +-- EP-16: node identity, correlation, and the feedback channel +-- --------------------------------------------------------------------------- + +-- | A self-looping record so @predict cellSig :: Program Cell Cell@ composes with +-- itself to make a multi-node program. +newtype Cell = Cell {cell :: Text} + deriving stock (Generic, Show, Eq) + +instance ToSchema Cell + +instance FromModel Cell + +instance ToPrompt Cell + +instance Validatable Cell + +cellSig :: Signature Cell Cell +cellSig = mkSignature "Echo the cell" + +-- | A two-@Predict@ chain: @Compose (Predict) (Predict)@. +chain2 :: Program Cell Cell +chain2 = predict cellSig >>> predict cellSig + +-- | A fixed response in fallback @[[ ## field ## ]]@ form, decoded by the lenient +-- parser into @Cell "echoed"@. +cellResp :: Response +cellResp = mkResponse "[[ ## cell ## ]]\nechoed\n[[ ## completed ## ]]\n" + +chkLen :: Program a b -> IO () +chkLen p = length (programNodePaths p) @?= length (foldParams p) + +-- | A total ensemble reducer (the program is only shape-inspected, never run). +firstOr :: [Cell] -> Cell +firstOr (c : _) = c +firstOr [] = Cell "" + +-- | The @LlmCallSpan@s of a tree, in start order. +llmSpansInOrder :: TraceTree -> [Span] +llmSpansInOrder t = sortOn startedAt [s | s <- Map.elems (spans t), kind s == LlmCallSpan] + +-- | Run a program under 'runProgramTraced' with the node-aware stack, returning the +-- decoded result and the captured tree (or a typed error). +runTraced :: Program Cell Cell -> Cell -> IO (Either ShikumiError (Cell, TraceTree)) +runTraced prog input = + runEff + . runErrorNoCallStack @ShikumiError + . runPrim + . runTime + . runTrace + . runCurrentNode + . runFixedLLM cellResp + . tracedNodeLLM + $ runProgramTraced prog input + +-- | M1: 'NodePath' identity, its agreement with @foldParams@, field metadata, and +-- JSON round-trips. +nodeTests :: TestTree +nodeTests = + testGroup + "node" + [ testCase "programNodePaths length agrees with foldParams across shapes" $ do + chkLen chain2 + chkLen (majorityVote 3 (TempFixed []) (predict cellSig)) + chkLen (majorityVoteBy 2 (TempFixed []) firstOr (predict cellSig)) + chkLen (parallel2 (predict cellSig) (predict cellSig)), + testCase "programNodePaths yields the expected structural paths for a chain" $ + programNodePaths chain2 @?= [NodePath [StepComposeL], NodePath [StepComposeR]], + testCase "nodeFields returns each node's input/output field names" $ do + let nf = map snd (nodeFields chain2) + map inputFieldNames nf @?= [["cell"], ["cell"]] + map outputFieldNames nf @?= [["cell"], ["cell"]], + testCase "mapParamsAt k edits the node at programNodePaths !! k (index law)" $ do + let edited = mapParamsAt 1 (\ps -> ps {instructionOverride = Just "x"}) chain2 + map instructionOverride (foldParams edited) @?= [Nothing, Just "x"], + testCase "NodePath and a nodePath-bearing SpanAttrs round-trip as JSON" $ do + let np = NodePath [StepComposeL, StepEnsemble 2] + decode (encode np) @?= Just np + let a = emptyAttrs {nodePath = Just np} + decode (encode a) @?= Just a + ] + +-- | M2: each model-call span is tagged with its issuing node's 'NodePath', and the +-- nodePath-bearing tree round-trips through the store at formatVersion 2. +correlateTests :: TestTree +correlateTests = + testGroup + "correlate" + [ testCase "each llm-call span carries its node's path, in programNodePaths order" $ do + res <- runTraced chain2 (Cell "x") + case res of + Left e -> assertFailure ("traced run failed: " <> show e) + Right (out, tree) -> do + out @?= Cell "echoed" + let llms = llmSpansInOrder tree + length llms @?= 2 + let paths = programNodePaths chain2 + map (nodePath . attrs) llms @?= [Just (paths !! 0), Just (paths !! 1)], + testCase "a nodePath-bearing tree round-trips through the store (formatVersion 2)" $ + withSystemTempDirectory "shikumi-trace-node" $ \dir -> do + res <- runTraced chain2 (Cell "x") + case res of + Left e -> assertFailure ("traced run failed: " <> show e) + Right (_, tree) -> do + let p = dir <> "/trace.json" + writeTraceFile p tree + loaded <- readTraceFile p + loaded @?= Right tree + ] + +-- | M3: the per-node feedback channel round-trips writes and reads. +feedbackTests :: TestTree +feedbackTests = + testGroup + "feedback" + [ testCase "feedbackFor returns critiques in attach order; untargeted node is empty" $ do + let p1 = NodePath [StepComposeL] + p2 = NodePath [StepComposeR] + p3 = NodePath [StepFMap] + (_, fb) <- + runEff . runPrim . runFeedback $ do + attachFeedback p1 "first" + attachFeedback p1 "second" + attachFeedback p2 "only" + feedbackFor p1 fb @?= ["first", "second"] + feedbackFor p2 fb @?= ["only"] + feedbackFor p3 fb @?= [], + testCase "FeedbackLog round-trips as JSON" $ do + (_, fb) <- + runEff . runPrim . runFeedback $ + attachFeedback (NodePath [StepComposeL]) "c" + decode (encode (fb :: FeedbackLog)) @?= Just fb + ] diff --git a/test/TraceFixtures.hs b/test/TraceFixtures.hs new file mode 100644 --- /dev/null +++ b/test/TraceFixtures.hs @@ -0,0 +1,103 @@ +-- | Shared, network-free fixtures for the EP-7 trace tests: a stub 'Model', a +-- request builder, canned responses, and two base @LLM@ interpreters — a /fixed/ +-- one that returns the same response for every call, and a /keyed/ one that picks +-- a response from the request 'Context' (so a two-stage pipeline gets distinct +-- responses, hence distinct cache keys). +module TraceFixtures + ( stubModel, + ctxFor, + optsFor, + mkResponse, + responseText, + lastUserText, + runFixedLLM, + runKeyedLLM, + runKeyedCountingLLM, + ) +where + +import Baikai + ( Api (Custom), + AssistantContent (..), + Context, + Message (UserMessage), + Model, + Options, + Response, + TextContent (..), + UserContent (UserText), + user, + _Context, + _Model, + _Options, + _Response, + _TextContent, + ) +import Control.Lens ((&), (.~), (^.)) +import Data.Generics.Labels () +import Data.IORef (IORef, atomicModifyIORef') +import Data.Text (Text) +import Data.Vector qualified as V +import Effectful (Eff, IOE, liftIO, (:>)) +import Effectful.Dispatch.Dynamic (interpret) +import Shikumi.LLM (LLM (..)) + +-- | A stub model with real-looking routing identity, so LM-call span labels read +-- @stub/stub-model@. +stubModel :: Model +stubModel = + _Model + & #provider .~ "stub" + & #modelId .~ "stub-model" + & #api .~ Custom "stub" + +-- | A one-user-turn request context carrying @t@ (varies the cache key per stage). +ctxFor :: Text -> Context +ctxFor t = _Context & #messages .~ V.singleton (user t) + +-- | Default options. +optsFor :: Options +optsFor = _Options + +-- | A response carrying @t@ as its single assistant text block, with small fixed +-- token counts and latency so span attributes are populated. +mkResponse :: Text -> Response +mkResponse t = + _Response + & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t)) + & #message . #usage . #inputTokens .~ 10 + & #message . #usage . #outputTokens .~ 5 + & #latencyMs .~ 12 + +-- | Concatenate the assistant text blocks of a response. +responseText :: Response -> Text +responseText resp = + mconcat [t | AssistantText (TextContent t) <- V.toList (resp ^. #message . #content)] + +-- | Concatenate the text of every user message in a request context. +lastUserText :: Context -> Text +lastUserText c = + mconcat + [ t + | UserMessage up <- V.toList (c ^. #messages), + UserText (TextContent t) <- V.toList (up ^. #content) + ] + +-- | A base @LLM@ interpreter that returns the same response for every completion. +runFixedLLM :: Response -> Eff (LLM : es) a -> Eff es a +runFixedLLM resp = interpret $ \_ -> \case + Complete {} -> pure resp + Stream {} -> pure [] + +-- | A base @LLM@ interpreter that picks a response from the request context. +runKeyedLLM :: (Context -> Response) -> Eff (LLM : es) a -> Eff es a +runKeyedLLM f = interpret $ \_ -> \case + Complete _ c _ -> pure (f c) + Stream {} -> pure [] + +-- | Like 'runKeyedLLM' but bumps a counter on every completion (used to assert how +-- many provider calls a live run made — and that a replay run makes zero). +runKeyedCountingLLM :: (IOE :> es) => IORef Int -> (Context -> Response) -> Eff (LLM : es) a -> Eff es a +runKeyedCountingLLM ref f = interpret $ \_ -> \case + Complete _ c _ -> liftIO (atomicModifyIORef' ref (\n -> (n + 1, ()))) >> pure (f c) + Stream {} -> pure []