diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,22 @@ ## Unreleased +## 0.2.0.0 - 2026-06-28 + +### Added + +- Built-in filesystem, shell, and web tool modules, plus helpers for assembling + built-in tool registries. +- `Shikumi.Tool.Env`, an execution-environment record for filesystem and process + backed tools. +- `Shikumi.Tool.Web`, a swappable HTTP/search client layer for web tools. + +### Changed + +- `ReActConfig` now includes a `compaction` field, and ReAct agents compact older + trajectory steps when usage approaches the model context window. +- Refreshed the internal `shikumi` bound for the `0.2` series. + ## 0.1.0.1 - 2026-06-21 ### Changed diff --git a/shikumi-tools.cabal b/shikumi-tools.cabal --- a/shikumi-tools.cabal +++ b/shikumi-tools.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: shikumi-tools -version: 0.1.0.1 +version: 0.2.0.0 synopsis: Typed tools and ReAct agents for shikumi LM programs (EP-11) @@ -49,18 +49,31 @@ Shikumi.CodeExec.ProgramOfThought Shikumi.CodeExec.Prompt Shikumi.Tool + Shikumi.Tool.Builtin + Shikumi.Tool.Builtin.Fs + Shikumi.Tool.Builtin.Shell + Shikumi.Tool.Builtin.Web + Shikumi.Tool.Env + Shikumi.Tool.Web build-depends: , aeson - , baikai >=0.2 && <0.3 - , base >=4.20 && <5 + , baikai >=0.2 && <0.3 + , base >=4.20 && <5 , bytestring , containers + , directory , effectful + , filepath , generic-lens - , lens ^>=5.3 - , shikumi ^>=0.1.0.1 - , text ^>=2.1 + , http-client + , http-client-tls + , http-types + , lens ^>=5.3 + , process + , regex-tdfa + , shikumi ^>=0.2.0.0 + , text ^>=2.1 , vector test-suite shikumi-tools-test @@ -71,26 +84,35 @@ ghc-options: -threaded -with-rtsopts=-N other-modules: AcceptanceSpec + BuiltinAcceptanceSpec CodeActSpec + CompactionSpec + EnvSpec Fixtures + FsSpec MockLLM ProgramOfThoughtSpec ProtocolSpec ReActSpec RestrictedSpec SchemaSpec + ShellSpec ToolSpec + WebSpec build-depends: , aeson , baikai >=0.2 && <0.3 , base + , bytestring , containers + , directory , effectful + , filepath , generic-lens , lens - , shikumi ^>=0.1.0.1 - , shikumi-tools ^>=0.1.0.1 + , shikumi ^>=0.2.0.0 + , shikumi-tools ^>=0.2.0.0 , tasty , tasty-hunit , text diff --git a/src/Shikumi/Agent/ReAct.hs b/src/Shikumi/Agent/ReAct.hs --- a/src/Shikumi/Agent/ReAct.hs +++ b/src/Shikumi/Agent/ReAct.hs @@ -58,6 +58,7 @@ Tool, ToolCall, ToolChoice (..), + Usage, flattenAssistantBlocks, user, _Context, @@ -77,9 +78,10 @@ import Data.Vector (Vector) import Data.Vector qualified as V import Effectful (Eff, (:>)) -import Effectful.Error.Static (Error, throwError) +import Effectful.Error.Static (Error, catchError, throwError) import GHC.Generics (Generic) import Shikumi.Adapter (ModelCapability (..), ToPrompt (toPrompt), attachSchema, capabilityFor) +import Shikumi.Compaction (CompactionConfig, compactTail, defaultCompactionConfig, usageExceedsWindow) import Shikumi.Error (ShikumiError (..)) import Shikumi.LLM (LLM, complete) import Shikumi.Program (Program (FMap), embed) @@ -139,13 +141,19 @@ -- | How an agent runs: the hard iteration cap and the protocol selector. data ReActConfig = ReActConfig { maxIters :: !Int, - protocol :: !ToolProtocol + protocol :: !ToolProtocol, + compaction :: !CompactionConfig } deriving stock (Show, Eq, Generic) -- | Six iterations, protocol auto-selected. defaultReActConfig :: ReActConfig -defaultReActConfig = ReActConfig {maxIters = 6, protocol = ProtocolAuto} +defaultReActConfig = + ReActConfig + { maxIters = 6, + protocol = ProtocolAuto, + compaction = defaultCompactionConfig + } -- --------------------------------------------------------------------------- -- Building agents @@ -184,8 +192,8 @@ i -> Eff es (o, Trajectory) reactLoop sig reg cfg i = do - traj <- loop 0 [] - o <- extract traj + traj0 <- loop 0 [] + (o, traj) <- extract traj0 pure (o, traj) where impl :: ProtocolImpl i o @@ -197,25 +205,63 @@ | 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 + (accForPrompt, resp) <- completeProposeRecover acc case parsePropose impl resp of - Left perr -> loop (iter + 1) (correctiveStep perr : acc) + Left perr -> loop (iter + 1) (correctiveStep perr : accForPrompt) Right (th, Finish) -> - pure (Trajectory (V.fromList (reverse (Step th Finish Nothing : acc))) TerminatedFinish) + pure (Trajectory (V.fromList (reverse (Step th Finish Nothing : accForPrompt))) 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) + acc' = Step th (CallTool nm args) (Just obs) : accForPrompt + acc'' <- compactAcc (resp ^. #model) (resp ^. #message . #usage) acc' + loop (iter + 1) 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 :: Trajectory -> Eff es (o, Trajectory) extract traj = do let (ctx, opts) = renderExtract impl i traj - resp <- complete _Model ctx opts - either throwError pure (parseExtract impl resp) + (trajForExtract, resp) <- + catchError + ((traj,) <$> complete _Model ctx opts) + ( \_cs -> \case + ContextWindowExceeded {} -> do + compacted <- forceCompactTrajectory traj + let (ctx', opts') = renderExtract impl i compacted + (compacted,) <$> complete _Model ctx' opts' + e -> throwError e + ) + o <- either throwError pure (parseExtract impl resp) + pure (o, trajForExtract) + completeProposeRecover :: [Step] -> Eff es ([Step], Response) + completeProposeRecover acc = do + let (ctx, opts) = renderPropose impl i (soFar acc) + catchError + ((acc,) <$> complete _Model ctx opts) + ( \_cs -> \case + ContextWindowExceeded {} -> do + compacted <- forceCompactAcc acc + let (ctx', opts') = renderPropose impl i (soFar compacted) + (compacted,) <$> complete _Model ctx' opts' + e -> throwError e + ) + + compactAcc :: Model -> Usage -> [Step] -> Eff es [Step] + compactAcc model usage acc + | usageExceedsWindow (compaction cfg) model usage = forceCompactAcc acc + | otherwise = pure acc + + forceCompactAcc :: [Step] -> Eff es [Step] + forceCompactAcc acc = + reverse <$> compactTail (compaction cfg) _Model renderStepLine summaryStep (reverse acc) + + forceCompactTrajectory :: Trajectory -> Eff es Trajectory + forceCompactTrajectory traj = do + compacted <- compactTail (compaction cfg) _Model renderStepLine summaryStep (V.toList (steps traj)) + pure (traj {steps = V.fromList compacted}) + -- A trajectory view of the steps gathered so far (termination is irrelevant -- for the propose render). soFar :: [Step] -> Trajectory @@ -233,6 +279,17 @@ Just ("Your previous reply was not a valid action JSON object: " <> perr <> ". Reply with exactly one action object.") } +summaryStep :: Text -> Step +summaryStep summaryText = + Step + { thought = "(compacted summary of earlier steps)", + action = CallTool "" Null, + observation = Just summaryText + } + +renderStepLine :: Step -> Text +renderStepLine s = renderTrajectory (Trajectory (V.singleton s) TerminatedFinish) + -- | 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 @@ -246,10 +303,10 @@ -- 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 + { 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. diff --git a/src/Shikumi/Tool.hs b/src/Shikumi/Tool.hs --- a/src/Shikumi/Tool.hs +++ b/src/Shikumi/Tool.hs @@ -223,5 +223,6 @@ SchemaMismatch t -> t ValidationFailure t -> t ProviderFailure t -> t + ContextWindowExceeded t -> t Timeout t -> t BudgetExceeded t -> t diff --git a/src/Shikumi/Tool/Builtin.hs b/src/Shikumi/Tool/Builtin.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Tool/Builtin.hs @@ -0,0 +1,40 @@ +-- | Built-in work-tool catalog. +module Shikumi.Tool.Builtin + ( builtinFsTools, + builtinWebTools, + builtinTools, + builtinRegistry, + module Shikumi.Tool.Builtin.Fs, + module Shikumi.Tool.Builtin.Shell, + module Shikumi.Tool.Builtin.Web, + ) +where + +import Shikumi.Tool (SomeTool (..), ToolRegistry, mkRegistry) +import Shikumi.Tool.Builtin.Fs +import Shikumi.Tool.Builtin.Shell +import Shikumi.Tool.Builtin.Web +import Shikumi.Tool.Env (ToolEnv) +import Shikumi.Tool.Web (WebClient) + +builtinFsTools :: ToolEnv -> [SomeTool] +builtinFsTools env = + [ SomeTool (readTool env), + SomeTool (writeTool env), + SomeTool (editTool env), + SomeTool (grepTool env), + SomeTool (bashTool env), + SomeTool (globTool env) + ] + +builtinWebTools :: WebClient -> [SomeTool] +builtinWebTools client = + [ SomeTool (webFetchTool client), + SomeTool (webSearchTool client) + ] + +builtinTools :: ToolEnv -> WebClient -> [SomeTool] +builtinTools env client = builtinFsTools env <> builtinWebTools client + +builtinRegistry :: ToolEnv -> WebClient -> ToolRegistry +builtinRegistry env client = mkRegistry (builtinTools env client) diff --git a/src/Shikumi/Tool/Builtin/Fs.hs b/src/Shikumi/Tool/Builtin/Fs.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Tool/Builtin/Fs.hs @@ -0,0 +1,474 @@ +{-# LANGUAGE RankNTypes #-} + +-- | Built-in filesystem tools over 'ToolEnv'. +-- +-- The search tools prefer host-provided @rg@/@fd@ through 'ToolEnv.exec' when +-- available, and fall back to a bounded in-process traversal through the same +-- environment record. The fallback skips noisy directories, binary files, and +-- large files, and caps traversal depth and result count. +module Shikumi.Tool.Builtin.Fs + ( ReadReq (..), + ReadResp (..), + WriteReq (..), + WriteResp (..), + EditReq (..), + EditResp (..), + GrepReq (..), + GrepMatch (..), + GrepResp (..), + GlobReq (..), + GlobResp (..), + readTool, + writeTool, + editTool, + grepTool, + globTool, + ) +where + +import Control.Lens ((^.)) +import Data.Aeson (ToJSON, Value (..), eitherDecodeStrict, object, (.=)) +import Data.Aeson.Key qualified as Key +import Data.Aeson.KeyMap qualified as KM +import Data.Aeson.Types (parseJSON, parseMaybe) +import Data.ByteString qualified as BS +import Data.Char (isAlphaNum) +import Data.Generics.Labels () +import Data.List (sort) +import Data.Maybe (mapMaybe) +import Data.Proxy (Proxy (..)) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Encoding qualified as TE +import Data.Text.Encoding.Error (lenientDecode) +import Effectful (Eff) +import Effectful.Error.Static (throwError) +import GHC.Generics (Generic) +import Shikumi.Adapter (ToPrompt) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.Schema (FromModel (..), ToSchema (..), fromModel) +import Shikumi.Tool (Tool, mkTool) +import Shikumi.Tool.Env + ( EnvRow, + ExecRequest (..), + ToolEnv (..), + ) +import Text.Regex.TDFA + ( CompOption (..), + Regex, + defaultCompOpt, + defaultExecOpt, + makeRegexOptsM, + matchTest, + (=~), + ) + +data ReadReq = ReadReq + { path :: !Text, + offset :: !(Maybe Int), + limit :: !(Maybe Int) + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +data ReadResp = ReadResp + { content :: !Text, + lineCount :: !Int, + truncated :: !Bool + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToJSON) + +data WriteReq = WriteReq + { path :: !Text, + content :: !Text + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +data WriteResp = WriteResp + { path :: !Text, + bytesWritten :: !Int + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToJSON) + +data EditReq = EditReq + { path :: !Text, + oldString :: !Text, + newString :: !Text, + replaceAll :: !(Maybe Bool) + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +data EditResp = EditResp + { path :: !Text, + replacements :: !Int + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToJSON) + +data GrepReq = GrepReq + { patternText :: !Text, + path :: !(Maybe Text), + glob :: !(Maybe Text), + ignoreCase :: !(Maybe Bool) + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToPrompt) + +instance ToSchema GrepReq where + toSchema _ = + toolInputSchema + [ ("pattern", toSchema (Proxy :: Proxy Text)), + ("path", toSchema (Proxy :: Proxy (Maybe Text))), + ("glob", toSchema (Proxy :: Proxy (Maybe Text))), + ("ignoreCase", toSchema (Proxy :: Proxy (Maybe Bool))) + ] + +instance FromModel GrepReq where + fromModelP _ = \case + Object obj -> + GrepReq + <$> requiredField "pattern" obj + <*> optionalField "path" obj + <*> optionalField "glob" obj + <*> optionalField "ignoreCase" obj + v -> Left (SchemaMismatch ("grep request must be an object, got " <> T.pack (show v))) + +data GrepMatch = GrepMatch + { file :: !Text, + line :: !Int, + text :: !Text + } + deriving stock (Generic, Show, Eq, Ord) + deriving anyclass (ToJSON) + +data GrepResp = GrepResp + { matches :: ![GrepMatch], + truncated :: !Bool + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToJSON) + +data GlobReq = GlobReq + { patternText :: !Text, + path :: !(Maybe Text) + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToPrompt) + +instance ToSchema GlobReq where + toSchema _ = + toolInputSchema + [ ("pattern", toSchema (Proxy :: Proxy Text)), + ("path", toSchema (Proxy :: Proxy (Maybe Text))) + ] + +instance FromModel GlobReq where + fromModelP _ = \case + Object obj -> + GlobReq + <$> requiredField "pattern" obj + <*> optionalField "path" obj + v -> Left (SchemaMismatch ("glob request must be an object, got " <> T.pack (show v))) + +data GlobResp = GlobResp + { paths :: ![Text], + truncated :: !Bool + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToJSON) + +readTool :: ToolEnv -> Tool ReadReq ReadResp +readTool env = + mkTool "read" "Read a file's contents." $ \req -> do + bs <- envReadFile env (req ^. #path) + let allLines = T.lines (decodeUtf8 bs) + start = max 0 (maybe 0 id (req ^. #offset)) + afterOffset = drop start allLines + selected = maybe afterOffset (`take` afterOffset) (req ^. #limit) + wasTruncated = start > 0 || length selected < length afterOffset + pure ReadResp {content = T.unlines selected, lineCount = length selected, truncated = wasTruncated} + +writeTool :: ToolEnv -> Tool WriteReq WriteResp +writeTool env = + mkTool "write" "Write UTF-8 text to a file." $ \req -> do + let bytes = TE.encodeUtf8 (req ^. #content) + envWriteFile env (req ^. #path) bytes + pure WriteResp {path = req ^. #path, bytesWritten = BS.length bytes} + +editTool :: ToolEnv -> Tool EditReq EditResp +editTool env = + mkTool "edit" "Replace text in a UTF-8 file." $ \req -> do + let old = req ^. #oldString + if T.null old + then throwValidation "edit: oldString must not be empty" + else do + original <- decodeUtf8 <$> envReadFile env (req ^. #path) + (edited, count) <- + if maybe False id (req ^. #replaceAll) + then pure (replaceAllCount old (req ^. #newString) original) + else pure (replaceFirstCount old (req ^. #newString) original) + if count == 0 + then throwValidation "edit: oldString not found" + else do + envWriteFile env (req ^. #path) (TE.encodeUtf8 edited) + pure EditResp {path = req ^. #path, replacements = count} + +grepTool :: ToolEnv -> Tool GrepReq GrepResp +grepTool env = + mkTool "grep" "Search files with a regex and return structured matches." $ \req -> do + root <- maybe (envCwd env) pure (req ^. #path) + fast <- commandPresent env "rg" + if fast + then do + resp <- grepFast env root req + case resp of + Just value -> pure value + Nothing -> grepFallback env root req + else grepFallback env root req + +globTool :: ToolEnv -> Tool GlobReq GlobResp +globTool env = + mkTool "glob" "Find files matching a glob pattern and return structured paths." $ \req -> do + root <- maybe (envCwd env) pure (req ^. #path) + fast <- commandPresent env "fd" + let pat = req ^. #patternText + if fast + then do + resp <- globFast env root pat + case resp of + Just value -> pure value + Nothing -> globFallback env root pat + else globFallback env root pat + +grepFast :: (EnvRow es) => ToolEnv -> Text -> GrepReq -> Eff es (Maybe GrepResp) +grepFast env root req = do + result <- + envExec + env + ExecRequest + { command = + T.unwords + [ "rg --json --color never --max-filesize 5M", + if maybe False id (req ^. #ignoreCase) then "-i" else "", + maybe "" (\g -> "-g " <> shellQuote g) (req ^. #glob), + "--", + shellQuote (req ^. #patternText), + shellQuote root + ], + cwd = Nothing, + stdin = Nothing, + timeoutMs = Just 30000 + } + if result ^. #exitCode == 0 || result ^. #exitCode == 1 + then pure (Just (parseRgJson (result ^. #stdout))) + else pure Nothing + +globFast :: (EnvRow es) => ToolEnv -> Text -> Text -> Eff es (Maybe GlobResp) +globFast env root pat = do + result <- + envExec + env + ExecRequest + { command = + T.unwords + [ "fd --glob --color never --exclude .git --exclude node_modules --exclude dist-newstyle --exclude .stack-work", + "--", + shellQuote pat, + shellQuote root + ], + cwd = Nothing, + stdin = Nothing, + timeoutMs = Just 30000 + } + if result ^. #exitCode == 0 + then do + let found = take maxResults (filter (not . T.null) (T.lines (result ^. #stdout))) + pure (Just GlobResp {paths = sort found, truncated = length found >= maxResults}) + else pure Nothing + +grepFallback :: (EnvRow es) => ToolEnv -> Text -> GrepReq -> Eff es GrepResp +grepFallback env root req = do + regex <- compileRegex (req ^. #patternText) (maybe False id (req ^. #ignoreCase)) + paths <- walkFiles env root + let globOk p = maybe True (`globMatches` p) (req ^. #glob) + matches <- + fmap concat $ + traverse + ( \filePath -> do + if globOk filePath + then grepFile env regex filePath + else pure [] + ) + paths + let limited = take maxResults matches + pure GrepResp {matches = limited, truncated = length matches > maxResults} + +globFallback :: (EnvRow es) => ToolEnv -> Text -> Text -> Eff es GlobResp +globFallback env root pat = do + paths <- walkFiles env root + let found = sort (take maxResults (filter (globMatches pat) paths)) + pure GlobResp {paths = found, truncated = length found >= maxResults} + +grepFile :: (EnvRow es) => ToolEnv -> Regex -> Text -> Eff es [GrepMatch] +grepFile env regex filePath = do + stat <- envStat env filePath + case stat of + Just s | s ^. #isFile && s ^. #size <= maxFileSize -> do + bs <- envReadFile env filePath + if isBinary bs + then pure [] + else do + let lineMatches = regexMatches regex + pure + [ GrepMatch {file = filePath, line = n, text = lineText} + | (n, lineText) <- zip [1 ..] (T.lines (decodeUtf8 bs)), + lineMatches lineText + ] + _ -> pure [] + +walkFiles :: (EnvRow es) => ToolEnv -> Text -> Eff es [Text] +walkFiles env root = go 0 root + where + go depth dir + | depth > maxDepth = pure [] + | otherwise = do + entries <- envReaddir env dir + fmap concat $ + traverse + ( \entry -> do + let child = joinPathText dir (entry ^. #name) + if entry ^. #isDir + then + if skipDirName (entry ^. #name) + then pure [] + else go (depth + 1) child + else pure [child] + ) + entries + +parseRgJson :: Text -> GrepResp +parseRgJson output = + let matches = take maxResults (mapMaybe parseLine (T.lines output)) + in GrepResp {matches, truncated = length matches >= maxResults} + where + parseLine line = do + Object obj <- either (const Nothing) Just (eitherDecodeStrict (TE.encodeUtf8 line)) + String "match" <- KM.lookup "type" obj + Object dat <- KM.lookup "data" obj + Object pathObj <- KM.lookup "path" dat + String file <- KM.lookup "text" pathObj + Object linesObj <- KM.lookup "lines" dat + String text <- KM.lookup "text" linesObj + Number lineNo <- KM.lookup "line_number" dat + n <- parseMaybe parseJSON (Number lineNo) + pure GrepMatch {file, line = n, text = T.dropWhileEnd (== '\n') text} + +compileRegex :: (EnvRow es) => Text -> Bool -> Eff es Regex +compileRegex pat ignoreCase = + case makeRegexOptsM defaultCompOpt {caseSensitive = not ignoreCase} defaultExecOpt (T.unpack pat) of + Just regex -> pure regex + Nothing -> throwValidation ("grep: invalid regex pattern " <> pat) + +regexMatches :: Regex -> Text -> Bool +regexMatches regex haystack = matchTest regex (T.unpack haystack) + +globMatches :: Text -> Text -> Bool +globMatches pat candidate = + T.unpack candidate =~ ("^" <> T.unpack (globToRegex pat) <> "$" :: String) + +globToRegex :: Text -> Text +globToRegex = go . T.unpack + where + go [] = "" + go ('*' : '*' : rest) = ".*" <> go rest + go ('*' : rest) = "[^/]*" <> go rest + go ('?' : rest) = "[^/]" <> go rest + go (c : rest) = T.pack (escapeRegex c) <> go rest + +replaceFirstCount :: Text -> Text -> Text -> (Text, Int) +replaceFirstCount old new src = + let (before, after) = T.breakOn old src + in if T.null after + then (src, 0) + else (before <> new <> T.drop (T.length old) after, 1) + +replaceAllCount :: Text -> Text -> Text -> (Text, Int) +replaceAllCount old new src = + let pieces = T.splitOn old src + count = length pieces - 1 + in (T.intercalate new pieces, count) + +commandPresent :: (EnvRow es) => ToolEnv -> Text -> Eff es Bool +commandPresent env cmd = do + result <- + envExec + env + ExecRequest + { command = "command -v " <> cmd, + cwd = Nothing, + stdin = Nothing, + timeoutMs = Just 5000 + } + pure (result ^. #exitCode == 0) + +decodeUtf8 :: BS.ByteString -> Text +decodeUtf8 = TE.decodeUtf8With lenientDecode + +isBinary :: BS.ByteString -> Bool +isBinary = BS.elem 0 . BS.take 8192 + +skipDirName :: Text -> Bool +skipDirName name = + name `elem` [".git", ".hg", ".svn", "node_modules", "dist-newstyle", ".stack-work", ".direnv"] + +joinPathText :: Text -> Text -> Text +joinPathText base name + | T.null base = name + | "/" `T.isSuffixOf` base = base <> name + | otherwise = base <> "/" <> name + +shellQuote :: Text -> Text +shellQuote t = "'" <> T.replace "'" "'\\''" t <> "'" + +escapeRegex :: Char -> String +escapeRegex c + | isAlphaNum c || c == '/' || c == '_' || c == '-' = [c] + | otherwise = ['\\', c] + +throwValidation :: (EnvRow es) => Text -> Eff es a +throwValidation = throwError . ValidationFailure + +toolInputSchema :: [(Text, Value)] -> Value +toolInputSchema fields = + object + [ "type" .= ("object" :: Text), + "properties" .= object (map (\(key, value) -> Key.fromText key .= value) fields), + "required" .= (["pattern"] :: [Text]), + "additionalProperties" .= False + ] + +requiredField :: (FromModel a) => Text -> KM.KeyMap Value -> Either ShikumiError a +requiredField key obj = + case KM.lookup (Key.fromText key) obj of + Nothing -> Left (MissingField key) + Just value -> fromModel value + +optionalField :: (FromModel a) => Text -> KM.KeyMap Value -> Either ShikumiError (Maybe a) +optionalField key obj = + case KM.lookup (Key.fromText key) obj of + Nothing -> Right Nothing + Just Null -> Right Nothing + Just value -> Just <$> fromModel value + +maxResults :: Int +maxResults = 1000 + +maxFileSize :: Integer +maxFileSize = 5 * 1024 * 1024 + +maxDepth :: Int +maxDepth = 25 diff --git a/src/Shikumi/Tool/Builtin/Shell.hs b/src/Shikumi/Tool/Builtin/Shell.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Tool/Builtin/Shell.hs @@ -0,0 +1,53 @@ +-- | Built-in shell tool over 'ToolEnv.exec'. +module Shikumi.Tool.Builtin.Shell + ( BashReq (..), + BashResp (..), + bashTool, + ) +where + +import Control.Lens ((^.)) +import Data.Aeson (ToJSON) +import Data.Generics.Labels () +import Data.Text (Text) +import GHC.Generics (Generic) +import Shikumi.Adapter (ToPrompt) +import Shikumi.Schema (FromModel, ToSchema) +import Shikumi.Tool (Tool, mkTool) +import Shikumi.Tool.Env (ExecRequest (..), ToolEnv (..)) + +data BashReq = BashReq + { command :: !Text, + cwd :: !(Maybe Text), + timeoutMs :: !(Maybe Int), + stdin :: !(Maybe Text) + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +data BashResp = BashResp + { exitCode :: !Int, + stdout :: !Text, + stderr :: !Text + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToJSON) + +bashTool :: ToolEnv -> Tool BashReq BashResp +bashTool env = + mkTool "bash" "Run a shell command and capture stdout, stderr, and the exit code." $ \req -> do + result <- + envExec + env + ExecRequest + { command = req ^. #command, + cwd = req ^. #cwd, + timeoutMs = req ^. #timeoutMs, + stdin = req ^. #stdin + } + pure + BashResp + { exitCode = result ^. #exitCode, + stdout = result ^. #stdout, + stderr = result ^. #stderr + } diff --git a/src/Shikumi/Tool/Builtin/Web.hs b/src/Shikumi/Tool/Builtin/Web.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Tool/Builtin/Web.hs @@ -0,0 +1,41 @@ +-- | Built-in web tools. +module Shikumi.Tool.Builtin.Web + ( FetchReq (..), + SearchReq (..), + webFetchTool, + webSearchTool, + ) +where + +import Control.Lens ((^.)) +import Data.Generics.Labels () +import Data.Text (Text) +import GHC.Generics (Generic) +import Shikumi.Adapter (ToPrompt) +import Shikumi.Schema (FromModel, ToSchema) +import Shikumi.Tool (Tool, mkTool) +import Shikumi.Tool.Web (FetchResult, SearchResult, WebClient (..)) + +data FetchReq = FetchReq + { url :: !Text, + maxBytes :: !(Maybe Int) + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +data SearchReq = SearchReq + { query :: !Text, + maxResults :: !(Maybe Int) + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +webFetchTool :: WebClient -> Tool FetchReq FetchResult +webFetchTool client = + mkTool "web_fetch" "Fetch the contents of a URL over HTTP." $ \req -> + webFetch client (req ^. #url) (req ^. #maxBytes) + +webSearchTool :: WebClient -> Tool SearchReq SearchResult +webSearchTool client = + mkTool "web_search" "Search the web with the configured search provider." $ \req -> + webSearch client (req ^. #query) (req ^. #maxResults) diff --git a/src/Shikumi/Tool/Env.hs b/src/Shikumi/Tool/Env.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Tool/Env.hs @@ -0,0 +1,178 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE RankNTypes #-} + +-- | Execution-environment operations for built-in work tools. +-- +-- Filesystem and process-backed tools should depend on this record instead of +-- calling the host operating system directly. A future sandbox can provide a +-- different 'ToolEnv' value while the tool definitions stay unchanged. +module Shikumi.Tool.Env + ( Path, + EnvRow, + ExecRequest (..), + ExecResult (..), + FileStat (..), + DirEntry (..), + ToolEnv (..), + localToolEnv, + ) +where + +import Control.Exception (IOException, try) +import Control.Lens ((^.)) +import Data.Aeson (ToJSON) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.Generics.Labels () +import Data.Text (Text) +import Data.Text qualified as T +import Effectful (Eff, (:>)) +import Effectful.Dispatch.Static (unsafeEff_) +import Effectful.Error.Static (Error, throwError) +import GHC.Generics (Generic) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.LLM (LLM) +import System.Directory qualified as Dir +import System.Exit (ExitCode (..)) +import System.Process qualified as Proc +import System.Timeout qualified as Timeout + +type Path = Text + +type EnvRow es = (LLM :> es, Error ShikumiError :> es) + +data ExecRequest = ExecRequest + { command :: !Text, + cwd :: !(Maybe Path), + stdin :: !(Maybe Text), + timeoutMs :: !(Maybe Int) + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToJSON) + +data ExecResult = ExecResult + { exitCode :: !Int, + stdout :: !Text, + stderr :: !Text + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToJSON) + +data FileStat = FileStat + { isFile :: !Bool, + isDir :: !Bool, + size :: !Integer + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToJSON) + +data DirEntry = DirEntry + { name :: !Text, + isDir :: !Bool + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToJSON) + +data ToolEnv = ToolEnv + { envExec :: forall es. (EnvRow es) => ExecRequest -> Eff es ExecResult, + envReadFile :: forall es. (EnvRow es) => Path -> Eff es ByteString, + envWriteFile :: forall es. (EnvRow es) => Path -> ByteString -> Eff es (), + envStat :: forall es. (EnvRow es) => Path -> Eff es (Maybe FileStat), + envReaddir :: forall es. (EnvRow es) => Path -> Eff es [DirEntry], + envExists :: forall es. (EnvRow es) => Path -> Eff es Bool, + envMkdir :: forall es. (EnvRow es) => Path -> Eff es (), + envRm :: forall es. (EnvRow es) => Path -> Eff es (), + envCwd :: forall es. (EnvRow es) => Eff es Path + } + +localToolEnv :: ToolEnv +localToolEnv = + ToolEnv + { envExec = localExec, + envReadFile = \path -> toolIO "readFile" (BS.readFile (T.unpack path)), + envWriteFile = \path bytes -> toolIO "writeFile" (BS.writeFile (T.unpack path) bytes), + envStat = localStat, + envReaddir = localReaddir, + envExists = \path -> toolIO "exists" (Dir.doesPathExist (T.unpack path)), + envMkdir = \path -> toolIO "mkdir" (Dir.createDirectoryIfMissing True (T.unpack path)), + envRm = localRm, + envCwd = T.pack <$> toolIO "cwd" Dir.getCurrentDirectory + } + +localExec :: (EnvRow es) => ExecRequest -> Eff es ExecResult +localExec req = do + let effectiveTimeoutMs = timeoutMsOrDefault (req ^. #timeoutMs) + result <- + toolIO "exec" $ + Timeout.timeout (effectiveTimeoutMs * 1000) $ + Proc.readCreateProcessWithExitCode process stdinText + case result of + Nothing -> + throwError $ + Timeout + ("tool env: exec timed out after " <> T.pack (show effectiveTimeoutMs) <> "ms") + Just (exit, out, err) -> + pure + ExecResult + { exitCode = exitCodeInt exit, + stdout = T.pack out, + stderr = T.pack err + } + where + process = + (Proc.proc "bash" ["-c", T.unpack (req ^. #command)]) + { Proc.cwd = T.unpack <$> (req ^. #cwd) + } + stdinText = maybe "" T.unpack (req ^. #stdin) + +localStat :: (EnvRow es) => Path -> Eff es (Maybe FileStat) +localStat path = toolIO "stat" $ do + let fp = T.unpack path + file <- Dir.doesFileExist fp + dir <- Dir.doesDirectoryExist fp + if file + then do + size <- Dir.getFileSize fp + pure (Just FileStat {isFile = True, isDir = False, size}) + else + if dir + then pure (Just FileStat {isFile = False, isDir = True, size = 0}) + else pure Nothing + +localReaddir :: (EnvRow es) => Path -> Eff es [DirEntry] +localReaddir path = toolIO "readdir" $ do + let dir = T.unpack path + names <- Dir.listDirectory dir + traverse + ( \entry -> do + isDir <- Dir.doesDirectoryExist (dir <> "/" <> entry) + pure DirEntry {name = T.pack entry, isDir} + ) + names + +localRm :: (EnvRow es) => Path -> Eff es () +localRm path = toolIO "rm" $ do + let fp = T.unpack path + file <- Dir.doesFileExist fp + dir <- Dir.doesDirectoryExist fp + if file + then Dir.removeFile fp + else + if dir + then Dir.removeDirectory fp + else pure () + +toolIO :: (EnvRow es) => Text -> IO a -> Eff es a +toolIO label action = do + result <- unsafeEff_ (try action) + case result of + Right value -> pure value + Left (err :: IOException) -> + throwError (ProviderFailure ("tool env: " <> label <> ": " <> T.pack (show err))) + +timeoutMsOrDefault :: Maybe Int -> Int +timeoutMsOrDefault = maybe 60000 id + +exitCodeInt :: ExitCode -> Int +exitCodeInt ExitSuccess = 0 +exitCodeInt (ExitFailure n) = n diff --git a/src/Shikumi/Tool/Web.hs b/src/Shikumi/Tool/Web.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Tool/Web.hs @@ -0,0 +1,139 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE RankNTypes #-} + +-- | Swappable HTTP operations for the built-in web tools. +-- +-- 'web_fetch' is functional with only a TLS 'Manager'. 'web_search' requires a +-- configured provider because search APIs need provider-specific credentials. +module Shikumi.Tool.Web + ( FetchResult (..), + SearchHit (..), + SearchResult (..), + SearchConfig (..), + WebClient (..), + localWebClient, + newTlsManager, + Manager, + ) +where + +import Control.Exception (try) +import Control.Lens ((^.)) +import Data.Aeson (FromJSON, ToJSON, eitherDecode) +import Data.ByteString.Char8 qualified as B8 +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 qualified as TE +import Data.Text.Encoding.Error (lenientDecode) +import Effectful (Eff) +import Effectful.Dispatch.Static (unsafeEff_) +import Effectful.Error.Static (throwError) +import GHC.Generics (Generic) +import Network.HTTP.Client + ( HttpException, + Manager, + httpLbs, + parseRequest, + responseBody, + responseHeaders, + responseStatus, + setQueryString, + ) +import Network.HTTP.Client.TLS qualified as TLS +import Network.HTTP.Types.Header (ResponseHeaders) +import Network.HTTP.Types.Status (statusCode) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.Tool.Env (EnvRow) + +data FetchResult = FetchResult + { status :: !Int, + contentType :: !Text, + body :: !Text, + truncated :: !Bool + } + deriving stock (Generic, Show, Eq) + deriving anyclass (FromJSON, ToJSON) + +data SearchHit = SearchHit + { title :: !Text, + url :: !Text, + snippet :: !Text + } + deriving stock (Generic, Show, Eq) + deriving anyclass (FromJSON, ToJSON) + +newtype SearchResult = SearchResult {hits :: [SearchHit]} + deriving stock (Generic, Show, Eq) + deriving anyclass (FromJSON, ToJSON) + +data SearchConfig = SearchConfig + { baseUrl :: !Text, + apiKey :: !Text + } + deriving stock (Generic, Show, Eq) + deriving anyclass (ToJSON) + +data WebClient = WebClient + { webFetch :: forall es. (EnvRow es) => Text -> Maybe Int -> Eff es FetchResult, + webSearch :: forall es. (EnvRow es) => Text -> Maybe Int -> Eff es SearchResult + } + +localWebClient :: Manager -> Maybe SearchConfig -> WebClient +localWebClient manager searchConfig = + WebClient + { webFetch = \url maxBytes -> localFetch manager url maxBytes, + webSearch = \query maxResults -> localSearch manager searchConfig query maxResults + } + +newTlsManager :: IO Manager +newTlsManager = TLS.newTlsManager + +localFetch :: (EnvRow es) => Manager -> Text -> Maybe Int -> Eff es FetchResult +localFetch manager url maxBytes = do + response <- httpIO "web_fetch" $ do + request <- parseRequest (T.unpack url) + httpLbs request manager + let maxBodyBytes = max 0 (maybe 100000 id maxBytes) + rawBody = responseBody response + limitedBody = LBS.take (fromIntegral maxBodyBytes) rawBody + truncated = LBS.length rawBody > fromIntegral maxBodyBytes + decodedBody = TE.decodeUtf8With lenientDecode (LBS.toStrict limitedBody) + pure + FetchResult + { status = statusCode (responseStatus response), + contentType = contentTypeOf (responseHeaders response), + body = decodedBody, + truncated + } + +localSearch :: (EnvRow es) => Manager -> Maybe SearchConfig -> Text -> Maybe Int -> Eff es SearchResult +localSearch _ Nothing _ _ = + throwError (ProviderFailure "web_search: no search provider configured") +localSearch manager (Just config) query maxResults = do + response <- httpIO "web_search" $ do + request0 <- parseRequest (T.unpack (config ^. #baseUrl)) + let request = + setQueryString + [ ("q", Just (TE.encodeUtf8 query)), + ("key", Just (TE.encodeUtf8 (config ^. #apiKey))), + ("limit", Just (B8.pack (show (maybe 10 id maxResults)))) + ] + request0 + httpLbs request manager + case eitherDecode (responseBody response) of + Right result -> pure result + Left err -> throwError (ProviderFailure ("web_search: could not decode search response: " <> T.pack err)) + +httpIO :: (EnvRow es) => Text -> IO a -> Eff es a +httpIO label action = do + result <- unsafeEff_ (try action) + case result of + Right value -> pure value + Left (err :: HttpException) -> + throwError (ProviderFailure (label <> ": " <> T.pack (show err))) + +contentTypeOf :: ResponseHeaders -> Text +contentTypeOf headers = + maybe "" (TE.decodeUtf8With lenientDecode) (lookup "Content-Type" headers) diff --git a/test/AcceptanceSpec.hs b/test/AcceptanceSpec.hs --- a/test/AcceptanceSpec.hs +++ b/test/AcceptanceSpec.hs @@ -7,6 +7,8 @@ module AcceptanceSpec (tests) where import Baikai (Response) +import Control.Lens ((&), (.~)) +import Data.Generics.Labels () import Data.Maybe (fromMaybe) import Data.Text qualified as T import Data.Vector qualified as V @@ -25,7 +27,6 @@ import MockLLM (runAgent) import Shikumi.Agent.ReAct ( Action (..), - ReActConfig (..), Step (..), Termination (..), ToolProtocol (..), @@ -51,7 +52,7 @@ res <- runAgent badArgsPromptScript - (reactWithTrajectory weatherSignature weatherRegistry defaultReActConfig {protocol = ProtocolPrompt}) + (reactWithTrajectory weatherSignature weatherRegistry (defaultReActConfig & #protocol .~ ProtocolPrompt)) weatherQuestion case res of Right (o :: WeatherResp, traj) -> do @@ -68,7 +69,7 @@ res <- runAgent script - (reactWithTrajectory weatherSignature weatherRegistry defaultReActConfig {protocol = proto}) + (reactWithTrajectory weatherSignature weatherRegistry (defaultReActConfig & #protocol .~ proto)) weatherQuestion case res of Right (o :: WeatherResp, traj) -> do diff --git a/test/BuiltinAcceptanceSpec.hs b/test/BuiltinAcceptanceSpec.hs new file mode 100644 --- /dev/null +++ b/test/BuiltinAcceptanceSpec.hs @@ -0,0 +1,140 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE OverloadedStrings #-} + +module BuiltinAcceptanceSpec (tests) where + +import Baikai (Response) +import Control.Exception (bracket) +import Control.Lens ((&), (.~)) +import Data.Aeson (object, (.=)) +import Data.Generics.Labels () +import Data.Maybe (fromMaybe) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.IO qualified as TIO +import Data.Vector qualified as V +import GHC.Generics (Generic) +import MockLLM (mkTextResponse, mkToolCallResponse, runAgent) +import Shikumi.Adapter (ToPrompt) +import Shikumi.Agent.ReAct + ( Action (..), + Step (..), + Termination (..), + ToolProtocol (..), + Trajectory (..), + defaultReActConfig, + reactWithTrajectory, + ) +import Shikumi.Schema (FromModel, ToSchema) +import Shikumi.Signature (Signature, mkSignature) +import Shikumi.Tool.Builtin (builtinRegistry) +import Shikumi.Tool.Env (localToolEnv) +import Shikumi.Tool.Web (FetchResult (..), SearchResult (..), WebClient (..)) +import System.Directory qualified as Dir +import System.FilePath (()) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +data WorkTask = WorkTask {task :: !Text} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToPrompt) + +data WorkAnswer = WorkAnswer {done :: !Bool} + deriving stock (Generic, Show, Eq) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +tests :: TestTree +tests = + testGroup + "BuiltinAcceptance" + [ testCase "builtinRegistry drives a scripted work cycle" $ + withTempDir "builtin" $ \root -> do + let file = T.pack (root "README.md") + TIO.writeFile (T.unpack file) "# Old Title\n\nbody\n" + res <- + runAgent + (script (T.pack root) file) + (reactWithTrajectory workSignature (builtinRegistry localToolEnv stubWebClient) (defaultReActConfig & #protocol .~ ProtocolNative)) + WorkTask {task = "Update the title and verify it."} + finalContent <- TIO.readFile (T.unpack file) + case res of + Left err -> assertFailure ("builtin agent failed: " <> show err) + Right (answer, traj) -> do + answer @?= WorkAnswer {done = True} + termination traj @?= TerminatedFinish + finalContent @?= "# New Title\n\nbody\n" + let stepList = V.toList (steps traj) + V.length (steps traj) @?= 6 + toolNames stepList @?= ["read", "edit", "bash", "grep", "web_fetch"] + case stepList of + [readStep, editStep, bashStep, grepStep, fetchStep, finishStep] -> do + assertObservation "read observes original content" "Old Title" readStep + assertObservation "edit reports replacement" "replacements" editStep + assertObservation "bash sees new title" "New Title" bashStep + assertObservation "grep sees new title" "New Title" grepStep + assertObservation "web_fetch sees stub body" "stub body" fetchStep + action finishStep @?= Finish + _ -> assertFailure "expected five tool steps and one finish step" + ] + +workSignature :: Signature WorkTask WorkAnswer +workSignature = mkSignature "Use the builtin tools to complete the requested repository work." + +script :: Text -> Text -> [Response] +script root file = + [ mkToolCallResponse "call_read" "read" (object ["path" .= file]), + mkToolCallResponse + "call_edit" + "edit" + ( object + [ "path" .= file, + "oldString" .= ("# Old Title" :: Text), + "newString" .= ("# New Title" :: Text) + ] + ), + mkToolCallResponse + "call_bash" + "bash" + (object ["command" .= ("cat " <> shellQuote file), "timeoutMs" .= (5000 :: Int)]), + mkToolCallResponse + "call_grep" + "grep" + ( object + [ "pattern" .= ("New Title" :: Text), + "path" .= root, + "glob" .= ("**/*.md" :: Text) + ] + ), + mkToolCallResponse "call_fetch" "web_fetch" (object ["url" .= ("https://example.test/docs" :: Text)]), + mkTextResponse "The requested work is complete.", + mkTextResponse "{\"done\": true}" + ] + +stubWebClient :: WebClient +stubWebClient = + WebClient + { webFetch = \_ _ -> + pure FetchResult {status = 200, contentType = "text/plain", body = "stub body", truncated = False}, + webSearch = \_ _ -> pure SearchResult {hits = []} + } + +toolNames :: [Step] -> [Text] +toolNames = foldMap $ \step -> case action step of + CallTool name _ -> [name] + Finish -> [] + +assertObservation :: String -> Text -> Step -> IO () +assertObservation label needle step = + assertBool label (needle `T.isInfixOf` fromMaybe "" (observation step)) + +withTempDir :: FilePath -> (FilePath -> IO a) -> IO a +withTempDir suffix action = do + tmp <- Dir.getTemporaryDirectory + let root = tmp ("shikumi-tools-" <> suffix) + bracket + (Dir.removePathForcibly root >> Dir.createDirectoryIfMissing True root >> pure root) + Dir.removePathForcibly + action + +shellQuote :: Text -> Text +shellQuote t = "'" <> T.replace "'" "'\\''" t <> "'" diff --git a/test/CompactionSpec.hs b/test/CompactionSpec.hs new file mode 100644 --- /dev/null +++ b/test/CompactionSpec.hs @@ -0,0 +1,179 @@ +module CompactionSpec (tests) where + +import Baikai (_Model, _Usage) +import Control.Lens ((&), (.~)) +import Data.Generics.Labels () +import Data.Text (Text) +import Data.Vector qualified as V +import Effectful (runEff) +import Effectful.Error.Static (runErrorNoCallStack) +import Fixtures + ( WeatherResp, + expectedWeather, + weatherQuestion, + weatherRegistry, + weatherSignature, + ) +import MockLLM + ( mkTextResponse, + mkUsageResponse, + runAgent, + runEffMock, + runMockLLMThrowingOn, + ) +import Shikumi.Agent.ReAct + ( Action (..), + ReActConfig (..), + Step (..), + Termination (..), + ToolProtocol (..), + Trajectory (..), + defaultReActConfig, + reactWithTrajectory, + ) +import Shikumi.Compaction + ( CompactionConfig (..), + compactTail, + defaultCompactionConfig, + overflowThreshold, + usageExceedsWindow, + ) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.Program (runProgram) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Compaction" + [ testCase "usageExceedsWindow flips at the boundary" $ do + let cfg = defaultCompactionConfig {reserveTokens = 100} + model = _Model & #contextWindow .~ 1000 + usage n = _Usage & #inputTokens .~ n + overflowThreshold cfg model @?= 900 + usageExceedsWindow cfg model (usage 899) @?= False + usageExceedsWindow cfg model (usage 900) @?= True + usageExceedsWindow (cfg {enabled = False}) model (usage 900) @?= False + usageExceedsWindow cfg _Model (usage 0) @?= False + overflowThreshold cfg (_Model & #contextWindow .~ 50) @?= 0, + testCase "compactTail folds older items and keeps the recent tail" $ do + let cfg = defaultCompactionConfig {keepRecent = 2} + res <- + runEffMock [mkTextResponse "S"] $ + compactTail cfg _Model id ("summary:" <>) (["e1", "e2", "e3", "e4", "e5", "e6"] :: [Text]) + res @?= Right ["summary:S", "e5", "e6"], + testCase "agent on tiny window compacts and completes" $ do + let cfg = + defaultReActConfig + { maxIters = 5, + protocol = ProtocolPrompt, + compaction = defaultCompactionConfig {reserveTokens = 10, keepRecent = 1} + } + model = _Model & #contextWindow .~ 100 + script = + [ mkUsageResponse model 10 (callReply "first"), + mkUsageResponse model 90 (callReply "second"), + mkTextResponse "summary of the first step", + mkUsageResponse model 10 finishReply, + mkTextResponse extractReply + ] + res <- runAgent script (reactWithTrajectory weatherSignature weatherRegistry cfg) weatherQuestion + case res of + Right (o :: WeatherResp, traj) -> do + o @?= expectedWeather + termination traj @?= TerminatedFinish + let ss = V.toList (steps traj) + assertBool "trajectory contains a summary step" (any isSummary ss) + case ss of + [summary, recent, finish] -> do + assertBool "first step is summary" (isSummary summary) + case action recent of + CallTool "get_weather" _ -> assertBool "recent step records an observation" (observation recent /= Nothing) + other -> assertFailure ("recent step should be get_weather, got " <> show other) + action finish @?= Finish + other -> assertFailure ("expected summary, recent tool step, finish; got " <> show other) + Left e -> assertFailure ("agent failed: " <> show e), + testCase "overflow error is caught, compacted, and retried once" $ do + let cfg = + defaultReActConfig + { maxIters = 5, + protocol = ProtocolPrompt, + compaction = defaultCompactionConfig {reserveTokens = 10, keepRecent = 0} + } + model = _Model & #contextWindow .~ 100 + script = + [ mkUsageResponse model 10 (callReply "first"), + mkTextResponse "reactive summary", + mkTextResponse finishReply, + mkTextResponse extractReply + ] + prog = reactWithTrajectory weatherSignature weatherRegistry cfg + res <- + runEff + . runErrorNoCallStack @ShikumiError + . runMockLLMThrowingOn [2] (ContextWindowExceeded "context length exceeded") script + $ runProgram prog weatherQuestion + case res of + Right (_ :: WeatherResp, traj) -> + assertBool "trajectory contains a summary step after recovery" (any isSummary (V.toList (steps traj))) + Left e -> assertFailure ("agent failed: " <> show e), + testCase "extract overflow is caught, compacted, and retried once" $ do + let cfg = + defaultReActConfig + { maxIters = 5, + protocol = ProtocolPrompt, + compaction = defaultCompactionConfig {reserveTokens = 10, keepRecent = 1} + } + model = _Model & #contextWindow .~ 100 + script = + [ mkUsageResponse model 10 (callReply "first"), + mkTextResponse finishReply, + mkTextResponse "extract recovery summary", + mkTextResponse extractReply + ] + prog = reactWithTrajectory weatherSignature weatherRegistry cfg + res <- + runEff + . runErrorNoCallStack @ShikumiError + . runMockLLMThrowingOn [3] (ContextWindowExceeded "context length exceeded") script + $ runProgram prog weatherQuestion + case res of + Right (_ :: WeatherResp, traj) -> + assertBool "returned trajectory contains extract recovery summary" (any isSummary (V.toList (steps traj))) + Left e -> assertFailure ("agent failed: " <> show e), + testCase "overflow retry is bounded to one retry" $ do + let cfg = + defaultReActConfig + { maxIters = 5, + protocol = ProtocolPrompt, + compaction = defaultCompactionConfig {reserveTokens = 10, keepRecent = 0} + } + model = _Model & #contextWindow .~ 100 + script = + [ mkUsageResponse model 10 (callReply "first"), + mkTextResponse "reactive summary" + ] + prog = reactWithTrajectory weatherSignature weatherRegistry cfg + res <- + runEff + . runErrorNoCallStack @ShikumiError + . runMockLLMThrowingOn [2, 4] (ContextWindowExceeded "context length exceeded") script + $ runProgram prog weatherQuestion + res @?= Left (ContextWindowExceeded "context length exceeded") + ] + +callReply :: Text -> Text +callReply th = + "{\"thought\": \"" + <> th + <> "\", \"action\": {\"tool\": \"get_weather\", \"args\": {\"city\": \"Paris\", \"units\": \"c\"}}}" + +finishReply :: Text +finishReply = "{\"thought\": \"I have the forecast.\", \"action\": {\"finish\": true}}" + +extractReply :: Text +extractReply = "{\"tempC\": 12.0, \"summary\": \"mild\"}" + +isSummary :: Step -> Bool +isSummary s = thought s == "(compacted summary of earlier steps)" diff --git a/test/EnvSpec.hs b/test/EnvSpec.hs new file mode 100644 --- /dev/null +++ b/test/EnvSpec.hs @@ -0,0 +1,89 @@ +{-# LANGUAGE OverloadedStrings #-} + +module EnvSpec (tests) where + +import Control.Lens ((^.)) +import Data.ByteString qualified as BS +import Data.Generics.Labels () +import Data.List (find) +import Data.Text qualified as T +import MockLLM (runEffMock) +import Shikumi.Tool.Env + ( DirEntry, + ExecRequest (..), + envCwd, + envExec, + envExists, + envMkdir, + envReadFile, + envReaddir, + envRm, + envStat, + envWriteFile, + localToolEnv, + ) +import System.Directory qualified as Dir +import System.FilePath (()) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Tool.Env" + [ testCase "localToolEnv can perform filesystem operations and exec" $ do + root <- freshTempDir + let file = T.pack (root "hello.txt") + nested = T.pack (root "nested") + result <- + runEffMock [] $ do + envWriteFile localToolEnv file "hello\n" + bytes <- envReadFile localToolEnv file + stat <- envStat localToolEnv file + entries <- envReaddir localToolEnv (T.pack root) + exists <- envExists localToolEnv file + envMkdir localToolEnv nested + nestedExists <- envExists localToolEnv nested + envRm localToolEnv nested + nestedGone <- not <$> envExists localToolEnv nested + execResult <- + envExec + localToolEnv + ExecRequest + { command = "echo hello", + cwd = Just (T.pack root), + stdin = Nothing, + timeoutMs = Just 5000 + } + cwd <- envCwd localToolEnv + pure (bytes, stat, entries, exists, nestedExists, nestedGone, execResult, cwd) + Dir.removePathForcibly root + case result of + Left err -> assertFailure ("localToolEnv failed: " <> show err) + Right (bytes, stat, entries, exists, nestedExists, nestedGone, execResult, cwd) -> do + bytes @?= "hello\n" + case stat of + Just fileStat -> do + fileStat ^. #isFile @?= True + fileStat ^. #isDir @?= False + fileStat ^. #size @?= fromIntegral (BS.length bytes) + Nothing -> assertFailure "expected stat for written file" + assertBool "readdir includes written file" (hasEntry "hello.txt" entries) + exists @?= True + nestedExists @?= True + nestedGone @?= True + execResult ^. #exitCode @?= 0 + assertBool "exec stdout includes hello" ("hello" `T.isInfixOf` (execResult ^. #stdout)) + assertBool "cwd returns an absolute path" ("/" `T.isPrefixOf` cwd) + ] + +freshTempDir :: IO FilePath +freshTempDir = do + tmp <- Dir.getTemporaryDirectory + let root = tmp "shikumi-tools-envspec" + Dir.removePathForcibly root + Dir.createDirectoryIfMissing True root + pure root + +hasEntry :: T.Text -> [DirEntry] -> Bool +hasEntry wanted = maybe False (const True) . find (\entry -> entry ^. #name == wanted) diff --git a/test/FsSpec.hs b/test/FsSpec.hs new file mode 100644 --- /dev/null +++ b/test/FsSpec.hs @@ -0,0 +1,161 @@ +{-# LANGUAGE OverloadedStrings #-} + +module FsSpec (tests) where + +import Control.Exception (bracket) +import Control.Lens ((^.)) +import Data.ByteString qualified as BS +import Data.List (sort) +import Data.Text (Text) +import Data.Text qualified as T +import MockLLM (runEffMock) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.Tool (Tool (..)) +import Shikumi.Tool.Builtin.Fs + ( EditReq (..), + GlobReq (..), + GrepMatch (..), + GrepReq (..), + ReadReq (..), + WriteReq (..), + editTool, + globTool, + grepTool, + readTool, + writeTool, + ) +import Shikumi.Tool.Env + ( ExecRequest (..), + ExecResult (..), + ToolEnv (..), + envCwd, + envExec, + envExists, + envMkdir, + envReadFile, + envReaddir, + envRm, + envStat, + envWriteFile, + localToolEnv, + ) +import System.Directory qualified as Dir +import System.FilePath (()) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Tool.Fs" + [ testCase "read/write/edit/grep/glob work through fast and fallback paths" $ + withTempDir "roundtrip" $ \root -> do + let file = T.pack (root "note.txt") + result <- + runEffMock [] $ do + rgPresent <- envExec localToolEnv (presentReq "rg") + fdPresent <- envExec localToolEnv (presentReq "fd") + written <- run (writeTool localToolEnv) WriteReq {path = file, content = "alpha\nold title\n"} + before <- run (readTool localToolEnv) ReadReq {path = file, offset = Nothing, limit = Nothing} + edited <- + run + (editTool localToolEnv) + EditReq {path = file, oldString = "old title", newString = "new title", replaceAll = Nothing} + after <- run (readTool localToolEnv) ReadReq {path = file, offset = Nothing, limit = Nothing} + fastGrep <- + run + (grepTool localToolEnv) + GrepReq {patternText = "new title", path = Just (T.pack root), glob = Just "**/*.txt", ignoreCase = Nothing} + fallbackGrep <- + run + (grepTool noFastToolEnv) + GrepReq {patternText = "new title", path = Just (T.pack root), glob = Just "**/*.txt", ignoreCase = Nothing} + fastGlob <- run (globTool localToolEnv) GlobReq {patternText = "**/*.txt", path = Just (T.pack root)} + fallbackGlob <- run (globTool noFastToolEnv) GlobReq {patternText = "**/*.txt", path = Just (T.pack root)} + pure (rgPresent, fdPresent, written, before, edited, after, fastGrep, fallbackGrep, fastGlob, fallbackGlob) + case result of + Left err -> assertFailure ("fs tools failed: " <> show err) + Right (rgPresent, fdPresent, written, before, edited, after, fastGrep, fallbackGrep, fastGlob, fallbackGlob) -> do + rgPresent ^. #exitCode @?= 0 + fdPresent ^. #exitCode @?= 0 + written ^. #bytesWritten @?= BS.length "alpha\nold title\n" + before ^. #content @?= "alpha\nold title\n" + edited ^. #replacements @?= 1 + after ^. #content @?= "alpha\nnew title\n" + normalizeMatches (fastGrep ^. #matches) @?= normalizeMatches (fallbackGrep ^. #matches) + assertBool "grep finds edited text" (any (\m -> m ^. #text == "new title") (fastGrep ^. #matches)) + assertBool "fast glob finds file" (file `elem` (fastGlob ^. #paths)) + assertBool "fallback glob finds file" (file `elem` (fallbackGlob ^. #paths)), + testCase "fallback grep skips noisy directories and binary files" $ + withTempDir "hardening" $ \root -> do + let visible = T.pack (root "visible.txt") + gitFile = T.pack (root ".git" "packed.txt") + nodeFile = T.pack (root "node_modules" "dep.txt") + binaryFile = T.pack (root "binary.bin") + result <- + runEffMock [] $ do + envWriteFile localToolEnv visible "SECRET visible\n" + envMkdir localToolEnv (T.pack (root ".git")) + envWriteFile localToolEnv gitFile "SECRET git\n" + envMkdir localToolEnv (T.pack (root "node_modules")) + envWriteFile localToolEnv nodeFile "SECRET node\n" + envWriteFile localToolEnv binaryFile (BS.pack [0, 83, 69, 67, 82, 69, 84]) + run + (grepTool noFastToolEnv) + GrepReq {patternText = "SECRET", path = Just (T.pack root), glob = Nothing, ignoreCase = Nothing} + case result of + Left err -> assertFailure ("fallback grep failed: " <> show err) + Right resp -> do + normalizeMatches (resp ^. #matches) @?= [GrepMatch {file = visible, line = 1, text = "SECRET visible"}] + resp ^. #truncated @?= False, + testCase "fallback grep reports malformed regex as ValidationFailure" $ + withTempDir "bad-regex" $ \root -> do + let visible = T.pack (root "visible.txt") + result <- + runEffMock [] $ do + envWriteFile localToolEnv visible "hello\n" + run + (grepTool noFastToolEnv) + GrepReq {patternText = "[", path = Just (T.pack root), glob = Nothing, ignoreCase = Nothing} + case result of + Left (ValidationFailure msg) -> assertBool "mentions invalid regex" ("invalid regex" `T.isInfixOf` msg) + other -> assertFailure ("expected ValidationFailure, got " <> show other) + ] + +withTempDir :: FilePath -> (FilePath -> IO a) -> IO a +withTempDir suffix action = do + tmp <- Dir.getTemporaryDirectory + let root = tmp ("shikumi-tools-fsspec-" <> suffix) + bracket + (Dir.removePathForcibly root >> Dir.createDirectoryIfMissing True root >> pure root) + Dir.removePathForcibly + action + +presentReq :: Text -> ExecRequest +presentReq command = + ExecRequest + { command = "command -v " <> command, + cwd = Nothing, + stdin = Nothing, + timeoutMs = Just 5000 + } + +noFastToolEnv :: ToolEnv +noFastToolEnv = + ToolEnv + { envExec = \req -> + if "command -v rg" `T.isPrefixOf` (req ^. #command) || "command -v fd" `T.isPrefixOf` (req ^. #command) + then pure ExecResult {exitCode = 1, stdout = "", stderr = ""} + else envExec localToolEnv req, + envReadFile = envReadFile localToolEnv, + envWriteFile = envWriteFile localToolEnv, + envStat = envStat localToolEnv, + envReaddir = envReaddir localToolEnv, + envExists = envExists localToolEnv, + envMkdir = envMkdir localToolEnv, + envRm = envRm localToolEnv, + envCwd = envCwd localToolEnv + } + +normalizeMatches :: [GrepMatch] -> [GrepMatch] +normalizeMatches = sort diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -3,14 +3,20 @@ module Main (main) where import AcceptanceSpec qualified +import BuiltinAcceptanceSpec qualified import CodeActSpec qualified +import CompactionSpec qualified +import EnvSpec qualified +import FsSpec qualified import ProgramOfThoughtSpec qualified import ProtocolSpec qualified import ReActSpec qualified import RestrictedSpec qualified import SchemaSpec qualified +import ShellSpec qualified import Test.Tasty (defaultMain, testGroup) import ToolSpec qualified +import WebSpec qualified main :: IO () main = @@ -19,10 +25,16 @@ "shikumi-tools" [ SchemaSpec.tests, ToolSpec.tests, + EnvSpec.tests, + WebSpec.tests, + FsSpec.tests, + ShellSpec.tests, ReActSpec.tests, ProtocolSpec.tests, AcceptanceSpec.tests, + BuiltinAcceptanceSpec.tests, RestrictedSpec.tests, ProgramOfThoughtSpec.tests, - CodeActSpec.tests + CodeActSpec.tests, + CompactionSpec.tests ] diff --git a/test/MockLLM.hs b/test/MockLLM.hs --- a/test/MockLLM.hs +++ b/test/MockLLM.hs @@ -7,15 +7,19 @@ -- exhausted script yields an empty text response (specs always script enough). module MockLLM ( runMockLLM, + runMockLLMThrowingOnce, + runMockLLMThrowingOn, runEffMock, runAgent, mkTextResponse, + mkUsageResponse, mkToolCallResponse, ) where import Baikai ( AssistantContent (..), + Model, Response, _Response, _TextContent, @@ -29,7 +33,8 @@ 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 Effectful.Error.Static (Error, runErrorNoCallStack, throwError) +import Numeric.Natural (Natural) import Shikumi.Error (ShikumiError) import Shikumi.LLM (LLM (..)) import Shikumi.Program (Program, runProgram) @@ -46,6 +51,38 @@ ) act +-- | Interpret @LLM@ like 'runMockLLM', but throw once on the first completion. +runMockLLMThrowingOnce :: + (IOE :> es, Error ShikumiError :> es) => + ShikumiError -> + [Response] -> + Eff (LLM : es) a -> + Eff es a +runMockLLMThrowingOnce = runMockLLMThrowingOn [1] + +-- | Interpret @LLM@ like 'runMockLLM', but throw on selected 1-based completion +-- calls. Useful for bounded-retry tests. +runMockLLMThrowingOn :: + (IOE :> es, Error ShikumiError :> es) => + [Int] -> + ShikumiError -> + [Response] -> + Eff (LLM : es) a -> + Eff es a +runMockLLMThrowingOn throwAt err script act = do + ref <- liftIO (newIORef script) + countRef <- liftIO (newIORef (0 :: Int)) + interpret + ( \_ -> \case + Complete {} -> do + n <- liftIO (atomicModifyIORef' countRef (\n0 -> let n1 = n0 + 1 in (n1, n1))) + if n `elem` throwAt + then throwError err + else 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@. @@ -73,6 +110,14 @@ mkTextResponse :: Text -> Response mkTextResponse t = _Response & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t)) + +-- | A text response that also carries a resolved model and input-token usage. +mkUsageResponse :: Model -> Natural -> Text -> Response +mkUsageResponse model inputTokens text = + mkTextResponse text + & #model .~ model + & #message . #usage . #inputTokens .~ inputTokens + & #message . #usage . #totalTokens .~ inputTokens -- | An assistant 'Response' carrying a single native tool-call block. mkToolCallResponse :: Text -> Text -> Value -> Response diff --git a/test/ProtocolSpec.hs b/test/ProtocolSpec.hs --- a/test/ProtocolSpec.hs +++ b/test/ProtocolSpec.hs @@ -21,8 +21,7 @@ ) import MockLLM (runAgent) import Shikumi.Agent.ReAct - ( ReActConfig (..), - Termination (..), + ( Termination (..), ToolProtocol (..), Trajectory (..), defaultReActConfig, @@ -40,12 +39,12 @@ nat <- runAgent nativeScript - (reactWithTrajectory weatherSignature weatherRegistry defaultReActConfig {protocol = ProtocolNative}) + (reactWithTrajectory weatherSignature weatherRegistry (defaultReActConfig & #protocol .~ ProtocolNative)) weatherQuestion pro <- runAgent promptScript - (reactWithTrajectory weatherSignature weatherRegistry defaultReActConfig {protocol = ProtocolPrompt}) + (reactWithTrajectory weatherSignature weatherRegistry (defaultReActConfig & #protocol .~ ProtocolPrompt)) weatherQuestion case (nat, pro) of (Right (oN :: WeatherResp, tN), Right (oP, tP)) -> do diff --git a/test/ReActSpec.hs b/test/ReActSpec.hs --- a/test/ReActSpec.hs +++ b/test/ReActSpec.hs @@ -3,6 +3,8 @@ -- never finishes stops at @maxIters@ and still extracts a best-effort answer. module ReActSpec (tests) where +import Control.Lens ((&), (.~)) +import Data.Generics.Labels () import Data.Vector qualified as V import Fixtures ( WeatherResp, @@ -16,7 +18,6 @@ import MockLLM (runAgent) import Shikumi.Agent.ReAct ( Action (..), - ReActConfig (..), Step (..), Termination (..), Trajectory (..), @@ -51,7 +52,7 @@ _ -> assertFailure "expected exactly two steps" Left e -> assertFailure ("agent failed: " <> show e), testCase "stops at maxIters with TerminatedMaxIters" $ do - let cfg = defaultReActConfig {maxIters = 1} + let cfg = defaultReActConfig & #maxIters .~ 1 res <- runAgent maxItersScript diff --git a/test/ShellSpec.hs b/test/ShellSpec.hs new file mode 100644 --- /dev/null +++ b/test/ShellSpec.hs @@ -0,0 +1,40 @@ +{-# LANGUAGE OverloadedStrings #-} + +module ShellSpec (tests) where + +import Control.Lens ((^.)) +import Data.Text qualified as T +import MockLLM (runEffMock) +import Shikumi.Tool (Tool (..)) +import Shikumi.Tool.Builtin.Shell (BashReq (..), bashTool) +import Shikumi.Tool.Env (localToolEnv) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Tool.Shell" + [ testCase "bash captures stdout and zero exit" $ do + result <- + runEffMock [] $ + run + (bashTool localToolEnv) + BashReq {command = "echo hi", cwd = Nothing, timeoutMs = Just 5000, stdin = Nothing} + case result of + Left err -> assertFailure ("bashTool failed: " <> show err) + Right resp -> do + resp ^. #exitCode @?= 0 + assertBool "stdout contains hi" ("hi" `T.isInfixOf` (resp ^. #stdout)), + testCase "bash returns stderr and non-zero exit as a value" $ do + result <- + runEffMock [] $ + run + (bashTool localToolEnv) + BashReq {command = "echo oops 1>&2; exit 3", cwd = Nothing, timeoutMs = Just 5000, stdin = Nothing} + case result of + Left err -> assertFailure ("bashTool failed: " <> show err) + Right resp -> do + resp ^. #exitCode @?= 3 + assertBool "stderr contains oops" ("oops" `T.isInfixOf` (resp ^. #stderr)) + ] diff --git a/test/WebSpec.hs b/test/WebSpec.hs new file mode 100644 --- /dev/null +++ b/test/WebSpec.hs @@ -0,0 +1,115 @@ +{-# LANGUAGE OverloadedStrings #-} + +module WebSpec (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 MockLLM (runEffMock) +import Shikumi.Tool (SomeTool (..), ToolRegistry, mkRegistry, runToolCall) +import Shikumi.Tool.Builtin.Web (webFetchTool, webSearchTool) +import Shikumi.Tool.Web + ( FetchResult (..), + SearchHit (..), + SearchResult (..), + WebClient (..), + localWebClient, + newTlsManager, + ) +import System.Environment (lookupEnv) +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.Web" + [ testCase "web_fetch returns a stubbed 200 observation" $ do + result <- + runEffMock [] $ + runToolCall + stubRegistry + (tc "web_fetch" (object ["url" .= ("https://example.test" :: Text)])) + case result of + Right (Right obs) -> do + assertBool "observation includes status" ("200" `T.isInfixOf` obs) + assertBool "observation includes body" ("stub body" `T.isInfixOf` obs) + other -> assertFailure ("expected web_fetch observation, got " <> show other), + testCase "web_fetch surfaces a 404 status as a value" $ do + result <- + runEffMock [] $ + runToolCall + notFoundRegistry + (tc "web_fetch" (object ["url" .= ("https://example.test/missing" :: Text)])) + case result of + Right (Right obs) -> do + assertBool "observation includes 404" ("404" `T.isInfixOf` obs) + assertBool "observation includes missing body" ("missing" `T.isInfixOf` obs) + other -> assertFailure ("expected 404 as a tool value, got " <> show other), + testCase "web_search returns stubbed hits" $ do + result <- + runEffMock [] $ + runToolCall + stubRegistry + (tc "web_search" (object ["query" .= ("shikumi" :: Text)])) + case result of + Right (Right obs) -> do + assertBool "observation includes title" ("Stub Result" `T.isInfixOf` obs) + assertBool "observation includes url" ("https://example.test/result" `T.isInfixOf` obs) + other -> assertFailure ("expected web_search observation, got " <> show other), + testCase "live web_fetch is gated by SHIKUMI_NET_TESTS" $ do + enabled <- lookupEnv "SHIKUMI_NET_TESTS" + case enabled of + Nothing -> pure () + Just _ -> do + manager <- newTlsManager + result <- + runEffMock [] $ + runToolCall + (mkRegistry [SomeTool (webFetchTool (localWebClient manager Nothing))]) + (tc "web_fetch" (object ["url" .= ("https://example.com" :: Text)])) + case result of + Right (Right obs) -> do + assertBool "live observation includes 200" ("200" `T.isInfixOf` obs) + assertBool "live observation has a body" ("Example Domain" `T.isInfixOf` obs) + Right (Left err) -> assertFailure ("live web_fetch returned tool error: " <> show err) + Left err -> assertFailure ("live web_fetch failed: " <> show err) + ] + +stubRegistry :: ToolRegistry +stubRegistry = mkRegistry [SomeTool (webFetchTool stubWebClient), SomeTool (webSearchTool stubWebClient)] + +notFoundRegistry :: ToolRegistry +notFoundRegistry = mkRegistry [SomeTool (webFetchTool notFoundWebClient)] + +stubWebClient :: WebClient +stubWebClient = + WebClient + { webFetch = \_ _ -> + pure FetchResult {status = 200, contentType = "text/plain", body = "stub body", truncated = False}, + webSearch = \_ _ -> + pure + SearchResult + { hits = + [ SearchHit + { title = "Stub Result", + url = "https://example.test/result", + snippet = "A stubbed search hit." + } + ] + } + } + +notFoundWebClient :: WebClient +notFoundWebClient = + WebClient + { webFetch = \_ _ -> + pure FetchResult {status = 404, contentType = "text/plain", body = "missing", truncated = False}, + webSearch = \_ _ -> pure SearchResult {hits = []} + }