R/tool-environment.R

In btw: A Toolkit for Connecting R and Large Language Models

#' Describe the contents of an environment
#'
#' @examples
#' env <- new.env()
#' env$cyl_6 <- mtcars[mtcars$cyl == 6, ]
#' env$gear_5 <- mtcars[mtcars$gear == 5, ]
#' btw_this(env)
#'
#' @param x An environment.
#' @param items Optional. A character vector of objects in the environment to
#'   describe.
#' @param ... Additional arguments are silently ignored.
#'
#' @returns
#' A string describing the environment contents with `#>` prefixing
#' each object's printed representation.
#'
#' @seealso [btw_tool_env_describe_environment()]
#'
#' @family btw formatting methods
#' @export
btw_this.environment <- function(x, ..., items = NULL) {
  btw_tool_env_describe_environment_impl(environment = x, items = items)@value
}

#' Tool: Describe an environment
#'
#' This tool can be used by the LLM to describe the contents of an R session,
#' i.e. the data frames and other objects loaded into the global environment.
#' This tool will only see variables that you've named and created in the
#' global environment, it cannot reach into package namespaces, see which
#' packages you have loaded, or access files on your computer.
#'
#' @examples
#' my_cars <- mtcars[mtcars$mpg > 25, ]
#' btw_tool_env_describe_environment("my_cars")
#'
#' @inheritParams btw_this.environment
#' @inheritParams btw_tool_docs_package_news
#'
#' @inherit btw_this.environment return
#'
#' @seealso [btw_this.environment()], [btw_tools()]
#' @family Tools
#' @export
btw_tool_env_describe_environment <- function(items, `_intent`) {}

btw_tool_env_describe_environment_impl <- function(
  items = NULL,
  ...,
  environment = global_env()
) {
  check_dots_empty()

  if (!is.environment(environment)) {
    # TODO: does the env name live in the global env?
    # is it in `search_envs`?
    cli::cli_abort("Not implemented yet.")
  }

  if (identical(trimws(items), "")) {
    items <- NULL
  }

  if (is_namespace(environment)) {
    cli::cli_abort(c(
      "Describing an entire package namespace is not supported",
      "i" = "Try choosing specific functions to describe: {.code btw('dplyr::mutate', 'dplyr::across')}."
    ))
  }

  res <- character()
  env_item_names <- ls(environment)
  if (!is.null(items)) {
    # Subset to `items`, keeping the order of `items`
    env_item_names <- intersect(items, env_item_names)
  }

  item_desc_prev <- NULL

  item_desc <- map(env_item_names, function(item_name) {
    item <- env_get(environment, item_name)

    if (identical(class(item), "character")) {
      # Only string literals passed through btw() hit `btw_this.character()`.
      # We rely on `dots_list()` turning `"foo"` into `list('"foo"' = "foo")`.
      item_name_dots_listed <- gsub("\\", "\\\\", item, fixed = TRUE)
      item_name_dots_listed <- sprintf('"%s"', item_name_dots_listed)
      if (!identical(item_name, item_name_dots_listed)) {
        item <- btw_returns_character(item)
      }
    }

    btw_this(item, caller_env = environment)
  })

  res <- c()
  for (i in seq_along(item_desc)) {
    desc <- item_desc[[i]]
    name <- env_item_names[[i]]
    is_user_prompt <- inherits(desc, "btw_user_prompt")

    if (i == 1) {
      res <- c(
        if (!is_user_prompt) c("## Context", ""),
        btw_item_with_description(name, desc)
      )
      next
    }

    is_adjacent_user_prompt <-
      is_user_prompt &&
      inherits(item_desc[[i - 1]], "btw_user_prompt")

    is_adjacent_user_context <-
      !is_user_prompt &&
      inherits(item_desc[[i - 1]], "btw_user_prompt")

    if (is_adjacent_user_prompt) {
      # Append text to previous prompt text
      res[length(res)] <- paste0(res[length(res)], "\n", desc)
    } else {
      res <- c(
        res,
        "",
        btw_item_with_description(
          name,
          desc,
          header = if (is_adjacent_user_context) "## Context"
        )
      )
    }
  }

  if (identical(res, c("## Context", ""))) {
    return(BtwToolResult(""))
  }

  BtwToolResult(res)
}

.btw_add_to_tools(
  name = "btw_tool_env_describe_environment",
  group = "env",
  tool = function() {
    ellmer::tool(
      function(items = NULL) {
        btw_tool_env_describe_environment_impl(items = items)
      },
      name = "btw_tool_env_describe_environment",
      description = "List and describe items in the R session's global environment.",
      annotations = ellmer::tool_annotations(
        title = "Object in Session",
        read_only_hint = TRUE,
        open_world_hint = FALSE,
        btw_can_register = function() TRUE
      ),
      arguments = list(
        items = ellmer::type_array(
          "The names of items to describe from the environment. Defaults to `NULL`, indicating all items.",
          items = ellmer::type_string(),
          required = FALSE
        )
      )
    )
  }
)

btw_item_with_description <- function(item_name, description, header = NULL) {
  if (inherits(description, "AsIs")) {
    return(description)
  }
  if (inherits(description, "btw_user_prompt")) {
    return(c("## User", description))
  }
  if (inherits(description, "btw_ignore")) {
    return(invisible())
  }
  if (inherits(description, "btw_captured")) {
    item_name <- gsub('^"|"$', '', item_name)
    item_name <- switch(
      item_name,
      "@last_error" = "last_error()",
      "@last_value" = ".Last.value",
      item_name
    )

    description <- md_code_block("r", item_name, paste("#>", description))
    item_name <- NULL
  }

  if (!is.null(header)) {
    header <- c(header, "")
  }
  paste(c(header, item_name, description), collapse = "\n")
}