ReplaceInFile.hs

module LLM.Tools.ReplaceInFile (replaceInFileToolTyped, ReplaceInFileToolArgs (..)) where

import Autodocodec qualified as AC
import Data.Aeson (FromJSON)
import Data.Text (Text)
import Data.Text qualified as T
import GHC.Generics (Generic)
import LLM.Core.Types (TypedTool (..))
import LLM.Tools.FsConfig (FsConfig, sandboxPath, withWorkspacePaths)
import LLM.Tools.FsLimits (maxReadBytes, readBoundedUtf8File, writeUtf8File)
import LLM.Tools.Internal.PathLock (withPathLock)
import LLM.Tools.Internal.TextEdit (countOccurrences, replaceFirst)

data ReplaceInFileToolArgs = ReplaceInFileToolArgs
  { _rifPath :: Text,
    _rifOld :: Text,
    _rifNew :: Text
  }
  deriving (Generic)
  deriving (FromJSON) via (AC.Autodocodec ReplaceInFileToolArgs)

instance AC.HasCodec ReplaceInFileToolArgs where
  codec :: AC.JSONCodec ReplaceInFileToolArgs
  codec =
    AC.object "replace a string in a file" $
      ReplaceInFileToolArgs
        <$> AC.requiredField "path" "Relative file path to write to" AC..= (._rifPath)
        <*> AC.requiredField "old" "The text content to be replaced in the file" AC..= (._rifOld)
        <*> AC.requiredField "new" "The replacement text content to write to the file" AC..= (._rifNew)

replaceInFileToolTyped :: FsConfig -> TypedTool ctx ReplaceInFileToolArgs
replaceInFileToolTyped cfg =
  TypedTool
    { ttoolName = "replace_in_file",
      ttoolDescription =
        "Replace the first occurrence of a string in a file. "
          <> "The 'old' string must appear exactly once in the file. "
          <> "Returns an error if the string is not found or appears more than once. "
          <> "Files larger than "
          <> T.pack (show maxReadBytes)
          <> " bytes are refused.",
      ttoolReadonly = False,
      ttoolExecute = const (withWorkspacePaths cfg . replaceExecTyped cfg)
    }

replaceExecTyped :: FsConfig -> ReplaceInFileToolArgs -> IO Text
replaceExecTyped cfg ReplaceInFileToolArgs {_rifPath = p, _rifOld, _rifNew} = do
  resolved <- sandboxPath cfg (T.unpack p)
  -- Hold the lock from the read to the write, so a concurrent edit of the
  -- same file cannot be lost.
  withPathLock resolved $ do
    contentE <- readBoundedUtf8File resolved p maxReadBytes
    case contentE of
      Left err -> pure $ "Error: " <> err
      Right content -> case countOccurrences _rifOld content of
        0 -> pure "Error: the 'old' string was not found in the file"
        1 -> do
          let replaced = replaceFirst _rifOld _rifNew content
          writeUtf8File resolved replaced
          pure $ "Successfully replaced text in " <> p
        n ->
          pure $
            "Error: the 'old' string was found "
              <> T.pack (show n)
              <> " times; it must appear exactly once"