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 typed tool registries and ReAct-style agents for shikumi programs. +- CodeAct, program-of-thought helpers, prompt utilities, and a restricted deterministic code interpreter for offline tests. diff --git a/shikumi-tools.cabal b/shikumi-tools.cabal new file mode 100644 --- /dev/null +++ b/shikumi-tools.cabal @@ -0,0 +1,97 @@ +cabal-version: 3.4 +name: shikumi-tools +version: 0.1.0.0 +synopsis: + Typed tools and ReAct agents for shikumi LM programs (EP-11) + +category: AI +description: + Typed tools and multi-step ReAct agents for shikumi. A @Tool i o@ is an ordinary + function over record types whose JSON-Schema argument shape is Generic-derived + (reusing the schema generator from the signatures plan) and which lowers to + baikai's untyped wire @Tool@ (MasterPlan integration point #8). A heterogeneous + @ToolRegistry@ erases the per-tool @i@/@o@ behind an existential while retaining + the decode/encode dictionaries, so a ReAct loop can dispatch by the name the model + emits at run time. On top of typed tools, @react@/@reactWithTrajectory@ build a + @Program i o@ whose embedded loop alternates propose -> dispatch -> observe until + the model finishes or a bound is hit, then extracts the typed answer — recording a + structured @Trajectory@ throughout. The whole agent is a first-class, composable, + inspectable shikumi program. Every test runs against a deterministic mock LM with + no network. + +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.Agent.ReAct + Shikumi.CodeExec.CodeAct + Shikumi.CodeExec.Interpreter + Shikumi.CodeExec.ProgramOfThought + Shikumi.CodeExec.Prompt + Shikumi.Tool + + build-depends: + , aeson + , baikai + , base >=4.20 && <5 + , bytestring + , containers + , effectful + , generic-lens + , lens ^>=5.3 + , shikumi ^>=0.1.0.0 + , text ^>=2.1 + , vector + +test-suite shikumi-tools-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: + AcceptanceSpec + CodeActSpec + Fixtures + MockLLM + ProgramOfThoughtSpec + ProtocolSpec + ReActSpec + RestrictedSpec + SchemaSpec + ToolSpec + + build-depends: + , aeson + , baikai + , base + , containers + , effectful + , generic-lens + , lens + , shikumi ^>=0.1.0.0 + , shikumi-tools ^>=0.1.0.0 + , tasty + , tasty-hunit + , text + , vector diff --git a/src/Shikumi/Agent/ReAct.hs b/src/Shikumi/Agent/ReAct.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Agent/ReAct.hs @@ -0,0 +1,475 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | ReAct ("Reason + Act") agents as first-class shikumi programs. +-- +-- A ReAct agent alternates between /thinking/ and /acting/: at each step the model +-- emits a short thought and either selects a tool with arguments or declares it is +-- finished. The runtime executes the selected tool (via "Shikumi.Tool"), records the +-- result as an /observation/, and asks the model again — until the model finishes or +-- a @maxIters@ bound is hit. A final /extract/ step turns the finished +-- 'Trajectory' into the typed output the caller wanted. +-- +-- The whole loop is wrapped in 'Shikumi.Program.Embed', so @react@ is an ordinary +-- @'Program' i o@: it runs under 'Shikumi.Program.runProgram' (the body needs only +-- @LLM@ + @Error ShikumiError@, integration point #4's exact row), is structurally +-- inspectable, and composes with every other shikumi program and combinator. +-- +-- Two tool /protocols/ are supported behind one internal interface ('ProtocolImpl'), +-- selected by a 'ToolProtocol' value ('ProtocolAuto' resolves per model via +-- "Shikumi.Adapter"'s 'capabilityFor'): a provider-native function-calling path +-- (baikai @Context.tools@ + @Options.toolChoice@, parsing @AssistantToolCall@ blocks) +-- and a prompt-based fallback that renders an explicit action grammar and parses the +-- model's text. The loop body is protocol-agnostic. +module Shikumi.Agent.ReAct + ( -- * The trajectory data model + Action (..), + Step (..), + Termination (..), + Trajectory (..), + + -- * Configuration + ToolProtocol (..), + ReActConfig (..), + defaultReActConfig, + + -- * Building agents + react, + reactWithTrajectory, + + -- * The protocol seam + ProtocolImpl (..), + resolveProtocolKind, + resolveProtocol, + + -- * Rendering (exposed for tests/inspection) + renderTrajectory, + ) +where + +import Baikai + ( AssistantContent (..), + Context, + Message, + Model, + Options, + Response, + TextContent (..), + Tool, + ToolCall, + ToolChoice (..), + flattenAssistantBlocks, + user, + _Context, + _Model, + _Options, + _ToolCall, + ) +import Control.Lens ((&), (.~), (^.)) +import Data.Aeson (Value (..), eitherDecodeStrict, encode) +import Data.Aeson.KeyMap qualified as KM +import Data.ByteString.Lazy qualified as LBS +import Data.Generics.Labels () +import Data.Proxy (Proxy (..)) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Encoding (decodeUtf8, encodeUtf8) +import Data.Vector (Vector) +import Data.Vector qualified as V +import Effectful (Eff, (:>)) +import Effectful.Error.Static (Error, throwError) +import GHC.Generics (Generic) +import Shikumi.Adapter (ModelCapability (..), ToPrompt (toPrompt), attachSchema, capabilityFor) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.LLM (LLM, complete) +import Shikumi.Program (Program (FMap), embed) +import Shikumi.Schema (FromModel, ToSchema, Validatable, parseOutput, toSchema) +import Shikumi.Signature (Signature, getInstruction) +import Shikumi.Tool (ToolRegistry, registryBaikai, registryTools, renderToolError, runToolCall, someToolDescription, someToolName, someToolSchema) + +-- --------------------------------------------------------------------------- +-- The trajectory data model +-- --------------------------------------------------------------------------- + +-- | What the model decided to do at a step: invoke a named tool with a raw JSON +-- arguments object, or declare it is finished. +data Action + = CallTool !Text !Value + | Finish + deriving stock (Show, Eq, Generic) + +-- | One recorded (thought, action, observation) step. @observation@ is 'Nothing' +-- for 'Finish'; for a tool call it is the tool result text or the rendered +-- 'Shikumi.Tool.ToolError'. +data Step = Step + { thought :: !Text, + action :: !Action, + observation :: !(Maybe Text) + } + deriving stock (Show, Eq, Generic) + +-- | Why the loop stopped. 'TerminatedBudget' is retained for forward +-- compatibility; the loop itself produces only 'TerminatedFinish' and +-- 'TerminatedMaxIters' (the budget ceiling is enforced one layer down by the +-- resilient @LLM@ interpreter, surfacing as a 'ShikumiError'). +data Termination + = TerminatedFinish + | TerminatedMaxIters !Int + | TerminatedBudget + deriving stock (Show, Eq, Generic) + +-- | The recorded sequence of steps plus the reason the loop stopped. +data Trajectory = Trajectory + { steps :: !(Vector Step), + termination :: !Termination + } + deriving stock (Show, Eq, Generic) + +-- --------------------------------------------------------------------------- +-- Configuration +-- --------------------------------------------------------------------------- + +-- | Which tool-calling protocol to use. 'ProtocolAuto' resolves per model. +data ToolProtocol + = ProtocolNative + | ProtocolPrompt + | ProtocolAuto + deriving stock (Show, Eq, Generic) + +-- | How an agent runs: the hard iteration cap and the protocol selector. +data ReActConfig = ReActConfig + { maxIters :: !Int, + protocol :: !ToolProtocol + } + deriving stock (Show, Eq, Generic) + +-- | Six iterations, protocol auto-selected. +defaultReActConfig :: ReActConfig +defaultReActConfig = ReActConfig {maxIters = 6, protocol = ProtocolAuto} + +-- --------------------------------------------------------------------------- +-- Building agents +-- --------------------------------------------------------------------------- + +-- | Build a ReAct agent as a @'Program' i o@ that returns the typed answer and +-- drops the trajectory. The ergonomic default. +react :: + forall i o. + (ToPrompt i, ToSchema o, FromModel o, Validatable o) => + Signature i o -> + ToolRegistry -> + ReActConfig -> + Program i o +react sig reg cfg = FMap fst (reactWithTrajectory sig reg cfg) + +-- | Build a ReAct agent that also returns the recorded 'Trajectory', for +-- evaluators/optimizers and tests that assert on the steps. +reactWithTrajectory :: + forall i o. + (ToPrompt i, ToSchema o, FromModel o, Validatable o) => + Signature i o -> + ToolRegistry -> + ReActConfig -> + Program i (o, Trajectory) +reactWithTrajectory sig reg cfg = embed (reactLoop sig reg cfg) + +-- | The agent loop, embedded into a 'Program' by 'reactWithTrajectory'. Runs in +-- exactly 'Shikumi.Program.runProgram'\'s effect row. +reactLoop :: + forall i o es. + (LLM :> es, Error ShikumiError :> es, ToPrompt i, ToSchema o, FromModel o, Validatable o) => + Signature i o -> + ToolRegistry -> + ReActConfig -> + i -> + Eff es (o, Trajectory) +reactLoop sig reg cfg i = do + traj <- loop 0 [] + o <- extract traj + pure (o, traj) + where + impl :: ProtocolImpl i o + impl = resolveProtocol (protocol cfg) _Model sig reg + + -- Propose -> dispatch -> observe, accumulating steps (newest first). + loop :: Int -> [Step] -> Eff es Trajectory + loop iter acc + | iter >= maxIters cfg = + pure (Trajectory (V.fromList (reverse acc)) (TerminatedMaxIters iter)) + | otherwise = do + let (ctx, opts) = renderPropose impl i (soFar acc) + resp <- complete _Model ctx opts + case parsePropose impl resp of + Left perr -> loop (iter + 1) (correctiveStep perr : acc) + Right (th, Finish) -> + pure (Trajectory (V.fromList (reverse (Step th Finish Nothing : acc))) TerminatedFinish) + Right (th, CallTool nm args) -> do + res <- runToolCall reg (mkToolCall nm args) + let obs = either renderToolError id res + loop (iter + 1) (Step th (CallTool nm args) (Just obs) : acc) + + -- The final extract call: render, issue, decode into @o@ (a decode failure + -- here is the agent's final-answer failure, surfaced as a 'ShikumiError'). + extract :: Trajectory -> Eff es o + extract traj = do + let (ctx, opts) = renderExtract impl i traj + resp <- complete _Model ctx opts + either throwError pure (parseExtract impl resp) + + -- A trajectory view of the steps gathered so far (termination is irrelevant + -- for the propose render). + soFar :: [Step] -> Trajectory + soFar acc = Trajectory (V.fromList (reverse acc)) TerminatedFinish + +-- | Build a synthetic step recording an unparseable proposal, so the corrective +-- text reaches the model on the next turn (the empty tool name is a sentinel that +-- is never dispatched). +correctiveStep :: Text -> Step +correctiveStep perr = + Step + { thought = "", + action = CallTool "" Null, + observation = + Just ("Your previous reply was not a valid action JSON object: " <> perr <> ". Reply with exactly one action object.") + } + +-- | A baikai 'ToolCall' built from a name and a raw arguments object (the call id +-- is irrelevant on the prompt path and synthesized on the native path). +mkToolCall :: Text -> Value -> ToolCall +mkToolCall nm args = _ToolCall & #name .~ nm & #arguments .~ args + +-- --------------------------------------------------------------------------- +-- The protocol seam +-- --------------------------------------------------------------------------- + +-- | The protocol-specific rendering/parsing the loop depends on. Both the native +-- and the prompt implementations satisfy this interface, so the loop body is +-- identical under either seam. +data ProtocolImpl i o = ProtocolImpl + { renderPropose :: i -> Trajectory -> (Context, Options), + parsePropose :: Response -> Either Text (Text, Action), + renderExtract :: i -> Trajectory -> (Context, Options), + parseExtract :: Response -> Either ShikumiError o + } + +-- | Resolve a (possibly 'ProtocolAuto') selection to a concrete kind for a model. +-- Conservative: a model without native tool support resolves to 'ProtocolPrompt', +-- because a silently-dropped native tool set (baikai does this for CLI providers) +-- would make the loop propose tools that never execute. +resolveProtocolKind :: ToolProtocol -> Model -> ToolProtocol +resolveProtocolKind ProtocolNative _ = ProtocolNative +resolveProtocolKind ProtocolPrompt _ = ProtocolPrompt +resolveProtocolKind ProtocolAuto m = case capabilityFor m of + NativeSchema -> ProtocolNative + PromptFallback -> ProtocolPrompt + +-- | Build the concrete 'ProtocolImpl' for a model, choosing the native or prompt +-- renderers from 'resolveProtocolKind'. +resolveProtocol :: + forall i o. + (ToPrompt i, ToSchema o, FromModel o, Validatable o) => + ToolProtocol -> + Model -> + Signature i o -> + ToolRegistry -> + ProtocolImpl i o +resolveProtocol proto m sig reg = case resolveProtocolKind proto m of + ProtocolNative -> nativeImpl sig reg + _ -> promptImpl sig reg + +-- --------------------------------------------------------------------------- +-- The prompt protocol +-- --------------------------------------------------------------------------- + +promptImpl :: + forall i o. + (ToPrompt i, ToSchema o, FromModel o, Validatable o) => + Signature i o -> + ToolRegistry -> + ProtocolImpl i o +promptImpl sig reg = + ProtocolImpl + { renderPropose = \i traj -> + let sys = + getInstruction sig + <> "\n\n" + <> toolMenu reg + <> "\n" + <> proposeGrammar + msg = user (taskBlock i <> "\n\n" <> historyBlock traj) + in (buildCtx sys [msg] V.empty Nothing, _Options), + parsePropose = \resp -> parseActionText (responseText resp), + renderExtract = \i traj -> + let sys = + getInstruction sig + <> "\n\n" + <> extractGuide (toSchema (Proxy @o)) + msg = user (taskBlock i <> "\n\n" <> historyBlock traj) + in (buildCtx sys [msg] V.empty Nothing, _Options), + parseExtract = \resp -> parseOutput (stripFences (responseText resp)) + } + +-- --------------------------------------------------------------------------- +-- The native protocol +-- --------------------------------------------------------------------------- + +nativeImpl :: + forall i o. + (ToPrompt i, ToSchema o, FromModel o, Validatable o) => + Signature i o -> + ToolRegistry -> + ProtocolImpl i o +nativeImpl sig reg = + ProtocolImpl + { renderPropose = \i traj -> + let sys = + getInstruction sig + <> "\n\nUse a tool when you need one, or answer directly when you have enough information." + msg = user (taskBlock i <> "\n\n" <> historyBlock traj) + opts = _Options & #toolChoice .~ Just ToolChoiceAuto + in (buildCtx sys [msg] (registryBaikai reg) Nothing, opts), + -- A tool-call block -> CallTool; no tool call (plain text) -> Finish. + parsePropose = \resp -> + case toolCallsOf resp of + (tc : _) -> Right (responseText resp, CallTool (tc ^. #name) (tc ^. #arguments)) + [] -> Right (responseText resp, Finish), + renderExtract = \i traj -> + let sys = + getInstruction sig + <> "\n\n" + <> extractGuide (toSchema (Proxy @o)) + msg = user (taskBlock i <> "\n\n" <> historyBlock traj) + -- Suppress tools so the model answers; attach the output schema (a + -- no-op until EP-2's responseFormat is wired in the local baikai). + opts = + attachSchema (toSchema (Proxy @o)) _Options + & #toolChoice .~ Just ToolChoiceNone + in (buildCtx sys [msg] V.empty Nothing, opts), + parseExtract = \resp -> parseOutput (stripFences (responseText resp)) + } + +-- --------------------------------------------------------------------------- +-- Rendering helpers +-- --------------------------------------------------------------------------- + +buildCtx :: Text -> [Message] -> Vector Tool -> Maybe ToolChoice -> Context +buildCtx sys msgs tools _ = + _Context + & #systemPrompt .~ Just sys + & #messages .~ V.fromList msgs + & #tools .~ tools + +taskBlock :: (ToPrompt i) => i -> Text +taskBlock i = "Task:\n" <> toPrompt i + +-- | Render the steps gathered so far as a readable history. +historyBlock :: Trajectory -> Text +historyBlock traj + | V.null (steps traj) = "Steps so far: (none)" + | otherwise = "Steps so far:\n" <> renderTrajectory traj + +-- | A human-readable rendering of a trajectory's steps. +renderTrajectory :: Trajectory -> Text +renderTrajectory traj = T.intercalate "\n" (zipWith one [1 :: Int ..] (V.toList (steps traj))) + where + one n s = + " " + <> T.pack (show n) + <> ". thought: " + <> thought s + <> "; action: " + <> renderAction (action s) + <> maybe "" ("; observation: " <>) (observation s) + renderAction Finish = "finish" + renderAction (CallTool nm args) = "call " <> nm <> " " <> encodeText args + +-- | The tool menu: each tool's name, description, and compact argument schema. +toolMenu :: ToolRegistry -> Text +toolMenu reg = case registryTools reg of + [] -> "You have no tools available." + ts -> "You can use these tools:\n" <> T.unlines (map one ts) + where + one st = + "- " + <> someToolName st + <> ": " + <> someToolDescription st + <> " (arguments schema: " + <> encodeText (someToolSchema st) + <> ")" + +-- | The action grammar instruction for the prompt protocol. +proposeGrammar :: Text +proposeGrammar = + T.unlines + [ "Reply with exactly one JSON object and nothing else, in one of these two forms:", + " {\"thought\": \"...\", \"action\": {\"tool\": \"\", \"args\": { ... }}}", + " {\"thought\": \"...\", \"action\": {\"finish\": true}}" + ] + +-- | The extract instruction: ask for one JSON object matching @o@'s schema. +extractGuide :: Value -> Text +extractGuide schema = + "Now produce the final answer as exactly one JSON object matching this schema, and nothing else:\n" + <> encodeText schema + +-- --------------------------------------------------------------------------- +-- Parsing helpers +-- --------------------------------------------------------------------------- + +-- | The concatenated text of every assistant text block. +responseText :: Response -> Text +responseText resp = + T.concat [t | AssistantText (TextContent t) <- V.toList (flattenAssistantBlocks resp)] + +-- | Every tool-call block in a response. +toolCallsOf :: Response -> [ToolCall] +toolCallsOf resp = [tc | AssistantToolCall tc <- V.toList (flattenAssistantBlocks resp)] + +-- | Parse the model's propose text into @(thought, action)@. Strips code fences, +-- decodes one JSON object, and reads its @thought@ and @action@ fields. +parseActionText :: Text -> Either Text (Text, Action) +parseActionText raw = + case eitherDecodeStrict (encodeUtf8 (stripFences raw)) of + Left e -> Left ("not valid JSON: " <> T.pack e) + Right v -> parseActionObject v + +-- | Read @{thought, action}@ from a decoded JSON value. +parseActionObject :: Value -> Either Text (Text, Action) +parseActionObject (Object o) = do + th <- case KM.lookup "thought" o of + Just (String t) -> Right t + Just _ -> Left "\"thought\" must be a string" + Nothing -> Right "" -- thought is optional + act <- case KM.lookup "action" o of + Just a -> parseAction a + Nothing -> Left "missing \"action\"" + pure (th, act) +parseActionObject _ = Left "expected a JSON object" + +-- | Read an action object: @{finish:true}@ -> 'Finish'; @{tool,args}@ -> 'CallTool'. +parseAction :: Value -> Either Text Action +parseAction (Object o) + | Just (Bool True) <- KM.lookup "finish" o = Right Finish + | Just (String nm) <- KM.lookup "tool" o = + Right (CallTool nm (maybe (Object KM.empty) id (KM.lookup "args" o))) + | otherwise = Left "action must be {\"finish\":true} or {\"tool\":..,\"args\":..}" +parseAction _ = Left "action must be a JSON object" + +-- | Strip a leading/trailing Markdown code fence (```… / ```json … ```), if any, +-- so a fenced JSON reply still decodes. Falls back to the input unchanged. +stripFences :: Text -> Text +stripFences t = + let trimmed = T.strip t + in case T.stripPrefix "```" trimmed of + Nothing -> trimmed + Just rest -> + let afterLang = dropFirstLine rest + in T.strip (dropClosingFence afterLang) + where + dropFirstLine s = T.drop 1 (T.dropWhile (/= '\n') s) + dropClosingFence s = fst (T.breakOn "```" s) + +-- | Compact-encode a JSON value to text. +encodeText :: Value -> Text +encodeText = decodeUtf8 . LBS.toStrict . encode diff --git a/src/Shikumi/CodeExec/CodeAct.hs b/src/Shikumi/CodeExec/CodeAct.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/CodeExec/CodeAct.hs @@ -0,0 +1,198 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | @codeAct@ (EP-27): a multi-turn loop where each /action/ is a code snippet that +-- may call provided tools, accumulating a 'Trajectory', then extracting the typed +-- answer. It mirrors DSPy's @CodeAct@ (a ReAct loop whose action is code) and reuses +-- 'Shikumi.Agent.ReAct'\'s @Trajectory@/@Step@/@Termination@ data model. +-- +-- The hermetic restricted interpreter does not itself call host functions, so tool +-- use is handled at the /protocol/ level: a snippet of the exact form +-- @call("\", \)@ is recognized by the loop and dispatched +-- through the typed 'Shikumi.Tool.runToolCall'; any other snippet is evaluated by the +-- sandbox and its output (or error) becomes the observation. A real subprocess +-- interpreter (EP-27 M4) would instead inject the tool functions into the +-- interpreter and call them natively, and may relax this convention. +-- +-- Like @react@/@programOfThought@, the whole loop is one 'Shikumi.Program.Embed' +-- node, so @codeAct@ carries no tunable parameters (@foldParams (codeAct sig reg) == +-- []@) and composes/serializes unchanged. +module Shikumi.CodeExec.CodeAct + ( CodeActConfig (..), + defaultCodeActConfig, + codeAct, + codeActWithTrajectory, + ) +where + +import Baikai (ToolCall, _Model, _ToolCall) +import Control.Lens ((&), (.~)) +import Data.Aeson (Value (..), eitherDecodeStrict) +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.Text.Encoding (encodeUtf8) +import Data.Vector qualified as V +import Effectful (Eff, (:>)) +import Effectful.Error.Static (Error, throwError) +import Shikumi.Adapter (ToPrompt (toPrompt), responseText) +import Shikumi.Agent.ReAct (Action (..), Step (..), Termination (..), Trajectory (..), renderTrajectory) +import Shikumi.CodeExec.Interpreter (CodeInterpreter (..), restrictedInterpreter) +import Shikumi.CodeExec.Prompt (encodeText, schemaInstruction, simpleContext, stripFences) +import Shikumi.Error (ShikumiError) +import Shikumi.LLM (LLM, complete) +import Shikumi.Program (Program (FMap), embed) +import Shikumi.Schema (FromModel, ToSchema, Validatable, parseOutput, toSchema) +import Shikumi.Signature (Signature, getInstruction) +import Shikumi.Tool (ToolRegistry, registryTools, renderToolError, runToolCall, someToolDescription, someToolName, someToolSchema) + +-- | How @codeAct@ runs: the hard iteration cap and the sandbox. +data CodeActConfig = CodeActConfig + { maxIters :: !Int, + interpreter :: !CodeInterpreter + } + +-- | Five iterations, the hermetic 'restrictedInterpreter'. +defaultCodeActConfig :: CodeActConfig +defaultCodeActConfig = CodeActConfig {maxIters = 5, interpreter = restrictedInterpreter} + +-- | Build a @codeAct@ agent returning the typed answer, dropping the trajectory. +codeAct :: + (ToPrompt i, ToSchema o, FromModel o, Validatable o) => + Signature i o -> + ToolRegistry -> + Program i o +codeAct sig reg = FMap fst (codeActWithTrajectory defaultCodeActConfig sig reg) + +-- | Build a @codeAct@ agent that also returns the recorded 'Trajectory'. +codeActWithTrajectory :: + (ToPrompt i, ToSchema o, FromModel o, Validatable o) => + CodeActConfig -> + Signature i o -> + ToolRegistry -> + Program i (o, Trajectory) +codeActWithTrajectory cfg sig reg = embed (codeActLoop cfg sig reg) + +-- | The loop body, in exactly the @Embed@ row. +codeActLoop :: + forall i o es. + (LLM :> es, Error ShikumiError :> es, ToPrompt i, ToSchema o, FromModel o, Validatable o) => + CodeActConfig -> + Signature i o -> + ToolRegistry -> + i -> + Eff es (o, Trajectory) +codeActLoop cfg sig reg i = do + traj <- loop 0 [] + o <- extract traj + pure (o, traj) + where + loop :: Int -> [Step] -> Eff es Trajectory + loop iter acc + | iter >= maxIters cfg = + pure (Trajectory (V.fromList (reverse acc)) (TerminatedMaxIters iter)) + | otherwise = do + let (ctx, opts) = simpleContext turnSys (turnUser acc) + resp <- complete _Model ctx opts + case parseCodeReply (responseText resp) of + Left perr -> loop (iter + 1) (correctiveStep perr : acc) + Right (code, finished) -> do + (act, obs) <- runAction code + let step = Step {thought = "", action = act, observation = Just obs} + if finished + then pure (Trajectory (V.fromList (reverse (step : acc))) TerminatedFinish) + else loop (iter + 1) (step : acc) + + -- A tool call snippet (call("name", args)) dispatches through runToolCall; any + -- other snippet runs in the sandbox. Both produce an observation text. + runAction :: Text -> Eff es (Action, Text) + runAction code = case parseCall code of + Just (nm, args) -> do + res <- runToolCall reg (mkToolCall nm args) + pure (CallTool nm args, either renderToolError id res) + Nothing -> do + r <- runCode (interpreter cfg) code + pure (CallTool "exec" (String code), either id id r) + + extract :: Trajectory -> Eff es o + extract traj = do + let prompt = "Task:\n" <> toPrompt i <> "\n\nTrajectory:\n" <> renderTrajectory traj + (ctx, opts) = simpleContext extractSys prompt + resp <- complete _Model ctx opts + either throwError pure (parseOutput (stripFences (responseText resp))) + + turnUser acc = + "Task:\n" + <> toPrompt i + <> "\n\n" + <> if null acc + then "No steps yet." + else "Steps so far:\n" <> renderTrajectory (Trajectory (V.fromList (reverse acc)) TerminatedFinish) + + turnSys = getInstruction sig <> "\n\n" <> toolMenu reg <> "\n\n" <> codeActGuide + extractSys = getInstruction sig <> "\n\n" <> schemaInstruction (toSchema (Proxy @o)) + +-- | Parse a model turn @{"code": "...", "finished": }@ (fences stripped). +parseCodeReply :: Text -> Either Text (Text, Bool) +parseCodeReply raw = case eitherDecodeStrict (encodeUtf8 (stripFences raw)) of + Left e -> Left ("not valid JSON: " <> T.pack e) + Right (Object o) -> do + code <- case KM.lookup "code" o of + Just (String c) -> Right c + _ -> Left "missing \"code\" string" + fin <- case KM.lookup "finished" o of + Just (Bool b) -> Right b + Nothing -> Right False + _ -> Left "\"finished\" must be a boolean" + pure (code, fin) + Right _ -> Left "expected a JSON object" + +-- | Recognize a tool-call snippet @call("name", )@, decoding the name and +-- arguments. Anything else returns 'Nothing' (it is handed to the sandbox). +parseCall :: Text -> Maybe (Text, Value) +parseCall code = do + afterOpen <- T.stripPrefix "call(" (T.strip code) + inner <- T.stripSuffix ")" (T.strip afterOpen) + case eitherDecodeStrict (encodeUtf8 ("[" <> inner <> "]")) of + Right (Array v) -> case V.toList v of + [String nm, args] -> Just (nm, args) + _ -> Nothing + _ -> Nothing + +-- | A baikai 'ToolCall' from a name and raw arguments object (the id is irrelevant). +mkToolCall :: Text -> Value -> ToolCall +mkToolCall nm args = _ToolCall & #name .~ nm & #arguments .~ args + +-- | A synthetic step recording an unparseable reply, fed back as an observation. +correctiveStep :: Text -> Step +correctiveStep perr = + Step + { thought = "", + action = CallTool "" Null, + observation = Just ("Your previous reply was not a valid {code, finished} object: " <> perr <> ".") + } + +-- | The tool menu rendered into the system prompt. +toolMenu :: ToolRegistry -> Text +toolMenu reg = case registryTools reg of + [] -> "You have no tools available." + ts -> "You can call these tools from your code as call(\"name\", {args}):\n" <> T.unlines (map one ts) + where + one st = + "- " + <> someToolName st + <> ": " + <> someToolDescription st + <> " (arguments schema: " + <> encodeText (someToolSchema st) + <> ")" + +-- | The action grammar for a code-act turn. +codeActGuide :: Text +codeActGuide = + "Each turn, reply with exactly one JSON object {\"code\": \"\", \"finished\": }. " + <> "The snippet runs in a restricted interpreter (arithmetic + - * /, strings with ++ and " + <> "len/upper/lower, lists with sum/length/concat). To call a tool, make the snippet exactly " + <> "call(\"\", ). Set finished to true once you have the answer." diff --git a/src/Shikumi/CodeExec/Interpreter.hs b/src/Shikumi/CodeExec/Interpreter.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/CodeExec/Interpreter.hs @@ -0,0 +1,314 @@ +{-# LANGUAGE RankNTypes #-} + +-- | The sandbox as a swappable /value/ (EP-27). +-- +-- A 'CodeInterpreter' takes a string of model-emitted code and returns either a +-- recoverable error message (fed back to the model) or the code's textual output. +-- It is captured in an 'Shikumi.Program.Embed' closure exactly as a +-- 'Shikumi.Tool.ToolRegistry' is — never added to the effect row — because the +-- @Embed@ body is fixed to the row @(LLM, Error ShikumiError)@. The hermetic +-- interpreters here are /pure/ (they ignore the effects and 'pure' their result), so +-- they fit inside any row. +-- +-- == SECURITY POSTURE +-- +-- Running model-emitted code is a genuine remote-code-execution risk. The hermetic +-- default 'restrictedInterpreter' is safe /by construction/: it parses a tiny DSL +-- and evaluates it purely — there is no syscall it can make, no filesystem, no +-- network, no unbounded loop primitive (a step/expression-size cap bounds runtime and +-- returns a typed error past it). The dangerous path is a real subprocess +-- interpreter (DSPy uses a @deno@/Pyodide sandbox); that is intentionally /not/ +-- shipped here. A future, gated, non-CI @subprocessInterpreter@ (EP-27 M4) would have +-- to enforce, at minimum: no network; no host filesystem access beyond a single fresh +-- scratch dir; no environment inheritance; CPU/memory/wall-clock limits (killing the +-- child and surfacing a 'Shikumi.Error.Timeout'); and loud failure (any limit breach +-- becomes a 'Shikumi.Error.ShikumiError', never a silent empty result). Because such +-- a subprocess needs @IOE@, which the @Embed@ row lacks, it could only be offered +-- through a separate @IOE@-bearing entry point, never inside the composable +-- @programOfThought@/@codeAct@ programs. +-- +-- == The restricted DSL +-- +-- A single expression (optionally prefixed @result = @), built from: integer and +-- rational literals and @+ - * /@ with parentheses; string literals with @++@ +-- concatenation and the functions @len@, @upper@, @lower@; and list literals @[a, b, +-- …]@ with @sum@, @length@, @concat@. The value of the expression is rendered to +-- text as the output. A parse error, an unknown identifier/function, a type error, +-- division by zero, or exceeding the step cap is returned as @Left \@. +module Shikumi.CodeExec.Interpreter + ( -- * The interpreter value + CodeInterpreter (..), + + -- * Hermetic implementations + restrictedInterpreter, + echoInterpreter, + + -- * The pure restricted evaluator (exposed for unit tests) + runRestricted, + ) +where + +import Control.Applicative ((<|>)) +import Data.Char (isAlpha, isAlphaNum, isDigit, isSpace) +import Data.Ratio (denominator, numerator) +import Data.Text (Text) +import Data.Text qualified as T +import Effectful (Eff, (:>)) +import Effectful.Error.Static (Error) +import Shikumi.Error (ShikumiError) +import Shikumi.LLM (LLM) +import Text.Read (readMaybe) + +-- | A sandbox as a plain value: run a code string, get back either a recoverable +-- error message (@Left@, fed to the model) or the program's textual output +-- (@Right@). A genuine infrastructure failure (e.g. a real subprocess dying) would +-- be thrown as a 'ShikumiError' from inside 'runCode'; the hermetic interpreters +-- never do that. +newtype CodeInterpreter = CodeInterpreter + { runCode :: forall es. (LLM :> es, Error ShikumiError :> es) => Text -> Eff es (Either Text Text) + } + +-- | The hermetic default: parse and evaluate the restricted DSL purely. No IO, no +-- network, no filesystem; a step cap bounds runtime. +restrictedInterpreter :: CodeInterpreter +restrictedInterpreter = CodeInterpreter (pure . runRestricted) + +-- | Echo the code back as its output (a trivial stub for wiring tests). +echoInterpreter :: CodeInterpreter +echoInterpreter = CodeInterpreter (pure . Right) + +-- --------------------------------------------------------------------------- +-- The restricted evaluator +-- --------------------------------------------------------------------------- + +-- | Evaluate one restricted-DSL expression, returning the rendered result (@Right@) +-- or a recoverable error message (@Left@). Pure and total. +runRestricted :: Text -> Either Text Text +runRestricted src = do + toks <- tokenize src + expr <- parseProgram toks + (v, _) <- eval stepCap expr + pure (renderVal v) + where + stepCap = 10000 + +-- --------------------------------------------------------------------------- +-- Values +-- --------------------------------------------------------------------------- + +data Val + = VNum !Rational + | VStr !Text + | VList ![Val] + deriving stock (Eq, Show) + +renderVal :: Val -> Text +renderVal (VNum r) + | denominator r == 1 = T.pack (show (numerator r)) + | otherwise = T.pack (show (fromRational r :: Double)) +renderVal (VStr s) = s +renderVal (VList xs) = "[" <> T.intercalate ", " (map renderVal xs) <> "]" + +-- --------------------------------------------------------------------------- +-- Tokenizer +-- --------------------------------------------------------------------------- + +data Tok + = TNum !Rational + | TStr !Text + | TIdent !Text + | TKwResult + | TPlus + | TMinus + | TStar + | TSlash + | TConcat + | TLParen + | TRParen + | TLBrack + | TRBrack + | TComma + | TEq + deriving stock (Eq, Show) + +tokenize :: Text -> Either Text [Tok] +tokenize = go . T.unpack + where + go [] = Right [] + go (c : cs) + | isSpace c = go cs + | isDigit c = + let (numStr, rest) = span (\x -> isDigit x || x == '.') (c : cs) + in case readNum numStr of + Just r -> (TNum r :) <$> go rest + Nothing -> Left ("bad number literal: " <> T.pack numStr) + | c == '"' = + let (str, rest) = break (== '"') cs + in case rest of + ('"' : rest') -> (TStr (T.pack str) :) <$> go rest' + _ -> Left "unterminated string literal" + | isAlpha c || c == '_' = + let (ident, rest) = span (\x -> isAlphaNum x || x == '_') (c : cs) + in (identTok ident :) <$> go rest + | c == '+' = case cs of + ('+' : rest) -> (TConcat :) <$> go rest + _ -> (TPlus :) <$> go cs + | c == '-' = (TMinus :) <$> go cs + | c == '*' = (TStar :) <$> go cs + | c == '/' = (TSlash :) <$> go cs + | c == '(' = (TLParen :) <$> go cs + | c == ')' = (TRParen :) <$> go cs + | c == '[' = (TLBrack :) <$> go cs + | c == ']' = (TRBrack :) <$> go cs + | c == ',' = (TComma :) <$> go cs + | c == '=' = (TEq :) <$> go cs + | otherwise = Left ("unexpected character: " <> T.singleton c) + identTok "result" = TKwResult + identTok s = TIdent (T.pack s) + readNum :: String -> Maybe Rational + readNum s = (fromInteger <$> (readMaybe s :: Maybe Integer)) <|> (toRational <$> (readMaybe s :: Maybe Double)) + +-- --------------------------------------------------------------------------- +-- Parser (recursive descent) +-- --------------------------------------------------------------------------- + +data Op = Add | Sub | Mul | Div | Concat + deriving stock (Eq, Show) + +data Expr + = ENum !Rational + | EStr !Text + | EList ![Expr] + | EBin !Op !Expr !Expr + | ECall !Text ![Expr] + deriving stock (Eq, Show) + +type P a = [Tok] -> Either Text (a, [Tok]) + +parseProgram :: [Tok] -> Either Text Expr +parseProgram toks0 = do + let toks = case toks0 of + (TKwResult : TEq : rest) -> rest + _ -> toks0 + (e, rest) <- parseExpr toks + case rest of + [] -> Right e + _ -> Left "unexpected trailing tokens" + +parseExpr :: P Expr +parseExpr = parseAdd + +parseAdd :: P Expr +parseAdd toks = do + (l, r1) <- parseMul toks + goAdd l r1 + where + goAdd l (TPlus : ts) = parseMul ts >>= \(r, ts') -> goAdd (EBin Add l r) ts' + goAdd l (TMinus : ts) = parseMul ts >>= \(r, ts') -> goAdd (EBin Sub l r) ts' + goAdd l (TConcat : ts) = parseMul ts >>= \(r, ts') -> goAdd (EBin Concat l r) ts' + goAdd l ts = Right (l, ts) + +parseMul :: P Expr +parseMul toks = do + (l, r1) <- parseFactor toks + goMul l r1 + where + goMul l (TStar : ts) = parseFactor ts >>= \(r, ts') -> goMul (EBin Mul l r) ts' + goMul l (TSlash : ts) = parseFactor ts >>= \(r, ts') -> goMul (EBin Div l r) ts' + goMul l ts = Right (l, ts) + +parseFactor :: P Expr +parseFactor (TNum n : ts) = Right (ENum n, ts) +parseFactor (TStr s : ts) = Right (EStr s, ts) +parseFactor (TLParen : ts) = do + (e, ts') <- parseExpr ts + case ts' of + (TRParen : ts'') -> Right (e, ts'') + _ -> Left "expected closing )" +parseFactor (TLBrack : ts) = parseSeq TRBrack ts >>= \(es, ts') -> Right (EList es, ts') +parseFactor (TIdent name : TLParen : ts) = parseSeq TRParen ts >>= \(args, ts') -> Right (ECall name args, ts') +parseFactor (TIdent name : ts) = Right (ECall name [], ts) +parseFactor _ = Left "expected an expression" + +-- | Parse a comma-separated sequence of expressions terminated by the given token. +parseSeq :: Tok -> P [Expr] +parseSeq close (t : ts) | t == close = Right ([], ts) +parseSeq close ts = do + (e, ts') <- parseExpr ts + case ts' of + (TComma : ts'') -> parseSeq close ts'' >>= \(es, ts3) -> Right (e : es, ts3) + (t : ts'') | t == close -> Right ([e], ts'') + _ -> Left "expected , or closing bracket" + +-- --------------------------------------------------------------------------- +-- Evaluator (with a step budget) +-- --------------------------------------------------------------------------- + +eval :: Int -> Expr -> Either Text (Val, Int) +eval fuel _ | fuel <= 0 = Left "step cap exceeded" +eval fuel (ENum n) = Right (VNum n, fuel - 1) +eval fuel (EStr s) = Right (VStr s, fuel - 1) +eval fuel (EList es) = do + (vs, f') <- evalList (fuel - 1) es + Right (VList vs, f') +eval fuel (EBin op a b) = do + (va, f1) <- eval (fuel - 1) a + (vb, f2) <- eval f1 b + v <- applyOp op va vb + Right (v, f2) +eval fuel (ECall name args) = do + (vs, f1) <- evalList (fuel - 1) args + v <- applyFn name vs + Right (v, f1) + +evalList :: Int -> [Expr] -> Either Text ([Val], Int) +evalList fuel [] = Right ([], fuel) +evalList fuel (e : es) = do + (v, f1) <- eval fuel e + (vs, f2) <- evalList f1 es + Right (v : vs, f2) + +applyOp :: Op -> Val -> Val -> Either Text Val +applyOp Add (VNum a) (VNum b) = Right (VNum (a + b)) +applyOp Sub (VNum a) (VNum b) = Right (VNum (a - b)) +applyOp Mul (VNum a) (VNum b) = Right (VNum (a * b)) +applyOp Div (VNum _) (VNum 0) = Left "division by zero" +applyOp Div (VNum a) (VNum b) = Right (VNum (a / b)) +applyOp Concat (VStr a) (VStr b) = Right (VStr (a <> b)) +applyOp Concat (VList a) (VList b) = Right (VList (a ++ b)) +applyOp op _ _ = Left ("type error: operator " <> opName op <> " applied to incompatible values") + +opName :: Op -> Text +opName = \case + Add -> "+" + Sub -> "-" + Mul -> "*" + Div -> "/" + Concat -> "++" + +applyFn :: Text -> [Val] -> Either Text Val +applyFn "len" [VStr s] = Right (VNum (fromIntegral (T.length s))) +applyFn "len" [VList xs] = Right (VNum (fromIntegral (length xs))) +applyFn "length" [VList xs] = Right (VNum (fromIntegral (length xs))) +applyFn "length" [VStr s] = Right (VNum (fromIntegral (T.length s))) +applyFn "upper" [VStr s] = Right (VStr (T.toUpper s)) +applyFn "lower" [VStr s] = Right (VStr (T.toLower s)) +applyFn "sum" [VList xs] = VNum . sum <$> traverse asNum xs +applyFn "concat" [VList xs] = concatVals xs +applyFn name [] = Left ("unknown identifier: " <> name) +applyFn name _ = Left ("unknown function or bad arguments: " <> name) + +asNum :: Val -> Either Text Rational +asNum (VNum n) = Right n +asNum _ = Left "type error: expected a number" + +concatVals :: [Val] -> Either Text Val +concatVals xs + | all isStr xs = Right (VStr (T.concat [s | VStr s <- xs])) + | all isList xs = Right (VList (concat [ys | VList ys <- xs])) + | otherwise = Left "type error: concat expects a list of all-strings or all-lists" + where + isStr (VStr _) = True + isStr _ = False + isList (VList _) = True + isList _ = False diff --git a/src/Shikumi/CodeExec/ProgramOfThought.hs b/src/Shikumi/CodeExec/ProgramOfThought.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/CodeExec/ProgramOfThought.hs @@ -0,0 +1,129 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | @programOfThought@ (EP-27): a program where the model writes code that +-- /computes/ the answer, the code runs in a sandbox, and the result flows back into +-- a typed answer. +-- +-- The loop mirrors DSPy's @ProgramOfThought.forward@: ask the model for a code +-- snippet that computes the answer; run it through the captured 'CodeInterpreter'; +-- if it errors, feed the error back and ask for a correction (up to @maxIters@ +-- attempts); once it runs, ask the model to extract the typed output @o@ from the +-- code and its output. The whole loop is a single 'Shikumi.Program.Embed' node, so +-- @programOfThought@ is an ordinary @'Program' i o@ that carries no tunable +-- parameters (the parameter-count invariant holds: @foldParams (programOfThought +-- sig) == []@) and composes/serializes exactly like @react@. +module Shikumi.CodeExec.ProgramOfThought + ( PoTConfig (..), + defaultPoTConfig, + programOfThought, + programOfThoughtWith, + ) +where + +import Baikai (_Model) +import Data.Proxy (Proxy (..)) +import Data.Text (Text) +import Data.Text qualified as T +import Effectful (Eff, (:>)) +import Effectful.Error.Static (Error, throwError) +import Shikumi.Adapter (ToPrompt (toPrompt), responseText) +import Shikumi.CodeExec.Interpreter (CodeInterpreter (..), restrictedInterpreter) +import Shikumi.CodeExec.Prompt (schemaInstruction, simpleContext, stripFences) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.LLM (LLM, complete) +import Shikumi.Program (Program, embed) +import Shikumi.Schema (FromModel, ToSchema, Validatable, parseOutput, toSchema) +import Shikumi.Signature (Signature, getInstruction) + +-- | How @programOfThought@ runs: the max number of code (re)generation attempts and +-- the sandbox to run snippets in. +data PoTConfig = PoTConfig + { maxIters :: !Int, + interpreter :: !CodeInterpreter + } + +-- | Three attempts, the hermetic 'restrictedInterpreter'. +defaultPoTConfig :: PoTConfig +defaultPoTConfig = PoTConfig {maxIters = 3, interpreter = restrictedInterpreter} + +-- | The ergonomic default: emit code, run it, fix on error up to @maxIters@, extract +-- the typed @o@. Uses 'defaultPoTConfig'. +programOfThought :: + (ToPrompt i, ToSchema o, FromModel o, Validatable o) => + Signature i o -> + Program i o +programOfThought = programOfThoughtWith defaultPoTConfig + +-- | @programOfThought@ with an explicit config (e.g. a different @maxIters@ or a +-- different 'CodeInterpreter'). +programOfThoughtWith :: + (ToPrompt i, ToSchema o, FromModel o, Validatable o) => + PoTConfig -> + Signature i o -> + Program i o +programOfThoughtWith cfg sig = embed (potLoop cfg sig) + +-- | The loop body, in exactly the @Embed@ row. +potLoop :: + forall i o es. + (LLM :> es, Error ShikumiError :> es, ToPrompt i, ToSchema o, FromModel o, Validatable o) => + PoTConfig -> + Signature i o -> + i -> + Eff es o +potLoop cfg sig i = go 0 Nothing + where + go :: Int -> Maybe (Text, Text) -> Eff es o + go iter mPrev = do + code <- propose mPrev + r <- runCode (interpreter cfg) code + case r of + Right out -> extract code out + Left err + | iter + 1 >= maxIters cfg -> + throwError + ( ProviderFailure + ("programOfThought: code failed after " <> tshow (iter + 1) <> " attempts: " <> err) + ) + | otherwise -> go (iter + 1) (Just (code, err)) + + -- Generate (first attempt) or regenerate (with the previous code + error). + propose :: Maybe (Text, Text) -> Eff es Text + propose mPrev = do + let (ctx, opts) = simpleContext (proposeSys mPrev) (toPrompt i) + resp <- complete _Model ctx opts + pure (stripFences (responseText resp)) + + -- Extract: hand the model the final code and its output, ask for the typed o. + extract :: Text -> Text -> Eff es o + extract code out = do + let prompt = toPrompt i <> "\n\nThe code:\n" <> code <> "\n\nIts output:\n" <> out + (ctx, opts) = simpleContext extractSys prompt + resp <- complete _Model ctx opts + either throwError pure (parseOutput (stripFences (responseText resp))) + + proposeSys Nothing = + getInstruction sig <> "\n\n" <> dslGuide <> "\nReply with only the code, no prose." + proposeSys (Just (prevCode, err)) = + getInstruction sig + <> "\n\n" + <> dslGuide + <> "\nYour previous code:\n" + <> prevCode + <> "\nfailed with error:\n" + <> err + <> "\nReply with only the corrected code." + + extractSys = getInstruction sig <> "\n\n" <> schemaInstruction (toSchema (Proxy @o)) + +-- | A description of the restricted language for the code-generation prompt. +dslGuide :: Text +dslGuide = + "Write a snippet in a restricted language that computes the answer. It supports " + <> "integer/rational arithmetic (+ - * /) with parentheses, string literals with ++ and " + <> "len/upper/lower, and lists with sum/length/concat. The value of the final expression " + <> "(optionally written `result = `) is the answer." + +tshow :: (Show a) => a -> Text +tshow = T.pack . show diff --git a/src/Shikumi/CodeExec/Prompt.hs b/src/Shikumi/CodeExec/Prompt.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/CodeExec/Prompt.hs @@ -0,0 +1,53 @@ +-- | Small prompt/parse helpers shared by the code-execution modules +-- ('Shikumi.CodeExec.ProgramOfThought', 'Shikumi.CodeExec.CodeAct'). These mirror +-- the equivalents in 'Shikumi.Agent.ReAct' (which keeps its own copies private); +-- they are factored here to avoid duplicating the few small helpers across the two +-- code-execution modules (EP-27). +module Shikumi.CodeExec.Prompt + ( simpleContext, + stripFences, + schemaInstruction, + encodeText, + ) +where + +import Baikai (Context, Options, user, _Context, _Options) +import Control.Lens ((&), (.~)) +import Data.Aeson (Value, encode) +import Data.ByteString.Lazy qualified as LBS +import Data.Generics.Labels () +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Encoding (decodeUtf8) +import Data.Vector qualified as V + +-- | A request context with a system prompt and a single user message, default +-- options. The neutral one-turn shape both code-execution loops use per model call. +simpleContext :: Text -> Text -> (Context, Options) +simpleContext sys userMsg = + ( _Context & #systemPrompt .~ Just sys & #messages .~ V.singleton (user userMsg), + _Options + ) + +-- | Strip a leading/trailing Markdown code fence (@```@ / @```json … ```@), if any, +-- so a fenced reply still decodes. Falls back to the trimmed input. (Same logic as +-- 'Shikumi.Agent.ReAct'\'s private @stripFences@.) +stripFences :: Text -> Text +stripFences t = + let trimmed = T.strip t + in case T.stripPrefix "```" trimmed of + Nothing -> trimmed + Just rest -> T.strip (dropClosingFence (dropFirstLine rest)) + where + dropFirstLine s = T.drop 1 (T.dropWhile (/= '\n') s) + dropClosingFence s = fst (T.breakOn "```" s) + +-- | An extraction instruction: ask for one JSON object matching the given schema. +schemaInstruction :: Value -> Text +schemaInstruction schema = + "Produce the final answer as exactly one JSON object matching this schema, and nothing else:\n" + <> encodeText schema + +-- | Compact-encode a JSON value to text. +encodeText :: Value -> Text +encodeText = decodeUtf8 . LBS.toStrict . encode diff --git a/src/Shikumi/Tool.hs b/src/Shikumi/Tool.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Tool.hs @@ -0,0 +1,227 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Typed tools and their lowering to baikai's wire tool — the MasterPlan's +-- integration point #8, owned here and nowhere else. +-- +-- A @'Tool' i o@ is a named function from an input record @i@ to an output @o@, +-- runnable in any effect stack that has the @LLM@ and @Error ShikumiError@ +-- capabilities (so a tool body may itself call sub-models and may signal failure +-- through the typed error channel — see the plan's Decision Log on why the row is +-- @(LLM, Error ShikumiError)@ rather than @IOE@). From a tool shikumi derives the +-- JSON Schema of @i@ (reusing 'Shikumi.Schema'), lowers it to baikai's untyped +-- @Baikai.Tool { name, description, parameters }@, decodes the model's argument +-- JSON into a typed @i@ (or a 'ToolError', never a crash), runs the body, and +-- encodes the @o@ result back to observation text. +-- +-- Heterogeneity is handled by the 'SomeTool' existential: it hides @i@/@o@ while +-- retaining, at wrap time, the schema/decode/encode dictionaries the wire boundary +-- needs. A 'ToolRegistry' is a name-keyed map of 'SomeTool's that a ReAct loop +-- dispatches against by the tool name the model emits at run time. +module Shikumi.Tool + ( -- * Typed tools + Tool (..), + mkTool, + toolSchemaOf, + lowerTool, + + -- * Heterogeneous tools and the registry + SomeTool (..), + someToolName, + someToolDescription, + someToolSchema, + lowerSomeTool, + runErased, + ToolRegistry, + mkRegistry, + registryLookup, + registryBaikai, + registryNames, + registryTools, + + -- * The typed error and the wire round-trip + ToolError (..), + renderToolError, + runToolCall, + ) +where + +import Baikai (ToolCall, _Tool) +import Baikai qualified as B +import Control.Lens ((&), (.~), (^.)) +import Data.Aeson (ToJSON, Value, encode) +import Data.ByteString.Lazy qualified as LBS +import Data.Generics.Labels () +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Proxy (Proxy (..)) +import Data.Text (Text) +import Data.Text.Encoding (decodeUtf8) +import Data.Vector (Vector) +import Data.Vector qualified as V +import Effectful (Eff, (:>)) +import Effectful.Error.Static (Error, catchError) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.LLM (LLM) +import Shikumi.Schema (FromModel, ToSchema, Validatable, fromModelChecked, toSchema) + +-- --------------------------------------------------------------------------- +-- The typed tool +-- --------------------------------------------------------------------------- + +-- | A typed tool: a named function from an input record @i@ to an output @o@. The +-- body is rank-2 over the effect row so the same value can run in any stack that +-- has @LLM@ and @Error ShikumiError@ — the exact row 'Shikumi.Program.runProgram' +-- provides, which is what lets a ReAct agent (built on the @Embed@ program node) +-- execute its tools without widening that row. +data Tool i o = Tool + { name :: !Text, + description :: !Text, + run :: !(forall es. (LLM :> es, Error ShikumiError :> es) => i -> Eff es o) + } + +-- | Build a tool from a name, a description, and an effectful body. A pure tool is +-- just @\\i -> pure (f i)@. The schema/decode/encode dictionaries are not needed +-- here — they are captured later, at 'SomeTool' wrap time. +mkTool :: + Text -> + Text -> + (forall es. (LLM :> es, Error ShikumiError :> es) => i -> Eff es o) -> + Tool i o +mkTool = Tool + +-- | The JSON Schema of a tool's input record, derived from @i@. +toolSchemaOf :: forall i o. (ToSchema i) => Tool i o -> Value +toolSchemaOf _ = toSchema (Proxy @i) + +-- | Lower a typed tool to baikai's untyped wire tool, putting the derived input +-- schema into @parameters@. This is the single sanctioned typed -> wire lowering. +lowerTool :: forall i o. (ToSchema i) => Tool i o -> B.Tool +lowerTool t = + _Tool + & #name .~ name t + & #description .~ description t + & #parameters .~ toolSchemaOf t + +-- --------------------------------------------------------------------------- +-- The existential and the registry +-- --------------------------------------------------------------------------- + +-- | A tool with its @i@/@o@ hidden, retaining the dictionaries needed to (a) +-- derive its input schema, (b) decode arguments from JSON into @i@, and (c) encode +-- its @o@ result to text — all captured at wrap time. +data SomeTool where + SomeTool :: + (ToSchema i, FromModel i, Validatable i, ToJSON o) => + Tool i o -> + SomeTool + +-- | The name of an erased tool. +someToolName :: SomeTool -> Text +someToolName (SomeTool t) = name t + +-- | The description of an erased tool. +someToolDescription :: SomeTool -> Text +someToolDescription (SomeTool t) = description t + +-- | The derived input JSON Schema of an erased tool. +someToolSchema :: SomeTool -> Value +someToolSchema (SomeTool t) = toolSchemaOf t + +-- | Lower an erased tool to baikai's wire tool. +lowerSomeTool :: SomeTool -> B.Tool +lowerSomeTool (SomeTool t) = lowerTool t + +-- | Run an erased tool against a raw JSON arguments object: decode to the hidden +-- @i@, run the body, encode the @o@ to text. Total: a decode failure becomes +-- 'ToolArgsInvalid' and a body that throws a 'ShikumiError' becomes 'ToolRunFailed' +-- — it never throws to the caller. +runErased :: + (LLM :> es, Error ShikumiError :> es) => + SomeTool -> + Value -> + Eff es (Either ToolError Text) +runErased (SomeTool t) args = + case fromModelChecked args of + Left err -> pure (Left (ToolArgsInvalid (name t) (shikumiErrorText err))) + Right i -> + (Right . encodeText <$> run t i) + `catchError` \_cs e -> pure (Left (ToolRunFailed (name t) (shikumiErrorText e))) + +-- | A heterogeneous tool set, keyed by tool name for O(log n) dispatch. +newtype ToolRegistry = ToolRegistry (Map Text SomeTool) + +-- | Build a registry from a list of erased tools (last wins on a name clash). +mkRegistry :: [SomeTool] -> ToolRegistry +mkRegistry = ToolRegistry . Map.fromList . map (\st -> (someToolName st, st)) + +-- | Look a tool up by the name the model emitted. +registryLookup :: Text -> ToolRegistry -> Maybe SomeTool +registryLookup nm (ToolRegistry m) = Map.lookup nm m + +-- | Every tool lowered to baikai's wire form, for @Context.tools@. +registryBaikai :: ToolRegistry -> Vector B.Tool +registryBaikai (ToolRegistry m) = V.fromList (map lowerSomeTool (Map.elems m)) + +-- | The names of the registered tools (the tool menu). +registryNames :: ToolRegistry -> [Text] +registryNames (ToolRegistry m) = Map.keys m + +-- | The erased tools, in name order. +registryTools :: ToolRegistry -> [SomeTool] +registryTools (ToolRegistry m) = Map.elems m + +-- --------------------------------------------------------------------------- +-- The typed error and the wire round-trip +-- --------------------------------------------------------------------------- + +-- | A tool-call failure carried as a /value/, never an exception. The agent feeds +-- the rendered text back to the model as an observation so it can recover, and +-- records it in the trajectory; only infrastructure faults bubble up as a +-- 'ShikumiError'. +data ToolError + = -- | the model named a tool the registry does not have + ToolNotFound !Text + | -- | tool name, decode-error message: the arguments did not match @i@ + ToolArgsInvalid !Text !Text + | -- | tool name, failure message: the body signalled a 'ShikumiError' + ToolRunFailed !Text !Text + deriving stock (Show, Eq) + +-- | Render a 'ToolError' as the observation text fed back to the model. +renderToolError :: ToolError -> Text +renderToolError = \case + ToolNotFound nm -> "Error: no such tool \"" <> nm <> "\"." + ToolArgsInvalid nm msg -> "Error: invalid arguments for \"" <> nm <> "\": " <> msg + ToolRunFailed nm msg -> "Error: tool \"" <> nm <> "\" failed: " <> msg + +-- | Find the named tool, decode, run, encode. The entry point the ReAct loop uses. +runToolCall :: + (LLM :> es, Error ShikumiError :> es) => + ToolRegistry -> + ToolCall -> + Eff es (Either ToolError Text) +runToolCall reg tc = + case registryLookup (tc ^. #name) reg of + Nothing -> pure (Left (ToolNotFound (tc ^. #name))) + Just st -> runErased st (tc ^. #arguments) + +-- --------------------------------------------------------------------------- +-- Helpers +-- --------------------------------------------------------------------------- + +-- | Compact-encode a JSON-serializable result to text for an observation. +encodeText :: (ToJSON a) => a -> Text +encodeText = decodeUtf8 . LBS.toStrict . encode + +-- | The human-readable payload of a 'ShikumiError' (its single 'Text' field). +shikumiErrorText :: ShikumiError -> Text +shikumiErrorText = \case + InvalidJSON t -> t + MissingField t -> "missing field " <> t + SchemaMismatch t -> t + ValidationFailure t -> t + ProviderFailure t -> t + Timeout t -> t + BudgetExceeded t -> t diff --git a/test/AcceptanceSpec.hs b/test/AcceptanceSpec.hs new file mode 100644 --- /dev/null +++ b/test/AcceptanceSpec.hs @@ -0,0 +1,81 @@ +-- | M4: the headline acceptance from the plan's Purpose, demonstrated entirely +-- against a mock LM (no live provider). It shows the four observable facts: a typed +-- 'WeatherResp' answer from a ReAct run; the recorded 'Trajectory' with the expected +-- steps and 'TerminatedFinish'; the frozen 'WeatherReq' schema; and malformed tool +-- arguments captured as a recorded 'ToolError' observation while the agent still +-- returns a typed value and throws nothing. +module AcceptanceSpec (tests) where + +import Baikai (Response) +import Data.Maybe (fromMaybe) +import Data.Text qualified as T +import Data.Vector qualified as V +import Fixtures + ( WeatherResp, + badArgsPromptScript, + expectedReqSchema, + expectedWeather, + nativeScript, + promptScript, + weatherQuestion, + weatherRegistry, + weatherSignature, + weatherTool, + ) +import MockLLM (runAgent) +import Shikumi.Agent.ReAct + ( Action (..), + ReActConfig (..), + Step (..), + Termination (..), + ToolProtocol (..), + Trajectory (..), + defaultReActConfig, + reactWithTrajectory, + ) +import Shikumi.Tool (toolSchemaOf) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Acceptance" + [ testCase "derived tool schema is the frozen WeatherReq schema" $ + toolSchemaOf weatherTool @?= expectedReqSchema, + testCase "typed tool + ReAct + mock LM end-to-end (prompt protocol)" $ + endToEnd ProtocolPrompt promptScript, + testCase "typed tool + ReAct + mock LM end-to-end (native protocol)" $ + endToEnd ProtocolNative nativeScript, + testCase "malformed args become a recorded ToolError, no crash" $ do + res <- + runAgent + badArgsPromptScript + (reactWithTrajectory weatherSignature weatherRegistry defaultReActConfig {protocol = ProtocolPrompt}) + weatherQuestion + case res of + Right (o :: WeatherResp, traj) -> do + o @?= expectedWeather + let obs = fromMaybe "" (observation (V.head (steps traj))) + assertBool "the bad tool call is recorded as an invalid-arguments observation" ("invalid arguments" `T.isInfixOf` obs) + Left e -> assertFailure ("agent should recover and return, but failed: " <> show e) + ] + +-- | Run the agent under a forced protocol against a script and assert the typed +-- answer, a two-step (tool call then finish) trajectory, and a clean finish. +endToEnd :: ToolProtocol -> [Response] -> Assertion +endToEnd proto script = do + res <- + runAgent + script + (reactWithTrajectory weatherSignature weatherRegistry defaultReActConfig {protocol = proto}) + weatherQuestion + case res of + Right (o :: WeatherResp, traj) -> do + o @?= expectedWeather + termination traj @?= TerminatedFinish + V.length (steps traj) @?= 2 + case action (V.head (steps traj)) of + CallTool nm _ -> nm @?= "get_weather" + other -> assertFailure ("first step should be a tool call, got " <> show other) + Left e -> assertFailure ("agent failed: " <> show e) diff --git a/test/CodeActSpec.hs b/test/CodeActSpec.hs new file mode 100644 --- /dev/null +++ b/test/CodeActSpec.hs @@ -0,0 +1,74 @@ +{-# LANGUAGE DeriveAnyClass #-} + +-- | EP-27 M3: @codeAct@ — a multi-turn loop whose action is a code snippet that may +-- call provided tools. Verified hermetically: a snippet calls a registered tool +-- (@addOne@) via the @call("name", args)@ convention, a second snippet computes a +-- value in the sandbox, the loop records a two-step 'Trajectory', finishes, and the +-- typed answer is extracted. +module CodeActSpec (tests) where + +import Data.Aeson (object, (.=)) +import Data.Text (Text) +import Data.Vector qualified as V +import GHC.Generics (Generic) +import MockLLM (mkTextResponse, runAgent) +import Shikumi.Adapter (ToPrompt) +import Shikumi.Agent.ReAct (Step (..), Termination (..), Trajectory (..)) +import Shikumi.CodeExec.CodeAct (codeActWithTrajectory, defaultCodeActConfig) +import Shikumi.CodeExec.Prompt (encodeText) +import Shikumi.Schema (FromModel, ToSchema) +import Shikumi.Signature (Signature, mkSignature) +import Shikumi.Tool (SomeTool (..), Tool, ToolRegistry, mkRegistry, mkTool) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertFailure, testCase, (@?=)) + +newtype Task = Task {task :: Text} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToPrompt) + +newtype CalcAnswer = CalcAnswer {value :: Int} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +newtype AddIn = AddIn {n :: Int} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel) + +-- | A tool that adds one to its argument. +addOneTool :: Tool AddIn Int +addOneTool = mkTool "addOne" "Add one to the integer n." (\(AddIn k) -> pure (k + 1)) + +registry :: ToolRegistry +registry = mkRegistry [SomeTool addOneTool] + +sig :: Signature Task CalcAnswer +sig = mkSignature "Answer the task, using tools from code when helpful." + +-- | A model turn JSON object @{code, finished}@. +turn :: Text -> Bool -> Text +turn code finished = encodeText (object ["code" .= code, "finished" .= finished]) + +tests :: TestTree +tests = + testGroup + "CodeActSpec" + [ testCase "a code snippet calls a tool, accumulates a 2-step trajectory, extracts the answer" $ do + let script = + [ mkTextResponse (turn "call(\"addOne\", {\"n\": 41})" False), + mkTextResponse (turn "result = 42" True), + mkTextResponse "{\"value\": 42}" + ] + out <- + runAgent + script + (codeActWithTrajectory defaultCodeActConfig sig registry) + (Task "add one to 41") + case out of + Right (answer, traj) -> do + answer @?= CalcAnswer 42 + V.length (steps traj) @?= 2 + termination traj @?= TerminatedFinish + -- the first step's observation is the tool result (addOne 41 = 42) + observation (V.head (steps traj)) @?= Just "42" + Left e -> assertFailure ("expected a typed answer, got " <> show e) + ] diff --git a/test/Fixtures.hs b/test/Fixtures.hs new file mode 100644 --- /dev/null +++ b/test/Fixtures.hs @@ -0,0 +1,166 @@ +{-# LANGUAGE DeriveAnyClass #-} + +-- | Shared fixtures for the @shikumi-tools@ specs: the weather records from the +-- plan's Purpose, their schema/decode/prompt instances, the typed @weatherTool@ and +-- its registry, the mock-LM scripts for both protocols, and the expected schema and +-- answer the specs assert against. +module Fixtures + ( -- * Records + WeatherReq (..), + WeatherResp (..), + AnswerWeatherQuestion (..), + + -- * Tool, signature, registry + weatherTool, + weatherRegistry, + weatherSignature, + + -- * Inputs / expected outputs + weatherQuestion, + expectedWeather, + weatherArgs, + expectedReqSchema, + + -- * Mock-LM scripts + promptScript, + nativeScript, + badArgsPromptScript, + maxItersScript, + ) +where + +import Baikai (Response) +import Data.Aeson (ToJSON, Value, object, (.=)) +import Data.Text (Text) +import GHC.Generics (Generic) +import MockLLM (mkTextResponse, mkToolCallResponse) +import Shikumi.Adapter (ToPrompt) +import Shikumi.Schema (FromModel, ToSchema) +import Shikumi.Signature (Signature, mkSignature) +import Shikumi.Tool (SomeTool (..), Tool, ToolRegistry, mkRegistry, mkTool) + +-- --------------------------------------------------------------------------- +-- Records +-- --------------------------------------------------------------------------- + +data WeatherReq = WeatherReq {city :: !Text, units :: !Text} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +data WeatherResp = WeatherResp {tempC :: !Double, summary :: !Text} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel, ToJSON, ToPrompt) + +newtype AnswerWeatherQuestion = AnswerWeatherQuestion {question :: Text} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToPrompt) + +-- --------------------------------------------------------------------------- +-- Tool, signature, registry +-- --------------------------------------------------------------------------- + +-- | A pure typed tool: it ignores its request and returns a fixed forecast (the +-- specs only need determinism, not a real lookup). +weatherTool :: Tool WeatherReq WeatherResp +weatherTool = + mkTool "get_weather" "Look up the current weather for a city." $ \_req -> + pure fixedForecast + +fixedForecast :: WeatherResp +fixedForecast = WeatherResp {tempC = 12.0, summary = "mild"} + +weatherRegistry :: ToolRegistry +weatherRegistry = mkRegistry [SomeTool weatherTool] + +weatherSignature :: Signature AnswerWeatherQuestion WeatherResp +weatherSignature = mkSignature "Answer the user's weather question, using tools when helpful." + +-- --------------------------------------------------------------------------- +-- Inputs / expected outputs +-- --------------------------------------------------------------------------- + +weatherQuestion :: AnswerWeatherQuestion +weatherQuestion = AnswerWeatherQuestion {question = "What's the weather in Paris?"} + +expectedWeather :: WeatherResp +expectedWeather = fixedForecast + +-- | The arguments the scripted model proposes for @get_weather@. +weatherArgs :: Value +weatherArgs = object ["city" .= ("Paris" :: Text), "units" .= ("c" :: Text)] + +-- | The exact JSON Schema 'Shikumi.Tool.toolSchemaOf' must produce for 'WeatherReq' +-- (frozen from "Shikumi.Schema"'s generator: both fields required, no additional +-- properties). +expectedReqSchema :: Value +expectedReqSchema = + object + [ "type" .= ("object" :: Text), + "properties" + .= object + [ "city" .= object ["type" .= ("string" :: Text)], + "units" .= object ["type" .= ("string" :: Text)] + ], + "required" .= (["city", "units"] :: [Text]), + "additionalProperties" .= False + ] + +-- --------------------------------------------------------------------------- +-- Mock-LM scripts +-- --------------------------------------------------------------------------- + +-- | The extract turn's reply: a JSON object decoding to 'expectedWeather'. +extractReply :: Text +extractReply = "{\"tempC\": 12.0, \"summary\": \"mild\"}" + +-- | A prompt-protocol propose reply: one action JSON object matching the grammar +-- 'Shikumi.Agent.ReAct' renders. @CallTool@ form. +proposeCallReply :: Text +proposeCallReply = + "{\"thought\": \"I should look up Paris.\", \"action\": {\"tool\": \"get_weather\", \"args\": {\"city\": \"Paris\", \"units\": \"c\"}}}" + +-- | A prompt-protocol propose reply with arguments missing the required @units@ +-- field, to exercise the typed 'Shikumi.Tool.ToolArgsInvalid' recovery path. +proposeBadArgsReply :: Text +proposeBadArgsReply = + "{\"thought\": \"Looking up Paris.\", \"action\": {\"tool\": \"get_weather\", \"args\": {\"city\": \"Paris\"}}}" + +-- | A prompt-protocol finish reply. +finishReply :: Text +finishReply = "{\"thought\": \"I have the forecast.\", \"action\": {\"finish\": true}}" + +-- | Prompt protocol: propose the tool call, finish, then extract. +promptScript :: [Response] +promptScript = + [ mkTextResponse proposeCallReply, + mkTextResponse finishReply, + mkTextResponse extractReply + ] + +-- | Native protocol: a tool-call block, a plain-text turn (no tool call = finish), +-- then the extract reply. +nativeScript :: [Response] +nativeScript = + [ mkToolCallResponse "call_1" "get_weather" weatherArgs, + mkTextResponse "I have the forecast now.", + mkTextResponse extractReply + ] + +-- | Prompt protocol with malformed tool arguments on the first turn: the tool +-- decode fails (recorded as an observation), the model then finishes, then +-- extracts. The agent still returns a typed answer. +badArgsPromptScript :: [Response] +badArgsPromptScript = + [ mkTextResponse proposeBadArgsReply, + mkTextResponse finishReply, + mkTextResponse extractReply + ] + +-- | A prompt-protocol script that never finishes: one tool-call propose, then the +-- extract reply. Run with @maxIters = 1@ it stops at the cap (after the single +-- proposal) and still extracts a best-effort answer. +maxItersScript :: [Response] +maxItersScript = + [ mkTextResponse proposeCallReply, + mkTextResponse extractReply + ] diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,28 @@ +-- | The @shikumi-tools@ test suite: typed tools, the ReAct loop, the protocol +-- seam, and the end-to-end acceptance — all network-free against a mock LM. +module Main (main) where + +import AcceptanceSpec qualified +import CodeActSpec qualified +import ProgramOfThoughtSpec qualified +import ProtocolSpec qualified +import ReActSpec qualified +import RestrictedSpec qualified +import SchemaSpec qualified +import Test.Tasty (defaultMain, testGroup) +import ToolSpec qualified + +main :: IO () +main = + defaultMain $ + testGroup + "shikumi-tools" + [ SchemaSpec.tests, + ToolSpec.tests, + ReActSpec.tests, + ProtocolSpec.tests, + AcceptanceSpec.tests, + RestrictedSpec.tests, + ProgramOfThoughtSpec.tests, + CodeActSpec.tests + ] diff --git a/test/MockLLM.hs b/test/MockLLM.hs new file mode 100644 --- /dev/null +++ b/test/MockLLM.hs @@ -0,0 +1,84 @@ +-- | A deterministic, network-free interpreter of the @LLM@ effect for the +-- @shikumi-tools@ specs, plus builders for the two kinds of scripted assistant +-- 'Response' a ReAct turn can produce: a plain-text reply (the prompt protocol and +-- the extract step) and a native tool-call block (the native protocol). +-- +-- Each 'Shikumi.LLM.Complete' pops the next scripted 'Response' in order; an +-- exhausted script yields an empty text response (specs always script enough). +module MockLLM + ( runMockLLM, + runEffMock, + runAgent, + mkTextResponse, + mkToolCallResponse, + ) +where + +import Baikai + ( AssistantContent (..), + Response, + _Response, + _TextContent, + _ToolCall, + ) +import Control.Lens ((&), (.~)) +import Data.Aeson (Value) +import Data.Generics.Labels () +import Data.IORef (IORef, atomicModifyIORef', newIORef) +import Data.Text (Text) +import Data.Vector qualified as V +import Effectful (Eff, IOE, liftIO, runEff, type (:>)) +import Effectful.Dispatch.Dynamic (interpret) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Shikumi.Error (ShikumiError) +import Shikumi.LLM (LLM (..)) +import Shikumi.Program (Program, runProgram) + +-- | Interpret @LLM@ by popping scripted 'Response's in order. The streaming op is +-- unused. An exhausted script yields an empty text response. +runMockLLM :: (IOE :> es) => [Response] -> Eff (LLM : es) a -> Eff es a +runMockLLM script act = do + ref <- liftIO (newIORef script) + interpret + ( \_ -> \case + Complete {} -> liftIO (pop ref) + Stream {} -> pure [] + ) + act + +-- | Discharge a network-free @LLM@ computation against a scripted mock: handle the +-- error channel, then @IO@. The whole stack is exactly the row a ReAct agent runs +-- in (@LLM@ + @Error ShikumiError@), plus @IOE@ at the bottom for the mock's @IORef@. +runEffMock :: + [Response] -> + Eff '[LLM, Error ShikumiError, IOE] a -> + IO (Either ShikumiError a) +runEffMock script = runEff . runErrorNoCallStack . runMockLLM script + +-- | Run a 'Program' (e.g. a @react@ agent) against a scripted mock LM. +runAgent :: + [Response] -> + Program i o -> + i -> + IO (Either ShikumiError o) +runAgent script prog i = runEffMock script (runProgram prog i) + +pop :: IORef [Response] -> IO Response +pop ref = atomicModifyIORef' ref step + where + step (x : xs) = (xs, x) + step [] = ([], mkTextResponse "") + +-- | An assistant 'Response' carrying @t@ as its single text block. +mkTextResponse :: Text -> Response +mkTextResponse t = + _Response & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t)) + +-- | An assistant 'Response' carrying a single native tool-call block. +mkToolCallResponse :: Text -> Text -> Value -> Response +mkToolCallResponse callId nm args = + _Response + & #message + . #content + .~ V.singleton + (AssistantToolCall (_ToolCall & #id_ .~ callId & #name .~ nm & #arguments .~ args)) diff --git a/test/ProgramOfThoughtSpec.hs b/test/ProgramOfThoughtSpec.hs new file mode 100644 --- /dev/null +++ b/test/ProgramOfThoughtSpec.hs @@ -0,0 +1,79 @@ +{-# LANGUAGE DeriveAnyClass #-} + +-- | EP-27 M2: @programOfThought@ — the model writes code, the sandbox runs it, the +-- result flows into a typed answer. Verified hermetically against the network-free +-- mock LM and the pure 'restrictedInterpreter'. The sandbox is shown to be +-- /load-bearing/ two ways: an always-fail interpreter makes the program give up +-- (proving the loop consults the interpreter's verdict), and an error-then-fix run +-- only reaches the answer because the interpreter really rejected @1 / 0@ and +-- accepted the correction. +module ProgramOfThoughtSpec (tests) where + +import Data.Text (Text) +import GHC.Generics (Generic) +import MockLLM (mkTextResponse, runAgent) +import Shikumi.Adapter (ToPrompt) +import Shikumi.CodeExec.Interpreter (CodeInterpreter (..)) +import Shikumi.CodeExec.ProgramOfThought (PoTConfig (..), programOfThought, programOfThoughtWith) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.Module (predict) +import Shikumi.Program (foldParams) +import Shikumi.Schema (FromModel, ToSchema) +import Shikumi.Signature (Signature, mkSignature) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +newtype Task = Task {task :: Text} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToPrompt, FromModel) + +newtype CalcAnswer = CalcAnswer {value :: Int} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +sig :: Signature Task CalcAnswer +sig = mkSignature "Compute the requested arithmetic result." + +-- | A sandbox that always rejects the code, to prove the loop consults it. +alwaysFail :: CodeInterpreter +alwaysFail = CodeInterpreter (\_ -> pure (Left "sandbox disabled")) + +tests :: TestTree +tests = + testGroup + "ProgramOfThoughtSpec" + [ testCase "solves an arithmetic task end-to-end (code -> run -> typed answer)" $ do + -- model emits the code, then (on extract) the typed answer JSON + out <- + runAgent + [mkTextResponse "37 * 19 + 6", mkTextResponse "{\"value\": 709}"] + (programOfThought sig) + (Task "multiply 37 by 19 and add 6") + out @?= Right (CalcAnswer 709), + testCase "a plain predict baseline returns the (wrong) direct guess" $ do + out <- + runAgent + [mkTextResponse "{\"value\": 700}"] + (predict sig) + (Task "multiply 37 by 19 and add 6") + out @?= Right (CalcAnswer 700), + testCase "the sandbox is load-bearing: an always-fail interpreter makes it give up" $ do + out <- + runAgent + [mkTextResponse "37 * 19 + 6", mkTextResponse "37 * 19 + 6"] + (programOfThoughtWith (PoTConfig {maxIters = 2, interpreter = alwaysFail}) sig) + (Task "multiply 37 by 19 and add 6") + case out of + Left (ProviderFailure _) -> pure () + other -> assertFailure ("expected ProviderFailure, got " <> show other), + testCase "error-then-fix: first snippet errors in the sandbox, second succeeds" $ do + -- 1/0 really errors in restrictedInterpreter -> regenerate -> 6 succeeds + out <- + runAgent + [mkTextResponse "1 / 0", mkTextResponse "6", mkTextResponse "{\"value\": 6}"] + (programOfThought sig) + (Task "compute six") + out @?= Right (CalcAnswer 6), + testCase "programOfThought carries no params (Embed node)" $ + assertBool "foldParams should be empty" (null (foldParams (programOfThought sig))) + ] diff --git a/test/ProtocolSpec.hs b/test/ProtocolSpec.hs new file mode 100644 --- /dev/null +++ b/test/ProtocolSpec.hs @@ -0,0 +1,63 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | M3: the native-vs-prompt protocol seam. The same @react@ program reaches the +-- same typed answer under either forced protocol, and 'ProtocolAuto' resolves +-- conservatively (a CLI model — whose tools baikai silently drops — picks the +-- prompt path; a native-capable model picks the native path). +module ProtocolSpec (tests) where + +import Baikai (Api (..), _Model) +import Control.Lens ((&), (.~)) +import Data.Generics.Labels () +import Data.Vector qualified as V +import Fixtures + ( WeatherResp, + expectedWeather, + nativeScript, + promptScript, + weatherQuestion, + weatherRegistry, + weatherSignature, + ) +import MockLLM (runAgent) +import Shikumi.Agent.ReAct + ( ReActConfig (..), + Termination (..), + ToolProtocol (..), + Trajectory (..), + defaultReActConfig, + reactWithTrajectory, + resolveProtocolKind, + ) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Protocol" + [ testCase "native and prompt paths agree on the typed answer" $ do + nat <- + runAgent + nativeScript + (reactWithTrajectory weatherSignature weatherRegistry defaultReActConfig {protocol = ProtocolNative}) + weatherQuestion + pro <- + runAgent + promptScript + (reactWithTrajectory weatherSignature weatherRegistry defaultReActConfig {protocol = ProtocolPrompt}) + weatherQuestion + case (nat, pro) of + (Right (oN :: WeatherResp, tN), Right (oP, tP)) -> do + oN @?= expectedWeather + oP @?= expectedWeather + termination tN @?= TerminatedFinish + termination tP @?= TerminatedFinish + V.length (steps tN) @?= 2 + V.length (steps tP) @?= 2 + _ -> assertFailure "both protocols should succeed", + testCase "ProtocolAuto picks prompt for a CLI model" $ + resolveProtocolKind ProtocolAuto (_Model & #api .~ AnthropicMessagesCli) @?= ProtocolPrompt, + testCase "ProtocolAuto picks native for a native-capable model" $ + resolveProtocolKind ProtocolAuto (_Model & #provider .~ "openai" & #api .~ OpenAIChatCompletions) @?= ProtocolNative + ] diff --git a/test/ReActSpec.hs b/test/ReActSpec.hs new file mode 100644 --- /dev/null +++ b/test/ReActSpec.hs @@ -0,0 +1,65 @@ +-- | M2: the ReAct loop as a 'Program', driven by a mock LM. A scripted tool-call +-- then finish yields the typed answer plus a two-step trajectory; a script that +-- never finishes stops at @maxIters@ and still extracts a best-effort answer. +module ReActSpec (tests) where + +import Data.Vector qualified as V +import Fixtures + ( WeatherResp, + expectedWeather, + maxItersScript, + promptScript, + weatherQuestion, + weatherRegistry, + weatherSignature, + ) +import MockLLM (runAgent) +import Shikumi.Agent.ReAct + ( Action (..), + ReActConfig (..), + Step (..), + Termination (..), + Trajectory (..), + defaultReActConfig, + reactWithTrajectory, + ) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "ReAct" + [ testCase "runs tool then finish; returns typed answer + trajectory" $ do + res <- + runAgent + promptScript + (reactWithTrajectory weatherSignature weatherRegistry defaultReActConfig) + weatherQuestion + case res of + Right (o :: WeatherResp, traj) -> do + o @?= expectedWeather + termination traj @?= TerminatedFinish + V.length (steps traj) @?= 2 + case V.toList (steps traj) of + [s1, s2] -> do + case action s1 of + CallTool nm _ -> nm @?= "get_weather" + other -> assertFailure ("step 1 should be a tool call, got " <> show other) + assertBool "tool step records an observation" (observation s1 /= Nothing) + action s2 @?= Finish + _ -> assertFailure "expected exactly two steps" + Left e -> assertFailure ("agent failed: " <> show e), + testCase "stops at maxIters with TerminatedMaxIters" $ do + let cfg = defaultReActConfig {maxIters = 1} + res <- + runAgent + maxItersScript + (reactWithTrajectory weatherSignature weatherRegistry cfg) + weatherQuestion + case res of + Right (o :: WeatherResp, traj) -> do + termination traj @?= TerminatedMaxIters 1 + o @?= expectedWeather + Left e -> assertFailure ("agent failed: " <> show e) + ] diff --git a/test/RestrictedSpec.hs b/test/RestrictedSpec.hs new file mode 100644 --- /dev/null +++ b/test/RestrictedSpec.hs @@ -0,0 +1,48 @@ +-- | EP-27 M1: the hermetic restricted-DSL evaluator behind 'restrictedInterpreter'. +-- Pure unit tests over 'runRestricted': it computes arithmetic/string/list results +-- and returns a @Left@ error message for division by zero, a parse error, or an +-- unknown function — the recoverable errors the @programOfThought@ loop feeds back. +module RestrictedSpec (tests) where + +import Data.Text (Text) +import Data.Text qualified as T +import Shikumi.CodeExec.Interpreter (runRestricted) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, testCase, (@?=)) + +assertLeft :: Either Text Text -> IO () +assertLeft r = assertBool ("expected Left, got " <> show r) (isLeft r) + where + isLeft (Left _) = True + isLeft _ = False + +assertLeftMentions :: Text -> Either Text Text -> IO () +assertLeftMentions needle = \case + Left msg -> assertBool ("error should mention " <> show needle <> ", got " <> show msg) (needle `T.isInfixOf` msg) + Right v -> assertBool ("expected Left mentioning " <> show needle <> ", got Right " <> show v) False + +tests :: TestTree +tests = + testGroup + "RestrictedSpec" + [ testCase "2 + 3 * 4 evaluates to 14 (precedence)" $ + runRestricted "2 + 3 * 4" @?= Right "14", + testCase "37 * 19 + 6 evaluates to 709" $ + runRestricted "37 * 19 + 6" @?= Right "709", + testCase "result = 42 evaluates to 42" $ + runRestricted "result = 42" @?= Right "42", + testCase "parentheses override precedence" $ + runRestricted "(2 + 3) * 4" @?= Right "20", + testCase "1 / 0 reports a division error" $ + assertLeftMentions "division" (runRestricted "1 / 0"), + testCase "len(\"hello\") evaluates to 5" $ + runRestricted "len(\"hello\")" @?= Right "5", + testCase "upper(\"ab\") ++ \"C\" concatenates to ABC" $ + runRestricted "upper(\"ab\") ++ \"C\"" @?= Right "ABC", + testCase "sum([1, 2, 3]) evaluates to 6" $ + runRestricted "sum([1, 2, 3])" @?= Right "6", + testCase "a malformed expression (extra paren) is a parse error" $ + assertLeft (runRestricted "upper(\"ab\") ++ \"C\")"), + testCase "an unknown function is an error" $ + assertLeftMentions "unknown" (runRestricted "frobnicate(3)") + ] diff --git a/test/SchemaSpec.hs b/test/SchemaSpec.hs new file mode 100644 --- /dev/null +++ b/test/SchemaSpec.hs @@ -0,0 +1,16 @@ +-- | M1: the typed tool's argument schema is the Generic-derived schema of its +-- input record, frozen against "Shikumi.Schema"'s real generator output. +module SchemaSpec (tests) where + +import Fixtures (expectedReqSchema, weatherTool) +import Shikumi.Tool (toolSchemaOf) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Schema" + [ testCase "derives WeatherReq argument schema" $ + toolSchemaOf weatherTool @?= expectedReqSchema + ] diff --git a/test/ToolSpec.hs b/test/ToolSpec.hs new file mode 100644 --- /dev/null +++ b/test/ToolSpec.hs @@ -0,0 +1,42 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | M1: the typed wire round-trip. A valid 'Baikai.ToolCall' decodes, runs, and +-- encodes its result; malformed arguments and an unknown tool name each yield a +-- typed 'ToolError' (a value), never an exception. +module ToolSpec (tests) where + +import Baikai (ToolCall, _ToolCall) +import Control.Lens ((&), (.~)) +import Data.Aeson (Value, object, (.=)) +import Data.Generics.Labels () +import Data.Text (Text) +import Data.Text qualified as T +import Fixtures (weatherArgs, weatherRegistry) +import MockLLM (runEffMock) +import Shikumi.Tool (ToolError (..), runToolCall) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +tc :: Text -> Value -> ToolCall +tc nm args = _ToolCall & #name .~ nm & #arguments .~ args + +tests :: TestTree +tests = + testGroup + "Tool" + [ testCase "decodes valid args and runs the body" $ do + res <- runEffMock [] (runToolCall weatherRegistry (tc "get_weather" weatherArgs)) + case res of + Right (Right obs) -> assertBool "observation mentions the forecast" ("mild" `T.isInfixOf` obs) + other -> assertFailure ("expected a tool result, got " <> show other), + testCase "returns ToolArgsInvalid for a missing required field" $ do + res <- runEffMock [] (runToolCall weatherRegistry (tc "get_weather" (object ["city" .= ("Paris" :: Text)]))) + case res of + Right (Left (ToolArgsInvalid nm _)) -> nm @?= "get_weather" + other -> assertFailure ("expected ToolArgsInvalid, got " <> show other), + testCase "returns ToolNotFound for an unknown name" $ do + res <- runEffMock [] (runToolCall weatherRegistry (tc "nope" (object []))) + case res of + Right (Left (ToolNotFound nm)) -> nm @?= "nope" + other -> assertFailure ("expected ToolNotFound, got " <> show other) + ]