diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,14 @@ ## Unreleased +## 0.1.1.0 - 2026-06-28 + +### Changed + +- Refreshed internal `shikumi` and `shikumi-cache` bounds for the current package + set. +- Updated tracing and replay internals to use label-based record access. + ## 0.1.0.1 - 2026-06-21 ### Changed diff --git a/shikumi-trace.cabal b/shikumi-trace.cabal --- a/shikumi-trace.cabal +++ b/shikumi-trace.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: shikumi-trace -version: 0.1.0.1 +version: 0.1.1.0 synopsis: Hierarchical tracing, observability, and deterministic replay for shikumi (EP-7) @@ -60,8 +60,8 @@ , generic-lens , lens ^>=5.3 , scientific - , shikumi ^>=0.1.0.1 - , shikumi-cache ^>=0.1.0.1 + , shikumi ^>=0.2.0.0 + , shikumi-cache ^>=0.1.1.0 , text ^>=2.1 , time , vector @@ -73,7 +73,7 @@ ghc-options: -threaded -with-rtsopts=-N build-depends: , base - , shikumi-trace ^>=0.1.0.1 + , shikumi-trace ^>=0.1.1.0 test-suite shikumi-trace-test import: common-options @@ -92,9 +92,9 @@ , generic-lens , lens ^>=5.3 , QuickCheck - , shikumi ^>=0.1.0.1 - , shikumi-cache ^>=0.1.0.1 - , shikumi-trace ^>=0.1.0.1 + , shikumi ^>=0.2.0.0 + , shikumi-cache ^>=0.1.1.0 + , shikumi-trace ^>=0.1.1.0 , tasty , tasty-hunit , tasty-quickcheck diff --git a/src/Shikumi/Trace.hs b/src/Shikumi/Trace.hs --- a/src/Shikumi/Trace.hs +++ b/src/Shikumi/Trace.hs @@ -60,7 +60,7 @@ Response, flattenAssistantBlocks, ) -import Control.Lens ((^.)) +import Control.Lens (at, ix, (%~), (&), (.~), (?~), (^.)) import Data.Aeson (FromJSON, FromJSONKey, ToJSON, ToJSONKey, Value, toJSON) import Data.Generics.Labels () import Data.List (sortOn) @@ -187,9 +187,9 @@ -- 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] + map (^. #spanId) $ + sortOn (\s -> (s ^. #startedAt, s ^. #spanId)) $ + [s | s <- Map.elems (t ^. #spans), (s ^. #parent) == Just sid] -- --------------------------------------------------------------------------- -- The effect @@ -233,11 +233,12 @@ -- | 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)) + { counter :: !(IORef Int), + stack :: !(IORef [SpanId]), + spans :: !(IORef (Map SpanId Span)), + root :: !(IORef (Maybe SpanId)) } + deriving stock (Generic) -- | Run a traced computation, returning its result and the finished tree. -- @@ -251,9 +252,9 @@ 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]}) + CurrentSpanId -> safeHead <$> readIORef (st ^. #stack) + BumpRetry -> modifyActive st (\a' -> a' & #retries %~ (+ 1)) + RecordToolCall tc -> modifyActive st (\a' -> a' & #toolCalls %~ (++ [tc])) AnnotateSpan f -> modifyActive st f WithSpan k lbl inner -> bracket @@ -274,15 +275,15 @@ -- 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)) + n <- atomicModifyIORef' (st ^. #counter) (\i -> (i + 1, i + 1)) let sid = SpanId ("span-" <> T.pack (show n)) - par <- safeHead <$> readIORef (tsStack st) + par <- safeHead <$> readIORef (st ^. #stack) now <- getCurrentTime let s = Span sid par k lbl now Nothing emptyAttrs - modifyIORef' (tsSpans st) (Map.insert sid s) - modifyIORef' (tsStack st) (sid :) + modifyIORef' (st ^. #spans) (at sid ?~ s) + modifyIORef' (st ^. #stack) (sid :) case par of - Nothing -> modifyIORef' (tsRoot st) (Just . fromMaybe sid) + Nothing -> modifyIORef' (st ^. #root) (Just . fromMaybe sid) Just _ -> pure () pure sid @@ -290,22 +291,22 @@ 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) + modifyIORef' (st ^. #spans) (ix sid . #endedAt ?~ now) + modifyIORef' (st ^. #stack) (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) + stk <- readIORef (st ^. #stack) case stk of - (sid : _) -> modifyIORef' (tsSpans st) (Map.adjust (\s -> s {attrs = f (attrs s)}) sid) + (sid : _) -> modifyIORef' (st ^. #spans) (ix sid . #attrs %~ f) [] -> 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) + sp <- readIORef (st ^. #spans) + r <- readIORef (st ^. #root) pure (TraceTree (fromMaybe (SpanId "") r) sp) safeHead :: [a] -> Maybe a @@ -337,17 +338,16 @@ 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)) - } + & #model ?~ (m ^. #modelId) + & #provider ?~ (m ^. #provider) + & #prompt ?~ requestToCanonicalValue m c o + & #response ?~ toJSON resp + & #latencyMs ?~ (resp ^. #latencyMs) + & #inputTokens ?~ (resp ^. #message . #usage . #inputTokens) + & #outputTokens ?~ (resp ^. #message . #usage . #outputTokens) + & #costUsd ?~ realToFrac (resp ^. #message . #usage . #cost . #usd :: Rational) + & #toolCalls .~ toolCallsOf resp + & #cacheKey ?~ unCacheKey (Key.cacheKey m c o) -- | Extract the tool calls from a response's assistant content blocks. toolCallsOf :: Response -> [ToolCallRecord] @@ -364,10 +364,10 @@ -- 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)) + | Map.null (t ^. #spans) = "(empty trace)\n" + | otherwise = T.concat (go (0 :: Int) (t ^. #root)) where - go depth sid = case Map.lookup sid (spans t) of + go depth sid = case Map.lookup sid (t ^. #spans) of Nothing -> [] Just s -> line depth s : concatMap (go (depth + 1)) (childrenOf t sid) line depth s = diff --git a/src/Shikumi/Trace/Program.hs b/src/Shikumi/Trace/Program.hs --- a/src/Shikumi/Trace/Program.hs +++ b/src/Shikumi/Trace/Program.hs @@ -42,6 +42,8 @@ ) where +import Control.Lens ((&), (.~)) +import Data.Generics.Labels () import Effectful (Dispatch (Dynamic), DispatchOf, Eff, Effect, (:>)) import Effectful.Dispatch.Dynamic (interpose, interpret, localSeqUnlift, send) import Effectful.Error.Static (Error) @@ -72,8 +74,7 @@ withSampleTemp, ) import Shikumi.Trace - ( SpanAttrs (..), - SpanKind (CombinatorSpan, LlmCallSpan, ModuleSpan), + ( SpanKind (CombinatorSpan, LlmCallSpan, ModuleSpan), Trace, annotateSpan, llmAttrs, @@ -137,7 +138,7 @@ Complete m c o -> withSpan LlmCallSpan (llmLabel m) $ do resp <- complete m c o mp <- askNode - annotateSpan (\_ -> (llmAttrs m c o resp) {nodePath = mp}) + annotateSpan (\_ -> llmAttrs m c o resp & #nodePath .~ mp) pure resp Stream m c o -> withSpan LlmCallSpan (llmLabel m) (stream m c o) diff --git a/src/Shikumi/Trace/Replay.hs b/src/Shikumi/Trace/Replay.hs --- a/src/Shikumi/Trace/Replay.hs +++ b/src/Shikumi/Trace/Replay.hs @@ -24,7 +24,7 @@ where import Baikai (Context, Model, Response) -import Control.Lens ((^.)) +import Control.Lens ((&), (.~), (^.)) import Data.Aeson (Result (Error, Success), Value, fromJSON) import Data.Generics.Labels () import Data.Map.Strict (Map) @@ -68,11 +68,11 @@ 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} + 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"} + 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 diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -471,12 +471,12 @@ 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 + 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} + let a = emptyAttrs & #nodePath .~ Just np decode (encode a) @?= Just a ]